Support hello-world example in wasm-patio example tests

1 files changed, 147 insertions, 41 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
index c884eee6..aa687d63 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
@@ -1,6 +1,7 @@
 package org.shedlang.compiler.backends.wasm;
 
 import org.pcollections.PMap;
+import org.pcollections.PStack;
 import org.pcollections.PVector;
 import org.shedlang.compiler.backends.wasm.structure.ast.instructions.*;
 import org.shedlang.compiler.backends.wasm.structure.ast.linking.WasmMemoryAddress;
@@ -14,6 +15,8 @@ import org.shedlang.compiler.patio.ast.types.*;
 import org.shedlang.compiler.util.Lists;
 import org.shedlang.compiler.util.P;
 
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.util.List;
 
 public class WasmPatioCodeGen {
@@ -36,7 +39,11 @@ public class WasmPatioCodeGen {
         for (var patioCompilationUnitElement : patioCompilationUnit.elements()) {
             switch (patioCompilationUnitElement) {
                 case PatioDataElementImmutable patioDataElementImmutable -> {
-                    throw new RuntimeException("TODO");
+                    wasmMemoryBuilder = wasmMemoryBuilder.add(
+                        patioDataElementImmutable.name(),
+                        patioDataValueToBytes(patioDataElementImmutable.value()),
+                        WasmAlignment.ALIGN_32_BIT
+                    ).memory();
                 }
                 case PatioDataElementMutable patioDataElementMutable -> {
                     wasmMemoryBuilder = wasmMemoryBuilder.add(
@@ -74,18 +81,19 @@ public class WasmPatioCodeGen {
 
                     var wasmResultTypes = patioTypeToWasmResultTypes(patioFunctionDefinition.returnType());
 
-                    var context = Context.forFunction(Lists.concat(wasmParams, wasmLocals));
-                    var wasmBody = generateInstructions(
+                    var context = generateInstructions(
                         patioFunctionDefinition.body(),
-                        context.clearInstructions()
+                        Context.forFunction(module, Lists.concat(wasmParams, wasmLocals))
                     );
 
+                    module = context.module;
+
                     module = module.addFunction(new WasmFunction(
                         name,
                         wasmParams,
                         wasmLocals,
                         wasmResultTypes,
-                        new WasmExpression(wasmBody)
+                        new WasmExpression(context.instructions)
                     ));
                 }
             }
@@ -96,7 +104,7 @@ public class WasmPatioCodeGen {
         return module;
     }
 
-    private static PVector<WasmInstruction> generateInstructions(
+    private static Context generateInstructions(
         List<PatioInstruction> patioInstructions,
         Context context
     ) {
@@ -104,7 +112,7 @@ public class WasmPatioCodeGen {
             context = generateInstruction(patioInstruction, context);
         }
 
-        return context.instructions;
+        return context;
     }
 
     private static Context generateInstruction(
@@ -149,6 +157,8 @@ public class WasmPatioCodeGen {
                     List.of(WasmValueType.I32)
                 );
 
+                context = context.addType(wasmFuncType);
+
                 context = context.addInstruction(new WasmCallIndirect(
                     // TODO: remove hard-coded zero
                     new WasmTableIndex(0),
@@ -171,10 +181,20 @@ public class WasmPatioCodeGen {
             case PatioIf patioIf -> {
                 context = generateOperand(patioIf.condition(), context);
 
+                context = context.enterBlock();
+                context = generateInstructions(patioIf.trueBody(), context);
+                var wasmTrueBody = context.instructions;
+                context = context.exitBlock();
+
+                context = context.enterBlock();
+                context = generateInstructions(patioIf.falseBody(), context);
+                var wasmFalseBody = context.instructions;
+                context = context.exitBlock();
+
                 context = context.addInstruction(new WasmIf(
                     new WasmBlockType(List.of()),
-                    generateInstructions(patioIf.trueBody(), context.clearInstructions()),
-                    generateInstructions(patioIf.falseBody(), context.clearInstructions())
+                    wasmTrueBody,
+                    wasmFalseBody
                 ));
 
                 yield context;
@@ -240,34 +260,6 @@ public class WasmPatioCodeGen {
         };
     }
 
-    private static WasmMemoryArg wasmMemoryArg(
-        PatioDataLayout patioDataLayout,
-        List<Integer> patioOffsets
-    ) {
-        var wasmOffset = 0;
-
-        for (var patioOffset : patioOffsets) {
-            switch (patioDataLayout) {
-                case PatioDataLayoutArray patioDataLayoutArray -> {
-                    wasmOffset += sizeOf(patioDataLayoutArray.element()) * patioOffset;
-                    patioDataLayout = patioDataLayoutArray.element();
-                }
-                case PatioDataLayoutScalar patioDataLayoutScalar -> {
-                    throw new RuntimeException("offset not allowed for scalars");
-                }
-                case PatioDataLayoutStruct patioDataLayoutStruct -> {
-                    for (var i = 0; i < patioOffset; i++) {
-                        wasmOffset += sizeOf(patioDataLayoutStruct.elements().get(i));
-                    }
-                    patioDataLayout = patioDataLayoutStruct.elements().get(patioOffset);
-                }
-            }
-        }
-
-        // TODO: alignment
-        return new WasmMemoryArg(wasmOffset, WasmAlignment.NONE);
-    }
-
     private static Context generateOperand(
         PatioOperand patioOperand,
         Context context
@@ -311,6 +303,84 @@ public class WasmPatioCodeGen {
         };
     }
 
+    private static WasmMemoryArg wasmMemoryArg(
+        PatioDataLayout patioDataLayout,
+        List<Integer> patioOffsets
+    ) {
+        var wasmOffset = 0;
+
+        for (var patioOffset : patioOffsets) {
+            switch (patioDataLayout) {
+                case PatioDataLayoutArray patioDataLayoutArray -> {
+                    wasmOffset += sizeOf(patioDataLayoutArray.element()) * patioOffset;
+                    patioDataLayout = patioDataLayoutArray.element();
+                }
+                case PatioDataLayoutScalar patioDataLayoutScalar -> {
+                    throw new RuntimeException("offset not allowed for scalars");
+                }
+                case PatioDataLayoutStruct patioDataLayoutStruct -> {
+                    for (var i = 0; i < patioOffset; i++) {
+                        wasmOffset += sizeOf(patioDataLayoutStruct.elements().get(i));
+                    }
+                    patioDataLayout = patioDataLayoutStruct.elements().get(patioOffset);
+                }
+            }
+        }
+
+        // TODO: alignment
+        return new WasmMemoryArg(wasmOffset, WasmAlignment.NONE);
+    }
+
+    private static byte[] patioDataValueToBytes(PatioDataValue value) {
+        var bytes = ByteBuffer.allocate(sizeOf(value));
+        bytes.order(ByteOrder.LITTLE_ENDIAN);
+
+        patioDataValueToBytes(value, bytes);
+
+        return bytes.array();
+    }
+
+    private static void patioDataValueToBytes(PatioDataValue value, ByteBuffer bytes) {
+        switch (value) {
+            case PatioDataValueArray patioDataValueArray -> {
+                for (var element : patioDataValueArray.elements()) {
+                    patioDataValueToBytes(element, bytes);
+                }
+            }
+            case PatioDataValueByte patioDataValueByte -> {
+                bytes.put(patioDataValueByte.value());
+            }
+            case PatioDataValueInt patioDataValueInt -> {
+                bytes.putInt(patioDataValueInt.value());
+            }
+            case PatioDataValueStruct patioDataValueStruct -> {
+                for (var element : patioDataValueStruct.elements()) {
+                    patioDataValueToBytes(element, bytes);
+                }
+            }
+        }
+    }
+
+    private static int sizeOf(PatioDataValue value) {
+        return switch (value) {
+            case PatioDataValueArray patioDataValueArray ->
+                patioDataValueArray.elements().stream()
+                    .map(element -> sizeOf(element))
+                    .reduce(0, Integer::sum);
+
+            case PatioDataValueByte patioDataValueByte ->
+                1;
+
+            case PatioDataValueInt patioDataValueInt ->
+                4;
+
+            case PatioDataValueStruct patioDataValueStruct ->
+                patioDataValueStruct.elements().stream()
+                    .map(element -> sizeOf(element))
+                    .reduce(0, Integer::sum);
+        };
+    }
+
     private static WasmLocal patioLocalToWasmLocal(PatioLocal local) {
         return new WasmLocal(local.name(), patioTypeToWasmValueType(local.type()));
     }
@@ -324,7 +394,7 @@ public class WasmPatioCodeGen {
     }
 
     private static class Context {
-        public static Context forFunction(List<WasmLocal> locals) {
+        public static Context forFunction(WasmModule module, List<WasmLocal> locals) {
             var patioLocalNameToWasmLocalIndex = P.<String, WasmLocalIndex>map();
 
             for (var local : locals) {
@@ -335,32 +405,68 @@ public class WasmPatioCodeGen {
             }
 
             return new Context(
+                module,
                 P.vector(),
+                P.stack(),
                 patioLocalNameToWasmLocalIndex
             );
         }
 
+        private final WasmModule module;
         private final PVector<WasmInstruction> instructions;
+        private final PStack<PVector<WasmInstruction>> instructionsStack;
         private final PMap<String, WasmLocalIndex> patioLocalNameToWasmLocalIndex;
 
         private Context(
+            WasmModule module,
             PVector<WasmInstruction> instructions,
+            PStack<PVector<WasmInstruction>> instructionsStack,
             PMap<String, WasmLocalIndex> patioLocalNameToWasmLocalIndex
         ) {
+            this.module = module;
             this.instructions = instructions;
+            this.instructionsStack = instructionsStack;
             this.patioLocalNameToWasmLocalIndex = patioLocalNameToWasmLocalIndex;
         }
 
+        private Context addType(WasmFuncType type) {
+            return new Context(
+                this.module.addType(type),
+                this.instructions,
+                this.instructionsStack,
+                this.patioLocalNameToWasmLocalIndex
+            );
+        }
+
         private Context addInstruction(WasmInstruction instruction) {
-            return new Context(this.instructions.plus(instruction), this.patioLocalNameToWasmLocalIndex);
+            return new Context(
+                this.module,
+                this.instructions.plus(instruction),
+                this.instructionsStack,
+                this.patioLocalNameToWasmLocalIndex
+            );
         }
 
         private WasmLocalIndex wasmLocalIndex(PatioLocalRef patioLocalRef) {
             return this.patioLocalNameToWasmLocalIndex.get(patioLocalRef.name());
         }
 
-        public Context clearInstructions() {
-            return new Context(P.vector(), this.patioLocalNameToWasmLocalIndex);
+        public Context enterBlock() {
+            return new Context(
+                this.module,
+                P.vector(),
+                this.instructionsStack.plus(this.instructions),
+                this.patioLocalNameToWasmLocalIndex
+            );
+        }
+
+        public Context exitBlock() {
+            return new Context(
+                this.module,
+                this.instructionsStack.get(0),
+                this.instructionsStack.minus(0),
+                this.patioLocalNameToWasmLocalIndex
+            );
         }
     }
 }