Allocate memory for closures

3 files changed, 91 insertions, 16 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenAlloc.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenAlloc.java
new file mode 100644
index 00000000..9d74e60c
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenAlloc.java
@@ -0,0 +1,32 @@
+package org.shedlang.compiler.backends.wasm;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmCall;
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmI32Const;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmFuncName;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmImport;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmImportDescriptor;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmValueType;
+
+import java.util.List;
+
+public class WasmCodeGenAlloc {
+    private WasmCodeGenAlloc() {
+    }
+
+    public static WasmCodeGenContext callAlloc(int size, WasmCodeGenContext context) {
+        var allocName = "shed__runtime__alloc";
+
+        return context
+            .addImport(new WasmImport(
+                "env",
+                "shed__runtime__alloc",
+                new WasmImportDescriptor.Func(
+                    allocName,
+                    new WasmFuncType(List.of(WasmValueType.I32), List.of(WasmValueType.I32))
+                )
+            ))
+            .addInstruction(new WasmI32Const(size))
+            .addInstruction(new WasmCall(new WasmFuncName(allocName)));
+    }
+}
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 41ae0ad3..2630f63c 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
@@ -34,10 +34,11 @@ public class WasmCodeGenClosures {
             List.of(CLOSURE_TYPE),
             List.of(WasmCodeGenTypes.VALUE_TYPE)
         );
+
         context = context.createFunction(functionName, closureFuncType, bodyContext -> {
-            var functionClosurePointerResult = bodyContext.generateLocal(WasmCodeGenLocalSpace.CLOSURE_POINTER);
-            var functionClosurePointer = functionClosurePointerResult.value();
-            bodyContext = functionClosurePointerResult.context();
+            var closurePointerResult = bodyContext.generateLocal(WasmCodeGenLocalSpace.CLOSURE_POINTER);
+            var closurePointer = closurePointerResult.value();
+            bodyContext = closurePointerResult.context();
 
             for (var freeVariable : freeVariables) {
                 var freeVariableResult = bodyContext.addShedLocal(freeVariable.name());
@@ -45,7 +46,7 @@ public class WasmCodeGenClosures {
                 bodyContext = freeVariableResult.context();
 
                 bodyContext = bodyContext
-                    .addInstruction(new WasmLocalGet(functionClosurePointer))
+                    .addInstruction(new WasmLocalGet(closurePointer))
                     .addInstruction(new WasmI32Load(new WasmMemoryArg(
                         // TODO: remove magic numbers
                         4 + freeVariable.index() * 4,
@@ -59,6 +60,17 @@ public class WasmCodeGenClosures {
             return bodyContext;
         });
 
+        var closurePointerResult = context.generateLocal(WasmCodeGenLocalSpace.CLOSURE_POINTER);
+        var closurePointer = closurePointerResult.value();
+        context = closurePointerResult.context();
+
+        // TODO: remove magic numbers
+        var closureSize = 4 + freeVariables.size() * 4;
+        context = WasmCodeGenAlloc.callAlloc(closureSize, context);
+        context = context.addInstruction(new WasmLocalSet(closurePointer));
+
+        context = context.addInstruction(new WasmLocalGet(closurePointer));
+
         return context;
     }
 
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
index 4d9180cb..82263a13 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
@@ -21,39 +21,57 @@ public record WasmCodeGenContext(
 ) {
     record ModuleContext(
         OrderedPSet<WasmFuncType> types,
-        OrderedPMap<String, WasmFunction> functionsByName,
+        OrderedPSet<WasmImport> imports,
+        PVector<WasmFunction> functions,
+        PSet<String> symbolNames,
         WasmMemoryBuilder memory
     ) {
         public static final ModuleContext EMPTY = new ModuleContext(
             OrderedPSet.empty(),
-            OrderedPMap.empty(),
+            OrderedPSet.empty(),
+            Empty.vector(),
+            Empty.set(),
             WasmMemoryBuilder.EMPTY
         );
 
         public ModuleContext addType(WasmFuncType funcType) {
             return new ModuleContext(
                 this.types.plus(funcType),
-                this.functionsByName,
+                this.imports,
+                this.functions,
+                this.symbolNames,
                 this.memory
             );
         }
 
-        public String generateFunctionName(WasmCodeGenGlobalSpace space, String name) {
-            return NameAllocator.generateName(
-                functionsByName::containsKey,
-                "shed__" + space,
-                Optional.of(name)
+        public ModuleContext addImport(WasmImport import_) {
+            return new ModuleContext(
+                this.types,
+                this.imports.plus(import_),
+                this.functions,
+                this.symbolNames,
+                this.memory
             );
         }
 
         public ModuleContext addFunction(WasmFunction function) {
             return new ModuleContext(
                 this.types.plus(function.type()),
-                this.functionsByName.plus(function.name(), function),
+                this.imports,
+                this.functions.plus(function),
+                this.symbolNames.plus(function.name()),
                 this.memory
             );
         }
 
+        public String generateSymbolName(WasmCodeGenGlobalSpace space, String name) {
+            return NameAllocator.generateName(
+                symbolNames::contains,
+                "shed__" + space,
+                Optional.of(name)
+            );
+        }
+
         public Result<WasmMemoryAddress> addData(
             WasmCodeGenGlobalSpace space,
             byte[] data,
@@ -63,7 +81,9 @@ public record WasmCodeGenContext(
 
             return new ModuleContext(
                 this.types,
-                this.functionsByName,
+                this.imports,
+                this.functions,
+                this.symbolNames,
                 result.memory()
             ).toResult(result.offset());
         }
@@ -81,11 +101,15 @@ public record WasmCodeGenContext(
                 ))
             ));
 
+            for (var import_ : imports) {
+                module = module.addImport(import_);
+            }
+
             for (var type : types) {
                 module = module.addType(type);
             }
 
-            for (var function : functionsByName.values()) {
+            for (var function : functions) {
                 module = module.addFunction(function);
             }
 
@@ -193,8 +217,15 @@ public record WasmCodeGenContext(
         );
     }
 
+    public WasmCodeGenContext addImport(WasmImport import_) {
+        return new WasmCodeGenContext(
+            this.moduleContext.addImport(import_),
+            this.functionContext
+        );
+    }
+
     public String generateFunctionName(WasmCodeGenGlobalSpace space, String name) {
-        return moduleContext.generateFunctionName(space, name);
+        return moduleContext.generateSymbolName(space, name);
     }
 
     public WasmCodeGenContext addFunction(WasmFunction function) {