Store function pointer in closure

8 files changed, 45 insertions, 2 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 2630f63c..081075e8 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
@@ -3,6 +3,7 @@ package org.shedlang.compiler.backends.wasm;
 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.WasmFuncName;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmLocalIndex;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmValueType;
@@ -69,6 +70,11 @@ public class WasmCodeGenClosures {
         context = WasmCodeGenAlloc.callAlloc(closureSize, context);
         context = context.addInstruction(new WasmLocalSet(closurePointer));
 
+        context = context
+            .addInstruction(new WasmLocalGet(closurePointer))
+            .addInstruction(new WasmTableOffsetConst(new WasmFuncName(functionName)))
+            .addInstruction(new WasmI32Store(new WasmMemoryArg(0, WasmAlignment.ALIGN_32_BIT)));
+
         context = context.addInstruction(new WasmLocalGet(closurePointer));
 
         return context;
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmI32Store.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmI32Store.java
new file mode 100644
index 00000000..8703b8c9
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmI32Store.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.instructions;
+
+public record WasmI32Store(WasmMemoryArg memoryArg) implements WasmInstruction {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java
index e0bce7b5..d9d2db58 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java
@@ -1,4 +1,4 @@
 package org.shedlang.compiler.backends.wasm.structure.ast.instructions;
 
-public sealed interface WasmInstruction permits WasmCall, WasmCallIndirect, WasmDrop, WasmI32Const, WasmI32Load, WasmI64Const, WasmLocalGet, WasmLocalSet, WasmMemoryAddressConst, WasmReturn {
+public sealed interface WasmInstruction permits WasmCall, WasmCallIndirect, WasmDrop, WasmI32Const, WasmI32Load, WasmI32Store, WasmI64Const, WasmLocalGet, WasmLocalSet, WasmMemoryAddressConst, WasmReturn, WasmTableOffsetConst {
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmTableOffsetConst.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmTableOffsetConst.java
new file mode 100644
index 00000000..4045a8ab
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmTableOffsetConst.java
@@ -0,0 +1,11 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.instructions;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmFuncName;
+
+/**
+ * This is a fake instruction that corresponds to an i32.const instruction for
+ * a table index. Encapsulating this in its own instruction avoids having to
+ * separately track relocation.
+ */
+public record WasmTableOffsetConst(WasmFuncName funcName) implements WasmInstruction {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmRelocatableIndex.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmRelocatableIndex.java
index b5035b34..eec53af7 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmRelocatableIndex.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmRelocatableIndex.java
@@ -6,5 +6,6 @@ import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
 public sealed interface WasmRelocatableIndex {
     record FuncIndex(WasmFuncName funcName) implements WasmRelocatableIndex {}
     record MemoryAddress(WasmMemoryAddress address) implements WasmRelocatableIndex {}
+    record TableElementIndex(WasmFuncName funcName) implements WasmRelocatableIndex {}
     record TypeIndex(WasmFuncType funcType) implements WasmRelocatableIndex {}
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
index aa2fff35..7306d929 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
@@ -61,6 +61,11 @@ public class WasmBinaryFormatInstructions {
                 writeMemoryArg(i32Load.memoryArg(), writer);
             }
 
+            case WasmI32Store i32Store -> {
+                writer.writeByte(0x36);
+                writeMemoryArg(i32Store.memoryArg(), writer);
+            }
+
             // == Numeric Instructions ==
 
             case WasmI32Const i32Const -> {
@@ -78,6 +83,11 @@ public class WasmBinaryFormatInstructions {
                 writer.writeByte(0x41);
                 writer.writeMemoryAddressConst(memoryAddressConst.address());
             }
+
+            case WasmTableOffsetConst tableIndexConst -> {
+                writer.writeByte(0x41);
+                writer.writeTableElementIndexConst(tableIndexConst.funcName());
+            }
         }
     }
 
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatLinking.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatLinking.java
index 68a9dfed..451229fc 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatLinking.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatLinking.java
@@ -74,6 +74,9 @@ public class WasmBinaryFormatLinking {
 
             case WasmRelocatableIndex.TypeIndex typeIndex ->
                 symbolTable.funcTypeToTypeIndex(typeIndex.funcType()).value();
+
+            case WasmRelocatableIndex.TableElementIndex tableElementIndex ->
+                symbolTable.funcNameToSymbolIndex(tableElementIndex.funcName()).value();
         };
 
         writer.writeByte(relocationEntry.type().id());
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java
index 507600f9..0d4bc68f 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java
@@ -71,7 +71,15 @@ public class WasmBinaryWriter {
     }
 
     public void writeTableIndex(WasmTableIndex tableIndex) {
-        writeUnsignedInt32(0);
+        writeUnsignedInt32(tableIndex.value());
+    }
+
+    public void writeTableElementIndexConst(WasmFuncName funcName) throws IOException {
+        writeRelocatableIndex(
+            WasmRelocationType.TABLE_INDEX_SLEB,
+            () -> writePaddedSignedInt32(0),
+            new WasmRelocatableIndex.TableElementIndex(funcName)
+        );
     }
 
     public void writeTypeIndex(WasmFuncType funcType) throws IOException {