Implement closures

[MaxGraey]

I decide to start discussion about simple closure implementations.

Some obvious (and may be naive) implementation is using generation class context which I prefer to demonstrate in following example:

declare function externalCall(a: i32, b: i32): void;

function range(a: i32, b: i32, fn: (n: i32) => void): void {
  if (a < b) {
    fn(a);
    range(a + 1, b, fn);
  }
}

export function test(n: i32): void {
  range(0, n, (i: i32) => {
    externalCall(i, n); // capture n
  });
}

which transform to:

// generated
class ClosureContext {
  fn: (ctx: usize, i: i32) => void;
  n: i32; // captured var
  // optinal "self: usize;" is closure instantiate inside instance class method
  parent: ClosureContext | null = null;
}

// generated
function lambdaFn(ctx: usize, i: i32): void {
  externalCall(i, changetype<ClosureContext>(ctx).n);
}

function range(a: i32, b: i32, ctx: usize): void {
  if (a < b) {
    changetype<ClosureContext>(ctx).fn(ctx, a); // replaced from "fn(a)";
    range(a + 1, b, ctx);
  }
}

export function test(n: i32): void {
  // insert
  let ctx = new ClosureContext();
  ctx.fn = lambdaFn;
  ctx.n = n;
  //
  range(0, n, changetype<usize>(ctx));
}

Closure and ClosureContext will not generated when no one variable was captured and use usual anonym functions.
ClosureContext will not created when only this was captured. In this case ctx param use to pass this reference.

Other discussions: #563

@dcodeIO @jtenner @willemneal Let me know what you think about this?

[jtenner]

I think we need to override the meaning of call indirect.

[jtenner]

How would this transform?

function add(a, b): callback {
  return () => a + b;
}

[MaxGraey]

@jtenner

class ClosureContext {
  fn: (ctx: usize) => i32;
  a: i32;
  b: i32;
  parent: ClosureContext | null = null;
}

function lambdaAdd(ctx: usize): i32 {
  return changetype<ClosureContext>(ctx).a + changetype<ClosureContext>(ctx).b;
}

function add(a: i32, b: i32): ClosureContext {
  let ctx = new ClosureContext();
  ctx.fn = lambdaAdd;
  ctx.a = a;
  ctx.b = b;
  return ctx;
}

[MaxGraey]

Instead

class ClosureContext {
  fn: (ctx: usize) => i32;
  b: i32;
  ...
}

we could just store index for indirect function table

class ClosureContext {
  fnIdx: usize;
  a: i32;
  ...
}

... 

call_indirect(ctx.fnIdx, ...args)

[jtenner]

How does this work with function dispatch?

For instance, let's say I pass the Closure Context out to js as a pointer. How can I re-call it?

[jtenner]

Is there a way to make a table and add entries to it?

[MaxGraey]

@jtenner actually you need pass only fn / fnIdx. After that just use same approach as for anonymus indirect function calls.

EDIT No you need unpack fn from returned ClosureContext object. Just provide another util for loader

[MaxGraey]

Next example:

function test(fn: (x: i32) => void): i32 {
  let n = 0;
  fn(x => { n = x });
  return n;
}

should generate:

class ClosureContext {
  fn: (ctx: usize, x: i32) => void;
  n: i32;
  parent: ClosureContext | null = null;
}

function lambdaFn(ctx: usize, x: i32): void {
   changetype<ClosureContext>(ctx).n = x;
}

function test(fn: (x: i32) => void): i32 {
  let n = 0;
  let ctx = new ClosureContext();
  ctx.fn = lambdaFn;
  ctx.n = n;
  fn(changetype<usize>(ctx));
  n = ctx.n;
  return n;
}

[jtenner]

Well I'm thinking about aspect collecting function pointers. How will aspect need to work with the function pointers?

[jtenner]

My guess is that passing around the closure context will cause problems with manually using call_indirect like aspect does.

Also, this closure method doesn't handle multiple references to the same local.

let a = 1;
let b = () => a;
let c = () => a += 1;

B and c will get different versions of a.

[MaxGraey]

@jtenner
You could get table and get function by index / ref. For example you have wasm:

(func $foo (result i32) (i32.const 1234))
(table (export "tbl") anyfunc (elem $foo))

than you js part:

WebAssembly.instantiateStreaming(fetch('main.wasm')).then(({ instance }) => {
  const table = instance.exports.tbl;
  console.log(table.get(0)());  // 1234
});

[jtenner]

It's likely we will need to allocate enclosed local values in a table or on the heap. Each pointer to those values will need to be stored in a table pointing to the heap.

This idea is Naive because the variables can no longer be treated like local variables because it's possible to modify local values before the function finishes executing.

[MaxGraey]

B and c will get different versions of a.

Why?

let a = 1;
let b = () => a;
let c = () => a += 1;

let br = b(); // 1
let cr = c(); // 2
// assert(a == 2)

Will generate:

class ClosureContextB { fn; a; }
class ClosureContextC { fn; a; }

function lambdaB(ctx: usize): i32 {
   return changetype<ClosureContextB>(ctx).a;
}

function lambdaC(ctx: usize): i32 {
   return changetype<ClosureContextC>(ctx).a += 1;
}

let a = 1;

let ctxB = new ClosureContextB(lambdaB, a);
let br = b(ctxB); // 1
// a = ctxB.a; // 1 unmodified so unnecessary

let ctxC = new ClosureContextB(lambdaC, a);
let cr = c(ctxC); // 2
a = ctxC.a; // 2

[jtenner]

I'm saying a in that example needs to exist on the heap for both b and c to access it, and the closure class needs to contain a Box<i32> that points to the heap location.

[MaxGraey]

No, we don't need pass plain types by boxed references. Also found pretty clear article. So ClosureContext (LexicalEnvironment) should be little bit modifier and also store reference to it's parent LexicalEnvironment.

[dcodeIO]

Regarding collection of closure contexts: Seems the idea here is to pass around a closure context (containing both the function index and the lexical scope) instead of just a function index. While that can be reference counted, it leads to a situation where something can be called with either a function index or a closure context, for example

function callIt(fn: () => void): void { fn(); }

function nop(): void {}
callIt(nop); // function index

let a: i32;
function ctx(): void { a = 1; }
callIt(ctx); // closure context

which means we'd either have to generate two versions of callIt (one taking a function index, one taking a closure context, otherwise doing the same) or doing only closure contexts, requiring every call to callIt to wrap function indexes in a temporary empty closure, which is an unnecessary allocation that is free'd immediately.

A better approach might be to utilize the multi-value spec, in that a closure is actually two values, a function index and a lexical environment, with the latter possibly being null.

[jtenner]

Yep. The issue I'm going to hit with a multivalue return is when these function pointers need to be utilized in JavaScript.

For instance, I want to call a describe function pointer that is nested. In order to do this from AssemblyScript, I need to export a function callIt that has 0 knowledge of lexical scope.

Edit: Could we have a primitive like lexscopeof(func)? That way I could at least manage it externally in javascript

[MaxGraey]

multi-value approach is definitely better but I'm not sure standalone VMs will support it. Its quite complicated even for compiler tools. (binaryen still not fully impement it). So having fallback (without multi-values support) it seems necessary even if not so efficient and require some allocation in heap

[dcodeIO]

Maybe one way to work around the allocation is to keep a singleton closure context per non-closure around in static memory. Like, if we know that the table has max 200 elements, make 200 dummies pre-populated with the function index and no lexical scope? Hmm

[jtenner]

Well at least 200 might not be enough. I can imagine taking advantage of this feature in aspect in very terrifying ways

[willemneal]

Firstly, I'd like to follow up on Daniel's example, would something like this work?

type context<T> = T extends Context | T extends Function
function callIt<context<T>>(fn: T): returnof<T> { //This function would need to change.
  if (isFunction<T>()){
    if isVoid<T>() {
     fn(); 
     return;
    }
    return fn();
  } 
  let ctx = changetype<Context>(fn);
  if (ctx.isVoid){ // Need to add a isVoid property.
     ctx.call()
     return
  }
  return ctx.call();
}

I read this article a year ago about how Elm handles first class functions in wasm: https://dev.to/briancarroll/elm-functions-in-webassembly-50ak

The big take away is that a function context is a function pointer, its current arity (or how many arguments it still has to take) and an ArrayBuffer of the arguments that have been passed.

When you have new function: let f = new Func(fn) and you the do f.call(..) You pass it the last parameter and you get a new function context which now takes one less argument. This continues until you have one argument left at which point the function pointer of the context is called. In the context of the article above, everything is immutable, which is why you get a clone of the context + the new argument. This way the intermediate functions that return functions of a smaller arity can be reused.

let add = (a: i32, b: 32): i32 => a + b;
let addOne = add(1); ==> this is now a function  `a: i32 => a + 1`;
let addTwo = add(2);
let two = addOne(1);
let three = addTwo(1);

This could look something like this:

class Func<Fn> {
   // fn: Fn; //function to be called when all arguments are present
   // airity: usize; // current airity

  get length(): usize {
    return lengthof<Fn>();
  }

  get currentArg(): usize {
    return this.length - this.arity;
  }
  
  constructor(public fn: Fn, public arity: usize, public args: ArrayBuffer){}

   static create<Fn>(fn: Fn, arity: usize = lengthof<Fn>(), args: new ArrayBuffer(sizeof<u64>() * lengthof<Fn>()): {
    //This isn't exactly right size the size of each parameter could vary. But for sake of simplicity let's assume they are all usize.
     //  Let's assume there is a builtin `paramType<Fn>(0)`
    type argType = paramType<Fn>(0);
    let func;
    if (arity > 1) ? paramType<Fn>(lengthof<Fn>() - arity + 1) : returnof<Fn>();
    let func = new Func<argType, Fn>(fn, arity, args);
    return func;
   }
 
  call<T>(arg: T): Func<Fn> | returnof<Fn>() {
    assert(arg instanceof paramType<Fn>(this.currentArg());
    if (arity == 0) { // arg is the last one and we can call the function 
      this.fn(...this.args); //This clearly needs to be a compiler operation to load the arguments as locals.  Or transform all paremeter references with offsets and pass the array/arraybuffer.
   }
   let args = this.args.clone();
   store<T>(args.dataStart + this.currentArg(), arg)
   return new Func<Fn>(fn, this.arity - 1, args);
  }
}

I know a lot of this isn't currently possible, just wanted to get the point across.

[dcodeIO]

Potential implementation idea:

• A function reference is represented by a function table index currently.
• Closures would be special function references with an attached context, means: managed allocations

Problem: The same function can be called with either a function table index or a closure pointer.
Solution: We know that allocations have an alignment of 16 bytes, so if we make all function table indexes with that alignment invalid (essentially skipping every 16th function table index, making it the null function), a callee can check whether it is dealing with a function table index or a closure pointer by means of:

if (fn & 15) {
  call_indirect(fn, ...);
} else {
  ctx = __retain(fn);
  call_indirect(ctx.index, ...);
  __release(ctx);
}

Here, ctx is a global referencing the current closure context that must be set before a closure is being executed. A closure context is a class-like managed structure of the form:

ClosureContext extends HEADER {
  index: i32;
  closedOverVariable1: Type1;
  ...
  closedOverVariableN: TypeN;
}

Compilation:

When encountering a closed-over variable

1. Remember that the current function is a closure
2. For each such variable, lay these out as of ClosureContext layout in advance
3. Replace what would be a local.get or local.set of each such variable with a load respectively store relative to current ctx.

When compiling a call to such a function

1. Allocate the closure context and assign it to ctx
2. Copy the function index, then each closed-over local's value to the closure context
3. Call the function normally
4. Copy the local values back to the caller's locals, or context if it is itself a closure
5. Release the closure context

Performance impact:

• Each closure call implies an allocation, retain and release, while non-closure calls do not
• Each indirect call implies a (fairly predictable) branch checking for function table index vs closure.
• Function references become refcounted only if !(val & 15) (means: is a closure pointer)

What am I missing? :)

[MaxGraey]

Isn't this approach cause to table fragmentation? I predict closures can be much much more often than ordinal indirect functions

[ghost]

Not an AS user, but if you want performance and fine low-level control, it might be best to have separate function types, or to change the meaning of arrow functions vs "normal" functions.
If all functions are functors, then wouldn't ASC assume JS imports are also functors?
This would also force explicit casting between functions and functors, ex:

Let's take this AS function:

function for_range(
    start: u32,
    stop: u32,
    callback: (n: u32) => void
) {
    for (let iterator = start; iterator < stop; ++iterator) {
        callback(iterator);
    }
}

if this were called with a "raw" function, ex, from the host side, then the code would break.

It might make more sense to do something more like:

function for_range(
    start: u32,
    stop: u32,
    callback: RawFunction<(n: u32) => void>
) {
    for (let iterator: u32 = start; iterator < stop; ++iterator) {
        callback(iterator);
    }
}

function foo(): void {
    let i: u32 = 0;

    for_range(0, 10, (n: u32): void => { log(++i) });
    // Error: "(n: u32) => void" is not assignable to RawFunction<(n: u32) => void>
}

Then again, that might not work out at all.

---

Also, you'd think that browsers, which have been optimized 24/7 to execute JS for the last decade, would know what locals to capture, and what not to, but their JiT compilers seem to make mistakes too often.

Explicit capturing might provide stronger guarantees about what is captured, by preventing potentially arbitrary compiler decisions.

[MaxGraey]

Also, you'd think that browsers, which have been optimized 24/7 to execute JS for the last decade, would know what locals to capture, and what not to, but their JiT compilers seem to make mistakes too often.

Found which function capture environment (has free variables) and actually is closure and which is not is pretty simple and fully precise and doing in ahead of time during analysis inside AS module. Main problem with functions / closures which cross wasm<->host boundary

[ghost]

...is pretty simple and fully precise and doing in ahead of time during analysis inside AS module.

That is a pretty bold assertion. ASC performance will likely slow down too.
Also, each closure context will be different, so passing functions around won't be possible, at least not without dynamic typing/type pruning.
Take that for_range function from earlier, each function would need to be called with a new closure type.

Next, consider this:

type ClosureStruct = {
        set: (n: i32) => void,
        get: () => i32
};

function foo(x: i32): ClosureStruct {
    return {
        set: (new_value: i32): void => {
            x = new_value;
        },
        get: (): i32 => x
    };
}

function bar(): void {
    const struct: ClosureStruct = foo(100);

    struct.set(10);
    struct.get(10);
}

Would this be accessing invalid data?
How much GC will this require on ASC's end?
AS might benefit from Wasm function reference binding.

[MaxGraey]

Take that for_range function from earlier, each function would need to be called with a new closure type.

For for_range we even don't need a closure. After lambda lifting it will be simple indirect/direct call.

[MaxGraey]

C++ can't do this usually. But Rust which have own intermediate representation called MIR provide similar transforms. See this here:
https://godbolt.org/z/4qsM68

[ghost]

C++ can't do this usually.

C++ couldn't do it because of the std::function class; generally std::function calls will be a magnitude worse than function calls.
An untyped lambda easily causes the same codegen as Rust did.
See here: https://godbolt.org/z/6Y8h4T
Although, of course, it differs between compilers.

But, good luck with whatever implementation that the ASC team uses for AS closures!

[MaxGraey]

C++ couldn't do it because of the std::function class; generally std::function calls will be a magnitude worse than function calls.

Ah right! I'm not very familiar with modern C++. But if you check llvm-ir you got a lot codegen initially. So all credit lies with LLVM in this case and its further optimizations.

But, good luck with whatever implementation that the ASC team uses for AS closures!

Thanks!

[Dudeplayz]

What's the current state of closures?

How can I workaround the following code snippet:

export function addClockEvent(cpu: CPU, callbackId: u32, cycles: u32): AVRClockEventCallback {
    return cpu.addClockEvent(() => callClockEventCallback(callbackId), cycles)
}

I'm porting a library to AS, but there is the point where a function is passed as a parameter. Because I can't pass a function to WASM, I am building a map with a callbackId that references the correct function on the JS side and calling this in an imported JS function. Without the closure I can't pass the parameter into it and modifying the function signature is not possible because it is deeply integrated.

[MaxGraey]

How can I workaround the following code snippet:

Just use temporary global variable:

let _callbackId: u32;

export function addClockEvent(cpu: CPU, callbackId: u32, cycles: u32): AVRClockEventCallback {
    _callbackId = callbackId;
    return cpu.addClockEvent(() => callClockEventCallback(_callbackId), cycles)
}

But this has some limitations. You can't use this approach for recursive calls and for serial calls

[Dudeplayz]

I will get multiple calls to this function, so the global variable would be overwriten every time it get called. Is there any other method to do it or have I misunderstood something? For me it currently looks like I have to change the callback signature which is very problematic and blocking. Having closures would be nice.

[urish]

Here's the approach I used in a similar situation. It replaces the callback parameter with an interface, so you can easily create classes in that implement the interface and store additional information (sort of manual closure).

interface CallbackInstance {
  execute(): void;
}

class ClockEventCallback implement CallbackInstance {
  constructor(readonly callbackId: u32) {}

  execute(): void {
    callClockEventCallback(this.callbackId); // or whatever custom login you want to do.
  }
}

Then change the signature of addClockEvent to accept a CallbackInstance and call its execute method:

  addClockEvent(callback: CallbackInstance , cycles: u32): void {
    // now store callback somewhere, and use callback.execute() to call it.
  }

Finally, call it like:

cpu.addClockEvent(new ClockEventCallback(whatever), 1000);

I hope this is helpful!

[Dudeplayz]

I used the approach @urish mentioned. I was hoping there is an alternative because it changes the signature of the method, which results in a changed API.

[Dudeplayz]

I got another error while replacing a closure with an interface, the code looks like the following:

//callReadHook and callWriteHook are imported module functions
export class ExternalCPUMemoryReadHook implements CPUMemoryReadHook {
    call(addr: u16): u8 {
        return callReadHook(addr);
    }
}

export class ExternalCPUMemoryWriteHook implements CPUMemoryHook {
    call(value: u8, oldValue: u8, addr: u16, mask: u8): boolean {
        return callWriteHook(value, oldValue, addr, mask);
    }
}

export function addReadHook(cpu: CPU, addr: u32): void {
    cpu.readHooks.set(addr, new ExternalCPUMemoryReadHook());
}

export function addWriteHook(cpu: CPU, addr: u32): void {
    cpu.writeHooks.set(addr, new ExternalCPUMemoryWriteHook());
}

// cpu.readHooks and cpu.writeHooks are of the following types:
export class CPUMemoryHooks extends Map<u32, CPUMemoryHook> {
}
export class CPUMemoryReadHooks extends Map<u32, CPUMemoryReadHook> {
}

Just the import breaks the build, removing the exported functions still crashes.
Here is the error log:

asc assembly/index.ts -b build/untouched.wasm -t build/untouched.wat -d build/module.d.ts --sourceMap C://Users/xxred/IdeaProjects/avr8js-research/build/untouched.wasm.map --debug --exportRuntime


▌ Whoops, the AssemblyScript compiler has crashed during buildTSD :-( 
▌ 
▌ Here is the stack trace hinting at the problem, perhaps it's useful?
▌ 
▌ AssertionError: assertion failed
▌     at assert (C:\Users\xxred\IdeaProjects\avr8js-research\node_modules\assemblyscript\dist\webpack:\assemblyscript\std\portable\index.js:201:11)
▌     at u.visitElement (C:\Users\xxred\IdeaProjects\avr8js-research\node_modules\assemblyscript\dist\webpack:\assemblyscript\src\definitions.ts:143:16)
▌     at u.visitFile (C:\Users\xxred\IdeaProjects\avr8js-research\node_modules\assemblyscript\dist\webpack:\assemblyscript\src\definitions.ts:80:14)
▌     at u.walk (C:\Users\xxred\IdeaProjects\avr8js-research\node_modules\assemblyscript\dist\webpack:\assemblyscript\src\definitions.ts:68:65)
▌     at u.build (C:\Users\xxred\IdeaProjects\avr8js-research\node_modules\assemblyscript\dist\webpack:\assemblyscript\src\definitions.ts:625:10)
▌     at Function.build (C:\Users\xxred\IdeaProjects\avr8js-research\node_modules\assemblyscript\dist\webpack:\assemblyscript\src\definitions.ts:376:36)
▌     at Object.t.buildTSD (C:\Users\xxred\IdeaProjects\avr8js-research\node_modules\assemblyscript\dist\webpack:\assemblyscript\src\index.ts:290:21)
▌     at C:\Users\xxred\IdeaProjects\avr8js-research\node_modules\assemblyscript\cli\asc.js:1173:34
▌     at measure (C:\Users\xxred\IdeaProjects\avr8js-research\node_modules\assemblyscript\cli\asc.js:1409:3)
▌     at Object.main (C:\Users\xxred\IdeaProjects\avr8js-research\node_modules\assemblyscript\cli\asc.js:1171:27)
▌ 
▌ If it refers to the dist files, try to 'npm install source-map-support' and
▌ run again, which should then show the actual code location in the sources.
▌ 
▌ If you see where the error is, feel free to send us a pull request. If not,
▌ please let us know: https://github.com/AssemblyScript/assemblyscript/issues
▌ 
▌ Thank you!

[tamusjroyce]

Any reason why closures can’t be rewritten as classes & pass through constructor. As far as I can see in the ast, there should be a 1 to 1 relationship going from closures & currying to classes. For example, Typescript classes to es3 translation results in closures.

[alienself]

Any updates on adding closure support? What is the blocker?

[luffyfly]

Any updates on adding closure support? What is the problem?

[luffyfly]

I planned to write a game engine use AS and I have programmed for a month, but now I give up AS and choose Rust because lack of supports for closures. Just let you know, It's not reasonable to not support closures in any case.

[gabriel-fallen]

@luffyfly as a Functional Programming enthusiast I totally agree that closures are super useful and handy. But are you sure closures in Rust are implemented for Wasm efficient enough for your game engine?

As a general observation, in the current state Wasm presents a poor target for Functional Programming: doesn't efficiently support heavy use of function pointers, closures and GC. Moreover AssemblyScript is not really a Functional Language. And if you're programming in OO style it's not hard to implement purpose-built closures out of objects (going as far as Strategy Pattern if really needed).

[luffyfly]

@gabriel-fallen I'm pretty sure about that, because When I call the fetch() function of the web host from Wasm, I have to use closures as callbacks to avoid blocking the main thread.

[MaxGraey]

I have to use closures as callbacks to avoid blocking the main thread.

Just to note. Closures and async (cooperative multitasking) are two independent things

[CryZe]

doesn't efficiently support heavy use of function pointers

Closures in Rust are rarely represented as function pointers and instead as anonymous structs that have a "call" method that is statically dispatched rather than dynamically in almost all cases. The method is almost always inlined too, to the point where almost always you can't even tell it ever was a closure to begin with. It's very strongly a zero cost abstraction there.

Though I'm also wondering why you say function pointers aren't well supported by Wasm. As far as I know they are just functions stored in the function table and then called via an call_indirect and an index into the table. I don't see a problem with that approach, other than having to know all the functions that are going to be called indirectly ahead of time, but that's not a problem for a compiled language.

GC though is definitely a problem with closures for languages that don't have an ownership system.

[luffyfly]

@CryZe I have tried function pointers, and It worked but It is so weird and difficult to use.

[luffyfly]

@MaxGraey If AS has aysnc implementations, I will not take closures.

[MaxGraey]

@luffyfly Again closures and aysnc is absolutely different things. Dsynchrony without calling "then" in returned Promise but through "await" does not use dispatch at all. You decide what you want.

[luffyfly]

@MaxGraey See #376 (comment).
Actually, I want both.

[butterunderflow]

What is the current state of the closure implementation? I've seen several PRs related to closures, but none of them have been merged. I'm curious about what is the biggest block of closure implementation.

[pchasco]

I am interested in using AssemblyScript, but must admit I am reluctant to use until closures are implemented. Yes, obviously there are workarounds, but they are inconvenient and add to code bloat. All TypeScript developers whom I assume this project hopes to pilfer are accustomed to the convenience of closures and to some degree of functional programming.

Many languages that target WASM already support closures, so I am confused as to what it is about AssemblyScript specifically that prevents closures from being implemented. I would be surprised if some limitation in WASM were a real roadblock. I saw some mentions of concerns around performance and allocations, which are reasonable, but I wonder if this is one of those situations where a bit of runtime efficiency is sacrificed for developer convenience? And the implementation wouldn't necessarily be the perfect implementation out the gate, so long as it is functionally correct. It could be refined or replaced later when new WASM features are introduced, or if a better implementation is worked out.

[alienself]

Any update on this? Another year without closures 😢

Will 2024 the year of closures and thread support in AssemblyScript?