feat: support first class function via builtin function

Author: HerrCai0907

src/builtins.ts

@@ -43,7 +43,8 @@ import {
   NodeKind,
   LiteralExpression,
   ArrayLiteralExpression,
-  IdentifierExpression
+  IdentifierExpression,
+  FunctionExpression
 } from "./ast";
 
 import {
@@ -194,6 +195,8 @@ export namespace BuiltinNames {
   export const instantiate = "~lib/builtins/instantiate";
   export const idof = "~lib/builtins/idof";
 
+  export const experimental_first_class_function = "~lib/builtins/experimental_first_class_function";
+
   export const i8 = "~lib/builtins/i8";
   export const i16 = "~lib/builtins/i16";
   export const i32 = "~lib/builtins/i32";
@@ -3609,6 +3612,26 @@ function builtin_unchecked(ctx: BuiltinFunctionContext): ExpressionRef {
 }
 builtinFunctions.set(BuiltinNames.unchecked, builtin_unchecked);
 
+// experimental_first_class_function(expr: *) -> *
+function builtin_experimental_first_class_function(ctx: BuiltinFunctionContext): ExpressionRef {
+  let compiler = ctx.compiler;
+  let module = compiler.module;
+  if (
+    checkTypeAbsent(ctx) |
+    checkArgsRequired(ctx, 1)
+  ) return module.unreachable();
+  let operand = ctx.operands[0];
+  if (operand.kind != NodeKind.Function) {
+    let prototype = ctx.prototype;
+    prototype.program.error(DiagnosticCode._0_expected, operand.range, "FunctionExpression");
+    return module.unreachable();
+  }
+  let functionExpression = <FunctionExpression>operand;
+  let expr = compiler.compileFirstClassFunction(functionExpression, ctx.contextualType, Constraints.None);
+  return expr;
+}
+builtinFunctions.set(BuiltinNames.experimental_first_class_function, builtin_experimental_first_class_function);
+
 // call_indirect<T?>(index: u32, ...args: *[]) -> T
 function builtin_call_indirect(ctx: BuiltinFunctionContext): ExpressionRef {
   let compiler = ctx.compiler;

src/compiler.ts

@@ -2054,20 +2054,23 @@ export class Compiler extends DiagnosticEmitter {
 
   // === Table ====================================================================================
 
+  private registerFunctionInTable(instance: Function): u32 {
+    // Add to the function table
+    let functionTable = this.functionTable;
+    let tableBase = this.options.tableBase;
+    if (!tableBase) tableBase = 1; // leave first elem blank
+    let index = tableBase + functionTable.length;
+    functionTable.push(instance);
+    return index;
+  }
+
   /** Ensures that a runtime counterpart of the specified function exists and returns its address. */
   ensureRuntimeFunction(instance: Function): i64 {
     assert(instance.is(CommonFlags.Compiled) && !instance.is(CommonFlags.Stub));
     let program = this.program;
     let memorySegment = instance.memorySegment;
     if (!memorySegment) {
-
-      // Add to the function table
-      let functionTable = this.functionTable;
-      let tableBase = this.options.tableBase;
-      if (!tableBase) tableBase = 1; // leave first elem blank
-      let index = tableBase + functionTable.length;
-      functionTable.push(instance);
-
+      let index = this.registerFunctionInTable(instance);
       // Create runtime function
       let rtInstance = assert(this.resolver.resolveClass(program.functionPrototype, [ instance.type ]));
       let buf = rtInstance.createBuffer();
@@ -7026,12 +7029,96 @@ export class Compiler extends DiagnosticEmitter {
     return module.unreachable();
   }
 
+  compileFirstClassFunction(
+    expression: FunctionExpression,
+    contextualType: Type,
+    constraints: Constraints
+  ): ExpressionRef {
+    let module = this.module;
+    let instance = this.doCompileFunctionExpression(expression, contextualType, constraints);
+    let currentType = this.currentType;
+    if (!instance) return module.unreachable(); 
+    let rtInstance = assert(this.resolver.resolveClass(this.program.functionPrototype, [ instance.type ]));
+    const functionIndexInTable = this.registerFunctionInTable(instance);
+    let ctor = this.ensureConstructor(rtInstance, expression);
+    let indexMember = assert(rtInstance.getMember("_index"));
+    assert(indexMember.kind == ElementKind.PropertyPrototype);
+    let indexProperty = assert((<PropertyPrototype>indexMember).instance);
+    let indexSetter = assert(indexProperty.setterInstance);
+    let tmp = this.currentFlow.getTempLocal(Type.i32);
+    let expr = module.block(null, [
+      this.makeCallDirect(
+        indexSetter,
+        [
+          module.local_tee(
+            tmp.index,
+            this.makeCallDirect(ctor, [module.i32(0)], expression, /*immediatelyDropped*/ false),
+            /*isManaged*/ true, // TODO(maybe can be false to optimize performance)
+            TypeRef.I32
+          ),
+          module.i32(functionIndexInTable),
+        ],
+        expression
+      ),
+      module.local_get(tmp.index, TypeRef.I32 )
+    ], TypeRef.I32);
+    this.currentType = currentType;
+    return expr;
+  }
+
   private compileFunctionExpression(
     expression: FunctionExpression,
     contextualType: Type,
     constraints: Constraints
   ): ExpressionRef {
-    let declaration = expression.declaration.clone(); // generic contexts can have multiple
+    let declaration = expression.declaration;
+    let isNamed = declaration.name.text.length > 0;
+    let module = this.module;
+    let flow = this.currentFlow;
+    let isSemanticallyAnonymous = !isNamed || contextualType != Type.void;
+
+    let instance = this.doCompileFunctionExpression(expression, contextualType, constraints);
+    if (!instance) return module.unreachable(); 
+    let offset = this.ensureRuntimeFunction(instance); // reports
+    let expr = this.options.isWasm64
+      ? module.i64(i64_low(offset), i64_high(offset))
+      : module.i32(i64_low(offset));
+
+    // add a constant local referring to the function if applicable
+    if (!isSemanticallyAnonymous) {
+      let fname = instance.name;
+      let existingLocal = flow.getScopedLocal(fname);
+      if (existingLocal) {
+        if (!existingLocal.declaration.range.source.isNative) {
+          this.errorRelated(
+            DiagnosticCode.Duplicate_identifier_0,
+            declaration.name.range,
+            existingLocal.declaration.name.range,
+            fname
+          );
+        } else { // scoped locals are shared temps that don't track declarations
+          this.error(
+            DiagnosticCode.Duplicate_identifier_0,
+            declaration.name.range, fname
+          );
+        }
+      } else {
+        let ftype = instance.type;
+        let local = flow.addScopedLocal(instance.name, ftype);
+        flow.setLocalFlag(local.index, LocalFlags.Constant | LocalFlags.Initialized);
+        expr = module.local_tee(local.index, expr, ftype.isManaged);
+      }
+    }
+
+    return expr;
+  }
+
+  private doCompileFunctionExpression(
+    expression: FunctionExpression,
+    contextualType: Type,
+    constraints: Constraints
+  ): Function | null {
+    let declaration = expression.declaration;
     assert(!declaration.typeParameters); // function expression cannot be generic
     let flow = this.currentFlow;
     let sourceFunction = flow.sourceFunction;
@@ -7047,7 +7134,6 @@ export class Compiler extends DiagnosticEmitter {
     );
     let instance: Function | null;
     let contextualTypeArguments = cloneMap(flow.contextualTypeArguments);
-    let module = this.module;
 
     // compile according to context. this differs from a normal function in that omitted parameter
     // and return types can be inferred and omitted arguments can be replaced with dummies.
@@ -7065,7 +7151,7 @@ export class Compiler extends DiagnosticEmitter {
           DiagnosticCode.Expected_0_arguments_but_got_1,
           expression.range, numParameters.toString(), numPresentParameters.toString()
         );
-        return module.unreachable();
+        return null;
       }
 
       // check non-omitted parameter types
@@ -7077,13 +7163,13 @@ export class Compiler extends DiagnosticEmitter {
             sourceFunction.parent,
             contextualTypeArguments
           );
-          if (!resolvedType) return module.unreachable();
+          if (!resolvedType) return null;
           if (!parameterTypes[i].isStrictlyAssignableTo(resolvedType)) {
             this.error(
               DiagnosticCode.Type_0_is_not_assignable_to_type_1,
               parameterNode.range, parameterTypes[i].toString(), resolvedType.toString()
             );
-            return module.unreachable();
+            return null;
           }
         }
         // any unused parameters are inherited but ignored
@@ -7097,7 +7183,7 @@ export class Compiler extends DiagnosticEmitter {
           sourceFunction.parent,
           contextualTypeArguments
         );
-        if (!resolvedType) return module.unreachable();
+        if (!resolvedType) return null;
         if (
           returnType == Type.void
             ? resolvedType != Type.void
@@ -7107,7 +7193,7 @@ export class Compiler extends DiagnosticEmitter {
             DiagnosticCode.Type_0_is_not_assignable_to_type_1,
             signatureNode.returnType.range, resolvedType.toString(), returnType.toString()
           );
-          return module.unreachable();
+          return null;
         }
       }
 
@@ -7120,20 +7206,20 @@ export class Compiler extends DiagnosticEmitter {
             DiagnosticCode._this_cannot_be_referenced_in_current_location,
             thisTypeNode.range
           );
-          return module.unreachable();
+          return null;
         }
         let resolvedType = this.resolver.resolveType(
           thisTypeNode,
           sourceFunction.parent,
           contextualTypeArguments
         );
-        if (!resolvedType) return module.unreachable();
+        if (!resolvedType) return null;
         if (!thisType.isStrictlyAssignableTo(resolvedType)) {
           this.error(
             DiagnosticCode.Type_0_is_not_assignable_to_type_1,
             thisTypeNode.range, thisType.toString(), resolvedType.toString()
           );
-          return module.unreachable();
+          return null;
         }
       }
 
@@ -7148,52 +7234,21 @@ export class Compiler extends DiagnosticEmitter {
       instance.flow.outer = flow;
       let worked = this.compileFunction(instance);
       this.currentType = contextualSignature.type;
-      if (!worked) return module.unreachable();
+      if (!worked) return null;
 
     // otherwise compile like a normal function
     } else {
       instance = this.resolver.resolveFunction(prototype, null, contextualTypeArguments);
-      if (!instance) return this.module.unreachable();
+      if (!instance) return null;
       instance.flow.outer = flow;
       let worked = this.compileFunction(instance);
       this.currentType = instance.signature.type;
-      if (!worked) return module.unreachable();
+      if (!worked) return null;
     }
-
-    let offset = this.ensureRuntimeFunction(instance); // reports
-    let expr = this.options.isWasm64
-      ? module.i64(i64_low(offset), i64_high(offset))
-      : module.i32(i64_low(offset));
-
-    // add a constant local referring to the function if applicable
-    if (!isSemanticallyAnonymous) {
-      let fname = instance.name;
-      let existingLocal = flow.getScopedLocal(fname);
-      if (existingLocal) {
-        if (!existingLocal.declaration.range.source.isNative) {
-          this.errorRelated(
-            DiagnosticCode.Duplicate_identifier_0,
-            declaration.name.range,
-            existingLocal.declaration.name.range,
-            fname
-          );
-        } else { // scoped locals are shared temps that don't track declarations
-          this.error(
-            DiagnosticCode.Duplicate_identifier_0,
-            declaration.name.range, fname
-          );
-        }
-      } else {
-        let ftype = instance.type;
-        let local = flow.addScopedLocal(instance.name, ftype);
-        flow.setLocalFlag(local.index, LocalFlags.Constant | LocalFlags.Initialized);
-        expr = module.local_tee(local.index, expr, ftype.isManaged);
-      }
-    }
-
-    return expr;
+    return instance;
   }
 
+
   /** Makes sure the enclosing source file of the specified expression has been compiled. */
   private maybeCompileEnclosingSource(expression: Expression): void {
     let internalPath = expression.range.source.internalPath;

src/diagnostics.ts

@@ -454,6 +454,7 @@ export abstract class DiagnosticEmitter {
     arg1: string | null = null,
     arg2: string | null = null
   ): void {
+    console.trace()
     this.emitDiagnostic(code, DiagnosticCategory.Error, range, null, arg0, arg1, arg2);
   }
 

std/assembly/builtins.ts

@@ -202,6 +202,10 @@ export declare function assert<T>(isTrueish: T, message?: string): T;
 @unsafe @builtin
 export declare function unchecked<T>(expr: T): T;
 
+// @ts-ignore: decorator
+@unsafe @builtin
+export declare function experimental_first_class_function<T>(expr: T): T;
+
 // @ts-ignore: decorator
 @unsafe @builtin
 export declare function call_indirect<T>(index: u32, ...args: auto[]): T;

std/assembly/index.d.ts

@@ -281,6 +281,8 @@ declare function decodeURI(str: string): string;
 /** Decodes a Uniform Resource Identifier (URI) component previously created by encodeURIComponent. */
 declare function decodeURIComponent(str: string): string;
 
+declare function experimental_first_class_function<T>(value: T): T & Function;
+
 /** Atomic operations. */
 declare namespace atomic {
   /** Atomically loads an integer value from memory and returns it. */