Use lambda instead of enter/exit function

2 files changed, 38 insertions, 32 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 bc21afbf..41ae0ad3 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
@@ -30,34 +30,34 @@ public class WasmCodeGenClosures {
             name
         );
 
-        var functionContext = context.enterFunction(new WasmFuncType(
+        var closureFuncType = new WasmFuncType(
             List.of(CLOSURE_TYPE),
             List.of(WasmCodeGenTypes.VALUE_TYPE)
-        ));
-
-        var functionClosurePointerResult = functionContext.generateLocal(WasmCodeGenLocalSpace.CLOSURE_POINTER);
-        var functionClosurePointer = functionClosurePointerResult.value();
-        functionContext = functionClosurePointerResult.context();
-
-        for (var freeVariable : freeVariables) {
-            var freeVariableResult = functionContext.addShedLocal(freeVariable.name());
-            var freeVariableIndex = freeVariableResult.value();
-            functionContext = freeVariableResult.context();
-
-            functionContext = functionContext
-                .addInstruction(new WasmLocalGet(functionClosurePointer))
-                .addInstruction(new WasmI32Load(new WasmMemoryArg(
-                    // TODO: remove magic numbers
-                    4 + freeVariable.index() * 4,
-                    WasmAlignment.ALIGN_32_BIT
-                )))
-                .addInstruction(new WasmLocalSet(freeVariableIndex));
-        }
-
-        functionContext = compileBody.apply(functionContext);
-
-        var function = functionContext.toFunction(functionName);
-        context = context.exitFunction(functionContext).addFunction(function);
+        );
+        context = context.createFunction(functionName, closureFuncType, bodyContext -> {
+            var functionClosurePointerResult = bodyContext.generateLocal(WasmCodeGenLocalSpace.CLOSURE_POINTER);
+            var functionClosurePointer = functionClosurePointerResult.value();
+            bodyContext = functionClosurePointerResult.context();
+
+            for (var freeVariable : freeVariables) {
+                var freeVariableResult = bodyContext.addShedLocal(freeVariable.name());
+                var freeVariableIndex = freeVariableResult.value();
+                bodyContext = freeVariableResult.context();
+
+                bodyContext = bodyContext
+                    .addInstruction(new WasmLocalGet(functionClosurePointer))
+                    .addInstruction(new WasmI32Load(new WasmMemoryArg(
+                        // TODO: remove magic numbers
+                        4 + freeVariable.index() * 4,
+                        WasmAlignment.ALIGN_32_BIT
+                    )))
+                    .addInstruction(new WasmLocalSet(freeVariableIndex));
+            }
+
+            bodyContext = compileBody.apply(bodyContext);
+
+            return bodyContext;
+        });
 
         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 f8b28d19..4d9180cb 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
@@ -13,6 +13,7 @@ import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmTableType;
 
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Function;
 
 public record WasmCodeGenContext(
     ModuleContext moduleContext,
@@ -257,8 +258,12 @@ public record WasmCodeGenContext(
         return new WasmTableIndex(0);
     }
 
-    public WasmCodeGenContext enterFunction(WasmFuncType funcType) {
-        return new WasmCodeGenContext(
+    public WasmCodeGenContext createFunction(
+        String name,
+        WasmFuncType funcType,
+        Function<WasmCodeGenContext, WasmCodeGenContext> create
+    ) {
+        var bodyContext = new WasmCodeGenContext(
             this.moduleContext,
             new FunctionContext(
                 funcType,
@@ -267,12 +272,13 @@ public record WasmCodeGenContext(
                 Empty.map()
             )
         );
-    }
 
-    public WasmCodeGenContext exitFunction(WasmCodeGenContext functionContext) {
-        // TODO: merge properly (or throw if changed?) (or use lambda?)
+        bodyContext = create.apply(bodyContext);
+
+        var function = bodyContext.functionContext.toFunction(name);
+
         return new WasmCodeGenContext(
-            functionContext.moduleContext,
+            bodyContext.moduleContext.addFunction(function),
             this.functionContext
         );
     }