Generate WASM functions for Shed functions

7 files changed, 158 insertions, 43 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
index 90d72969..aedb2c55 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
@@ -218,6 +218,7 @@ public record LlvmCodeGenContext(
     public LlvmCodeGenResult<List<LlvmBasicBlock>> exitFunction(
         LlvmCodeGenContext functionContext
     ) {
+        // TODO: merge top-level entities
         return new LlvmCodeGenContext(
             typeInfo,
             stack,
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
index 5926af9c..2a380d35 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
@@ -59,9 +59,19 @@ public class WasmCodeGen {
                     context
                 );
             }
-            case DefineFunctionInstruction defineFunctionInstruction -> {
-                throw new UnsupportedOperationException("TODO");
-            }
+            case DefineFunctionInstruction defineFunction ->
+                WasmCodeGenClosures.create(
+                    defineFunction.name(),
+                    defineFunction.freeVariables(),
+                    functionContext -> {
+                        for (var bodyInstruction : defineFunction.body()) {
+                            functionContext = generateInstruction(bodyInstruction, functionContext);
+                        }
+                        return functionContext;
+                    },
+                    context
+                );
+
             case DiscardInstruction discard ->
                 context.addInstruction(new WasmDrop());
 
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 a2d34df7..c1a56d73 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
@@ -1,9 +1,7 @@
 package org.shedlang.compiler.backends.wasm;
 
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmCallIndirect;
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmI32Load;
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmLocalGet;
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmMemoryArg;
+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;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmLocalIndex;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
@@ -11,6 +9,7 @@ import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmValueType;
 
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 
 import static org.shedlang.compiler.util.Iterables.sortedBy;
 
@@ -20,6 +19,48 @@ public class WasmCodeGenClosures {
 
     private static final WasmValueType CLOSURE_TYPE = WasmValueType.I32;
 
+    public static WasmCodeGenContext create(
+        String name,
+        FreeVariables freeVariables,
+        Function<WasmCodeGenContext, WasmCodeGenContext> compileBody,
+        WasmCodeGenContext context
+    ) {
+        var functionName = context.generateFunctionName(
+            WasmCodeGenGlobalSpace.FUNCTION,
+            name
+        );
+
+        var functionContext = context.enterFunction(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);
+
+        context = context.exitFunction(functionContext.toFunction(functionName));
+
+        return context;
+    }
+
     public static WasmCodeGenContext call(
         WasmLocalIndex closure,
         List<WasmLocalIndex> positionalArgs,
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 37396d2f..b7120624 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
@@ -6,23 +6,28 @@ import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmExpres
 import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmInstruction;
 import org.shedlang.compiler.backends.wasm.structure.ast.linking.WasmMemoryAddress;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.*;
-import org.shedlang.compiler.backends.wasm.structure.ast.types.*;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmLimits;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmReferenceType;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmTableType;
 
 import java.util.List;
 import java.util.Optional;
 
 public record WasmCodeGenContext(
     OrderedPSet<WasmFuncType> types,
-    PVector<WasmFunction> functions,
+    OrderedPMap<String, WasmFunction> functionsByName,
     WasmMemoryBuilder memory,
+    WasmFuncType funcType,
     PVector<WasmInstruction> instructions,
     OrderedPMap<String, WasmLocal> locals,
     PMap<String, WasmLocalIndex> shedLocalNameToLocalIndex
 ) {
     public static final WasmCodeGenContext EMPTY = new WasmCodeGenContext(
         OrderedPSet.empty(),
-        Empty.vector(),
+        OrderedPMap.empty(),
         WasmMemoryBuilder.EMPTY,
+        new WasmFuncType(List.of(), List.of()),
         Empty.vector(),
         OrderedPMap.empty(),
         Empty.map()
@@ -30,12 +35,13 @@ public record WasmCodeGenContext(
 
     public WasmCodeGenContext addType(WasmFuncType funcType) {
         return new WasmCodeGenContext(
-            types.plus(funcType),
-            functions,
-            memory,
-            instructions,
-            locals,
-            shedLocalNameToLocalIndex
+            this.types.plus(funcType),
+            this.functionsByName,
+            this.memory,
+            this.funcType,
+            this.instructions,
+            this.locals,
+            this.shedLocalNameToLocalIndex
         );
     }
 
@@ -47,38 +53,41 @@ public record WasmCodeGenContext(
         var result = memory.add(space, data, alignment);
 
         return new WasmCodeGenContext(
-            types,
-            functions,
+            this.types,
+            this.functionsByName,
             result.memory(),
-            instructions,
-            locals,
-            shedLocalNameToLocalIndex
+            this.funcType,
+            this.instructions,
+            this.locals,
+            this.shedLocalNameToLocalIndex
         ).toResult(result.offset());
     }
 
     public WasmCodeGenContext addInstruction(WasmInstruction instruction) {
         return new WasmCodeGenContext(
-            types,
-            functions,
-            memory,
-            instructions.plus(instruction),
-            locals,
-            shedLocalNameToLocalIndex
+            this.types,
+            this.functionsByName,
+            this.memory,
+            this.funcType,
+            this.instructions.plus(instruction),
+            this.locals,
+            this.shedLocalNameToLocalIndex
         );
     }
 
-    public WasmCodeGenContext toFunction() {
+    public WasmCodeGenContext toFunction(String name) {
         var function = new WasmFunction(
-            "f",
-            List.of(),
-            locals.values().stream().toList(),
-            List.of(),
+            name,
+            locals.values().stream().limit(funcType.params().size()).toList(),
+            locals.values().stream().skip(funcType.params().size()).toList(),
+            funcType.results(),
             new WasmExpression(instructions)
         );
         return new WasmCodeGenContext(
-            types.plus(function.type()),
-            functions.plus(function),
-            memory,
+            this.types.plus(function.type()),
+            this.functionsByName.plus(function.name(), function),
+            this.memory,
+            new WasmFuncType(List.of(), List.of()),
             Empty.vector(),
             OrderedPMap.empty(),
             Empty.map()
@@ -102,7 +111,7 @@ public record WasmCodeGenContext(
             module = module.addType(type);
         }
 
-        for (var function : functions) {
+        for (var function : functionsByName.values()) {
             module = module.addFunction(function);
         }
 
@@ -116,8 +125,9 @@ public record WasmCodeGenContext(
 
         return new WasmCodeGenContext(
             context.types,
-            context.functions,
+            context.functionsByName,
             context.memory,
+            context.funcType,
             context.instructions,
             context.locals,
             context.shedLocalNameToLocalIndex.plus(name, localIndex)
@@ -150,19 +160,56 @@ public record WasmCodeGenContext(
             name
         );
         return new WasmCodeGenContext(
-            types,
-            functions,
-            memory,
-            instructions,
-            locals.plus(localName, new WasmLocal(localName, WasmCodeGenTypes.VALUE_TYPE)),
-            shedLocalNameToLocalIndex
+            this.types,
+            this.functionsByName,
+            this.memory,
+            this.funcType,
+            this.instructions,
+            this.locals.plus(
+                localName,
+                new WasmLocal(localName, WasmCodeGenTypes.VALUE_TYPE)
+            ),
+            this.shedLocalNameToLocalIndex
         ).toResult(localIndex);
     }
 
+    public String generateFunctionName(WasmCodeGenGlobalSpace space, String name) {
+        return NameAllocator.generateName(
+            functionsByName::containsKey,
+            "shed__" + space,
+            Optional.of(name)
+        );
+    }
+
     public WasmTableIndex indirectFunctionTableIndex() {
         return new WasmTableIndex(0);
     }
 
+    public WasmCodeGenContext enterFunction(WasmFuncType funcType) {
+        return new WasmCodeGenContext(
+            this.types,
+            this.functionsByName,
+            this.memory,
+            funcType,
+            Empty.vector(),
+            OrderedPMap.empty(),
+            Empty.map()
+        );
+    }
+
+    public WasmCodeGenContext exitFunction(WasmCodeGenContext functionContext) {
+        // TODO: merge properly (or throw if changed?) (or use lambda?)
+        return new WasmCodeGenContext(
+            functionContext.types,
+            functionContext.functionsByName,
+            functionContext.memory,
+            this.funcType,
+            this.instructions,
+            this.locals,
+            this.shedLocalNameToLocalIndex
+        );
+    }
+
     public <T> WasmCodeGenResult<T> toResult(T value) {
         return new WasmCodeGenResult<>(value, this);
     }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java
index 4b1b5293..97e51b68 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java
@@ -1,6 +1,7 @@
 package org.shedlang.compiler.backends.wasm;
 
 public enum WasmCodeGenGlobalSpace {
+    FUNCTION("function"),
     DATA_STRING("string");
 
     private final String name;
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java
index eeda312a..83bf9459 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java
@@ -3,6 +3,8 @@ package org.shedlang.compiler.backends.wasm;
 public enum WasmCodeGenLocalSpace {
     CALL_ARG_POSITIONAL("callArgPositional"),
     CALL_ARG_NAMED("callArgNamed"),
+    CLOSURE_POINTER("closurePointer"),
+    CLOSURE_FREE_VARIABLE("closureFreeVariable"),
     LOCAL("local"),
     RECEIVER("receiver");
 
diff --git a/src/main/java/org/shedlang/compiler/util/Iterables.java b/src/main/java/org/shedlang/compiler/util/Iterables.java
index 38730594..ef392893 100644
--- a/src/main/java/org/shedlang/compiler/util/Iterables.java
+++ b/src/main/java/org/shedlang/compiler/util/Iterables.java
@@ -1,5 +1,6 @@
 package org.shedlang.compiler.util;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.List;
@@ -17,6 +18,18 @@ public class Iterables {
         void call() throws E;
     }
 
+    public static <T, U> List<U> filterToType(Iterable<T> iterable, Class<U> clazz) {
+        var result = new ArrayList<U>();
+
+        for (var element : iterable) {
+            if (clazz.isInstance(element)) {
+                result.add(clazz.cast(element));
+            }
+        }
+
+        return result;
+    }
+
     public static <T, E extends Exception> void forEachInterspersed(
         Iterable<T> values,
         OnElement<T, E> onElement,