fix bug from using one apply function

Author: jonasongg

src/wasm-compiler/builderGenerator.ts

@@ -29,9 +29,10 @@ import {
   SET_LEX_ADDR_FX,
   SET_PAIR_HEAD_FX,
   SET_PAIR_TAIL_FX,
+  SET_PARAM_FX,
 } from "./constants";
 import { f64, global, i32, i64, local, mut, wasm } from "./wasm-util/builder";
-import { WasmCall, WasmInstruction } from "./wasm-util/types";
+import { WasmCall, WasmInstruction, WasmRaw } from "./wasm-util/types";
 
 const builtInFunctions: {
   name: string;
@@ -102,8 +103,11 @@ type Binding = { name: string; tag: "local" | "nonlocal" };
 
 // all expressions compile to a call to a function like makeX, get/setLexAddress, arithOp, etc.
 // so expressions return WasmCalls. (every expression results in i32 i64)
+// WasmRaw is for function calls
 export class BuilderGenerator
-  implements StmtNS.Visitor<WasmInstruction>, ExprNS.Visitor<WasmCall>
+  implements
+    StmtNS.Visitor<WasmInstruction>,
+    ExprNS.Visitor<WasmCall | WasmRaw>
 {
   private strings: [string, number][] = [];
   private heapPointer = 0;
@@ -256,9 +260,12 @@ export class BuilderGenerator
           .func("$main")
           .results(...(undroppedInstr ? [i32, i64] : []))
           .body(
-            wasm
-              .call(ALLOC_ENV_FX)
-              .args(i32.const(globalEnvLength), i32.const(0)),
+            global.set(
+              CURR_ENV,
+              wasm
+                .call(ALLOC_ENV_FX)
+                .args(i32.const(globalEnvLength), i32.const(0), i32.const(0))
+            ),
 
             ...builtInFuncsDeclarations,
 
@@ -408,42 +415,35 @@ export class BuilderGenerator
       );
   }
 
-  visitCallExpr(expr: ExprNS.Call): WasmCall {
+  visitCallExpr(expr: ExprNS.Call): WasmRaw {
     const callee = this.visit(expr.callee);
+    const args = expr.args.map((arg) => this.visit(arg));
 
-    // because we're not actually passing in arguments to apply, the lex addresses from the args
-    // are not accurate.
-    // they assume that we're not yet in the new function call's environment, because that's how it is
-    // from the compiler's perspective.
-    // but here, we're already in the new environment because PRE_APPLY was called. so we have to
-    // increment the depth from each GET_LEX_ADDRESS call from each arg by 1.
-    const args: WasmCall[] = expr.args
-      .map((arg) => this.visit(arg))
-      .map((arg) => ({
-        ...arg,
-        ...(arg.function === GET_LEX_ADDR_FX.name && {
-          arguments: arg.arguments.map((a, i) => ({
-            ...a,
-
-            // first argument is depth
-            ...(i === 0 &&
-              a.op === "i32.const" && { value: a.value + BigInt(1) }),
-          })),
-        }),
-      }));
-
-    return wasm.call(APPLY_FX_NAME).args(
-      global.get(CURR_ENV),
-      wasm.call(PRE_APPLY_FX).args(callee, i32.const(args.length)),
-
-      // these are not arguments, but they don't produce values, so it's ok to insert them here
-      // this is to maintain the return type of WasmCall
-      // this is equivalent to setting lex addresses after calling PRE_APPLY, but before APPLY
-
-      ...[...Array(args.length).keys()].map((i) =>
-        wasm.call(SET_LEX_ADDR_FX).args(i32.const(0), i32.const(i), args[i])
-      )
-    );
+    // PRE_APPLY returns (1, 2) callee tag and value, (3) pointer to new environment
+    // APPLY expects (1) pointer to return environment, (2, 3) callee tag and value
+
+    // we call PRE_APPLY first, which verifies the callee is a closure and arity matches
+    // AND creates a new environment for the function call, but does not set CURR_ENV yet
+    // this is so that we can set the arguments in the new environment first
+
+    // this means we can't use SET_LEX_ADDR_FX because it uses CURR_ENV internally
+    // so we manually set the arguments in the new environment using SET_PARAM_FX
+
+    // the SET_PARAM function returns the env address after setting the parameter
+    // so we can chain the calls together
+    return wasm.raw`
+(global.get ${CURR_ENV})
+(call ${PRE_APPLY_FX.name} ${callee} (i32.const ${args.length}))
+
+${args.map(
+  (arg, i) =>
+    wasm.raw`
+(i32.const ${i * 12}) (i32.add) ${arg} (call ${SET_PARAM_FX.name})`
+)}
+
+(global.set ${CURR_ENV})
+(call ${APPLY_FX_NAME})
+`;
   }
 
   visitReturnStmt(stmt: StmtNS.Return): WasmInstruction {

src/wasm-compiler/constants.ts

@@ -605,22 +605,29 @@ export const COMPARISON_OP_FX = wasm
   );
 
 // *3*4 because each variable has a tag and payload = 3 words = 12 bytes; +4 because parentEnv is stored at start of env
+// we initialise only local variables to UNBOUND, NOT parameters.
+// this is because have already set parameters in the new environment before calling this function.
 export const ALLOC_ENV_FX = wasm
   .func("$_alloc_env")
-  .params({ $size: i32, $parent: i32 })
+  .params({ $size: i32, $parent: i32, $arity: i32 })
+  .results(i32)
   .body(
-    global.set(CURR_ENV, global.get(HEAP_PTR)),
-    i32.store(global.get(CURR_ENV), local.get("$parent")),
+    global.get(HEAP_PTR), // return the start of the new env, set CURR_ENV AFTER
+    i32.store(global.get(HEAP_PTR), local.get("$parent")),
     global.set(HEAP_PTR, i32.add(global.get(HEAP_PTR), i32.const(4))),
 
-    memory.fill(global.get(HEAP_PTR), i32.const(TYPE_TAG.UNBOUND), i32.mul(local.get("$size"), i32.const(12))),
+    memory.fill(
+      i32.add(global.get(HEAP_PTR), i32.mul(local.get("$arity"), i32.const(12))),
+      i32.const(TYPE_TAG.UNBOUND),
+      i32.mul(i32.sub(local.get("$size"), local.get("$arity")), i32.const(12))
+    ),
     global.set(HEAP_PTR, i32.add(global.get(HEAP_PTR), i32.mul(local.get("$size"), i32.const(12))))
   );
 
 export const PRE_APPLY_FX = wasm
   .func("$_pre_apply")
   .params({ $tag: i32, $val: i64, $arity: i32 })
-  .results(i32, i64)
+  .results(i32, i64, i32)
   .body(
     wasm
       .if(i32.ne(local.get("$tag"), i32.const(TYPE_TAG.CLOSURE)))
@@ -632,14 +639,15 @@ export const PRE_APPLY_FX = wasm
       )
       .then(wasm.call("$_log_error").args(i32.const(ERROR_MAP.FUNC_WRONG_ARITY[0])), wasm.unreachable()),
 
+    local.get("$tag"),
+    local.get("$val"),
     wasm
       .call(ALLOC_ENV_FX)
       .args(
         i32.and(i32.wrap_i64(i64.shr_u(local.get("$val"), i64.const(32))), i32.const(255)),
-        i32.wrap_i64(local.get("$val"))
-      ),
-    local.get("$tag"),
-    local.get("$val")
+        i32.wrap_i64(local.get("$val")),
+        local.get("$arity")
+      )
   );
 
 export const APPLY_FX_NAME = "$_apply";
@@ -673,10 +681,6 @@ export const GET_LEX_ADDR_FX = wasm
           i32.load(i32.add(i32.add(local.get("$env"), i32.const(4)), i32.mul(local.get("$index"), i32.const(12))))
         ),
 
-        wasm.call("$_log_int").args(i64.extend_i32_u(local.get("$depth"))),
-        wasm.call("$_log_int").args(i64.extend_i32_u(local.get("$index"))),
-        wasm.call("$_log_int").args(i64.extend_i32_u(local.get("$tag"))),
-
         wasm
           .if(i32.eq(local.get("$tag"), i32.const(TYPE_TAG.UNBOUND)))
           .then(wasm.call("$_log_error").args(i32.const(ERROR_MAP.UNBOUND[0])), wasm.unreachable()),
@@ -725,6 +729,17 @@ export const SET_LEX_ADDR_FX = wasm
     wasm.unreachable()
   );
 
+export const SET_PARAM_FX = wasm
+  .func("$_set_param")
+  .params({ $addr: i32, $tag: i32, $value: i64 })
+  .results(i32)
+  .body(
+    i32.store(i32.add(local.get("$addr"), i32.const(4)), local.get("$tag")),
+    i64.store(i32.add(local.get("$addr"), i32.const(8)), local.get("$value")),
+
+    local.get("$addr")
+  );
+
 export const nativeFunctions = [
   MAKE_INT_FX,
   MAKE_FLOAT_FX,
@@ -738,6 +753,7 @@ export const nativeFunctions = [
   GET_PAIR_TAIL_FX,
   SET_PAIR_HEAD_FX,
   SET_PAIR_TAIL_FX,
+  LOG_FX,
   NEG_FX,
   ARITHMETIC_OP_FX,
   STRING_COMPARE_FX,
@@ -746,5 +762,5 @@ export const nativeFunctions = [
   PRE_APPLY_FX,
   GET_LEX_ADDR_FX,
   SET_LEX_ADDR_FX,
-  LOG_FX,
+  SET_PARAM_FX,
 ];