Store free variables in closure

1 files changed, 27 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
index 081075e8..8a5eb207 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
@@ -1,5 +1,6 @@
 package org.shedlang.compiler.backends.wasm;
 
+import org.shedlang.compiler.ast.typed.FreeVariable;
 import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.backends.wasm.structure.ast.instructions.*;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmAlignment;
@@ -48,11 +49,7 @@ public class WasmCodeGenClosures {
 
                 bodyContext = bodyContext
                     .addInstruction(new WasmLocalGet(closurePointer))
-                    .addInstruction(new WasmI32Load(new WasmMemoryArg(
-                        // TODO: remove magic numbers
-                        4 + freeVariable.index() * 4,
-                        WasmAlignment.ALIGN_32_BIT
-                    )))
+                    .addInstruction(new WasmI32Load(freeVariableMemoryArg(freeVariable)))
                     .addInstruction(new WasmLocalSet(freeVariableIndex));
             }
 
@@ -73,13 +70,36 @@ public class WasmCodeGenClosures {
         context = context
             .addInstruction(new WasmLocalGet(closurePointer))
             .addInstruction(new WasmTableOffsetConst(new WasmFuncName(functionName)))
-            .addInstruction(new WasmI32Store(new WasmMemoryArg(0, WasmAlignment.ALIGN_32_BIT)));
+            .addInstruction(new WasmI32Store(functionPointerMemoryArg()));
+
+        for (var freeVariable : freeVariables) {
+            context = context
+                .addInstruction(new WasmLocalGet(closurePointer))
+                .addInstruction(new WasmLocalGet(context.shedLocal(freeVariable.name())))
+                .addInstruction(new WasmI32Store(freeVariableMemoryArg(freeVariable)));
+        }
 
         context = context.addInstruction(new WasmLocalGet(closurePointer));
 
         return context;
     }
 
+    private static WasmMemoryArg functionPointerMemoryArg() {
+        return new WasmMemoryArg(
+            // TODO: remove magic numbers
+            0,
+            WasmAlignment.ALIGN_32_BIT
+        );
+    }
+
+    private static WasmMemoryArg freeVariableMemoryArg(FreeVariable freeVariable) {
+        return new WasmMemoryArg(
+            // TODO: remove magic numbers
+            4 + freeVariable.index() * 4,
+            WasmAlignment.ALIGN_32_BIT
+        );
+    }
+
     public static WasmCodeGenContext call(
         WasmLocalIndex closure,
         List<WasmLocalIndex> positionalArgs,
@@ -111,6 +131,6 @@ public class WasmCodeGenClosures {
     ) {
         return context
             .addInstruction(new WasmLocalGet(closure))
-            .addInstruction(new WasmI32Load(new WasmMemoryArg(0, WasmAlignment.ALIGN_32_BIT)));
+            .addInstruction(new WasmI32Load(functionPointerMemoryArg()));
     }
 }