Add proper namespace import for WASM

8 files changed, 136 insertions, 18 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
index 50c05a9c..179e6cf6 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
@@ -1,12 +1,13 @@
 package org.shedlang.compiler.backends.llvm;
 
 import org.shedlang.compiler.ast.typed.FreeVariables;
-import org.shedlang.compiler.backends.llvm.ir.ast.*;
+import org.shedlang.compiler.backends.llvm.ir.ast.LlvmFunctionDeclaration;
+import org.shedlang.compiler.backends.llvm.ir.ast.LlvmGlobalDefinition;
+import org.shedlang.compiler.backends.llvm.ir.ast.LlvmModule;
+import org.shedlang.compiler.backends.llvm.ir.ast.LlvmParam;
 import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmReturn;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperand;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperandInt;
-import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperandPtrToInt;
-import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableGlobal;
 import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypes;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.typechecker.TypeInfo;
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmBackend.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmBackend.java
index 48a8d5ee..dc6fd468 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmBackend.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmBackend.java
@@ -4,6 +4,7 @@ import org.shedlang.compiler.backends.Backend;
 import org.shedlang.compiler.backends.wasm.structure.WasmBinaryFormat;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmModule;
 import org.shedlang.compiler.config.ProjectConfig;
+import org.shedlang.compiler.filesystem.ShedCompilerFiles;
 import org.shedlang.compiler.stackir.instructions.Instruction;
 import org.shedlang.compiler.typechecker.TypeInfo;
 import org.shedlang.compiler.types.NamespaceName;
@@ -56,6 +57,8 @@ public class WasmBackend implements Backend {
             var mainObjectPath = writeObjectFile(mainModule, targetDir.resolve("bins"), bin.name());
             ldArgs.add(mainObjectPath.toString());
 
+            ldArgs.add(ShedCompilerFiles.getRoot().resolve("runtime/build-wasm/libshed.o").toString());
+
             run(ldArgs);
         }
     }
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 15bf531b..64835427 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
@@ -4,6 +4,7 @@ import org.shedlang.compiler.backends.wasm.structure.ast.instructions.*;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmLocalIndex;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmModule;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmValueType;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.typechecker.TypeInfo;
 import org.shedlang.compiler.types.NamespaceName;
@@ -22,10 +23,59 @@ public class WasmCodeGen {
     public static WasmModule generateMain(NamespaceName main, TypeInfo typeInfo) {
         var context = WasmCodeGenContext.fromTypeInfo(typeInfo);
 
+        // TODO: put these in a separate module
+        context = context.createFunction(
+            "wasm_memory_grow",
+            new WasmFuncType(List.of(WasmValueType.I32), List.of(WasmValueType.I32)),
+            bodyContext -> {
+                return bodyContext
+                    .generateLocal(WasmCodeGenLocalSpace.LOCAL, "pages").context()
+                    .addInstruction(new WasmLocalGet(new WasmLocalIndex(0)))
+                    .addInstruction(new WasmMemoryGrow());
+            }
+        );
+
         context = context.createFunction(
             "_start",
             new WasmFuncType(List.of(), List.of()),
-            initContext -> initContext
+            mainContext -> {
+                var namespaceResult = mainContext.generateLocal(
+                    WasmCodeGenLocalSpace.NAMESPACE,
+                    WasmCodeGenNamespaces.namespaceNameToWasmIdentifier(main)
+                );
+                var namespace = namespaceResult.value();
+                mainContext = namespaceResult.context();
+
+                mainContext = WasmCodeGenNamespaces.importValue(main, mainContext);
+                mainContext = mainContext.addInstruction(new WasmLocalSet(namespace));
+
+                var fieldResult = mainContext.generateLocal(
+                    WasmCodeGenLocalSpace.MAIN
+                );
+                var field = fieldResult.value();
+                mainContext = fieldResult.context();
+
+                mainContext = WasmCodeGenNamespaces.fieldLoad(
+                    main,
+                    namespace,
+                    // TODO: check there is a main function
+                    "main",
+                    mainContext
+                );
+                mainContext = mainContext.addInstruction(new WasmLocalSet(field));
+
+                // TODO: make code gen functions consistent in whether they
+                //  return locals or put values on the stack
+                mainContext = WasmCodeGenClosures.call(
+                    field,
+                    List.of(),
+                    List.of(),
+                    mainContext
+                );
+                mainContext = mainContext.addInstruction(new WasmDrop());
+
+                return mainContext;
+            }
         );
 
         return context.toModule();
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 aa130318..71c90174 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java
@@ -7,6 +7,8 @@ public enum WasmCodeGenLocalSpace {
     CLOSURE_FREE_VARIABLE("closureFreeVariable"),
     FIELD_VALUE("fieldValue"),
     LOCAL("local"),
+    MAIN("main"),
+    NAMESPACE("namespace"),
     RECEIVER("receiver");
 
     private final String name;
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenNamespaces.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenNamespaces.java
index a2faaacd..f2fcbebf 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenNamespaces.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenNamespaces.java
@@ -1,12 +1,8 @@
 package org.shedlang.compiler.backends.wasm;
 
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmI32Store;
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmLocalGet;
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmMemoryAddressConst;
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmMemoryArg;
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.*;
 import org.shedlang.compiler.backends.wasm.structure.ast.linking.WasmMemoryAddress;
-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.modules.*;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
 import org.shedlang.compiler.types.NamespaceName;
 
@@ -17,6 +13,23 @@ public class WasmCodeGenNamespaces {
     private WasmCodeGenNamespaces() {
     }
 
+    public static WasmCodeGenContext importValue(
+        NamespaceName namespaceName,
+        WasmCodeGenContext context
+    ) {
+        var importFuncName = new WasmFuncName(globalNameNamespaceImport(namespaceName));
+        return context
+            .addImport(new WasmImport(
+                "env",
+                importFuncName.value(),
+                new WasmImportDescriptor.Func(
+                    importFuncName.value(),
+                    new WasmFuncType(List.of(), List.of(WasmCodeGenTypes.VALUE_TYPE))
+                )
+            ))
+            .addInstruction(new WasmCall(importFuncName));
+    }
+
     public static WasmCodeGenContext define(
         NamespaceName namespaceName,
         Function<WasmCodeGenContext, WasmCodeGenContext> generateInitBody,
@@ -35,6 +48,15 @@ public class WasmCodeGenNamespaces {
             globalNameNamespaceInit(namespaceName),
             new WasmFuncType(List.of(), List.of()),
             generateInitBody
+        ).createFunction(
+            globalNameNamespaceImport(namespaceName),
+            new WasmFuncType(List.of(), List.of(WasmCodeGenTypes.VALUE_TYPE)),
+            bodyContext -> {
+                // TODO: only init when necessary
+                return bodyContext
+                    .addInstruction(new WasmCall(new WasmFuncName(globalNameNamespaceInit(namespaceName))))
+                    .addInstruction(new WasmMemoryAddressConst(namespaceMemoryAddress(namespaceName, bodyContext)));
+            }
         );
     }
 
@@ -44,17 +66,30 @@ public class WasmCodeGenNamespaces {
         WasmLocalIndex fieldValue,
         WasmCodeGenContext context
     ) {
-        var namespaceTypeInfo = context.typeInfo().namespaceTypeInfo(namespaceName);
-
         var namespaceAddress = namespaceMemoryAddress(namespaceName, context);
 
         return context
             .addInstruction(new WasmMemoryAddressConst(namespaceAddress))
             .addInstruction(new WasmLocalGet(fieldValue))
-            .addInstruction(new WasmI32Store(new WasmMemoryArg(
-                // TODO: remove magic number
-                namespaceTypeInfo.fieldIndex(fieldName) * 4,
-                WasmAlignment.ALIGN_32_BIT
+            .addInstruction(new WasmI32Store(fieldMemoryArg(
+                namespaceName,
+                fieldName,
+                context
+            )));
+    }
+
+    public static WasmCodeGenContext fieldLoad(
+        NamespaceName namespaceName,
+        WasmLocalIndex namespace,
+        String fieldName,
+        WasmCodeGenContext context
+    ) {
+        return context
+            .addInstruction(new WasmLocalGet(namespace))
+            .addInstruction(new WasmI32Load(fieldMemoryArg(
+                namespaceName,
+                fieldName,
+                context
             )));
     }
 
@@ -63,6 +98,20 @@ public class WasmCodeGenNamespaces {
         return context.memoryAddress(dataName);
     }
 
+    private static WasmMemoryArg fieldMemoryArg(
+        NamespaceName namespaceName,
+        String fieldName,
+        WasmCodeGenContext context
+    ) {
+        var namespaceTypeInfo = context.typeInfo().namespaceTypeInfo(namespaceName);
+
+        return new WasmMemoryArg(
+            // TODO: remove magic number
+            namespaceTypeInfo.fieldIndex(fieldName) * 4,
+            WasmAlignment.ALIGN_32_BIT
+        );
+    }
+
     private static String globalNameNamespaceValue(NamespaceName namespaceName) {
         return "shed__namespace__" + namespaceNameToWasmIdentifier(namespaceName);
     }
@@ -71,7 +120,11 @@ public class WasmCodeGenNamespaces {
         return "shed__namespaceInit__" + namespaceNameToWasmIdentifier(namespaceName);
     }
 
-    private static String namespaceNameToWasmIdentifier(NamespaceName namespaceName) {
+    private static String globalNameNamespaceImport(NamespaceName namespaceName) {
+        return "shed__namespaceImport__" + namespaceNameToWasmIdentifier(namespaceName);
+    }
+
+    public static String namespaceNameToWasmIdentifier(NamespaceName namespaceName) {
         return String.join("_", namespaceName.parts());
     }
 }
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 46f7e0e7..ed1a1cf9 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, WasmI32Store, WasmI64Const, WasmLocalGet, WasmLocalSet, WasmMemoryAddressConst, WasmReturn, WasmTableElementIndexConst {
+public sealed interface WasmInstruction permits WasmCall, WasmCallIndirect, WasmDrop, WasmI32Const, WasmI32Load, WasmI32Store, WasmI64Const, WasmLocalGet, WasmLocalSet, WasmMemoryAddressConst, WasmMemoryGrow, WasmReturn, WasmTableElementIndexConst {
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmMemoryGrow.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmMemoryGrow.java
new file mode 100644
index 00000000..0b184fe9
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmMemoryGrow.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.instructions;
+
+public record WasmMemoryGrow() implements WasmInstruction {
+}
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 77c54489..45403459 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
@@ -66,6 +66,11 @@ public class WasmBinaryFormatInstructions {
                 writeMemoryArg(i32Store.memoryArg(), writer);
             }
 
+            case WasmMemoryGrow memoryGrow -> {
+                writer.writeByte(0x40);
+                writer.writeByte(0x00);
+            }
+
             // == Numeric Instructions ==
 
             case WasmI32Const i32Const -> {