Bodge together code gen from Patio to Wasm

2 files changed, 363 insertions, 0 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
new file mode 100644
index 00000000..78fa39ba
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
@@ -0,0 +1,356 @@
+package org.shedlang.compiler.backends.wasm;
+
+import org.pcollections.PMap;
+import org.pcollections.PVector;
+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.*;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmValueType;
+import org.shedlang.compiler.patio.ast.*;
+import org.shedlang.compiler.patio.ast.instructions.*;
+import org.shedlang.compiler.patio.ast.operands.*;
+import org.shedlang.compiler.patio.ast.types.*;
+import org.shedlang.compiler.util.Lists;
+import org.shedlang.compiler.util.P;
+
+import java.util.List;
+
+public class WasmPatioCodeGen {
+    public static final String ALLOC_NAME = "shed__runtime__alloc";
+
+    public static WasmModule generateCompilationUnit(
+        PatioCompilationUnit patioCompilationUnit
+    ) {
+        var module = WasmModule.EMPTY.addImport(new WasmImport(
+            "env",
+            ALLOC_NAME,
+            new WasmImportDescriptor.Func(
+                ALLOC_NAME,
+                new WasmFuncType(List.of(WasmValueType.I32), List.of(WasmValueType.I32))
+            )
+        ));
+
+        var wasmMemoryBuilder = WasmMemoryBuilder.EMPTY;
+
+        for (var patioCompilationUnitElement : patioCompilationUnit.elements()) {
+            switch (patioCompilationUnitElement) {
+                case PatioDataElement patioDataElement -> {
+                    wasmMemoryBuilder = wasmMemoryBuilder.add(
+                        patioDataElement.name(),
+                        new byte[sizeOf(patioDataElement.layout())],
+                        WasmAlignment.ALIGN_32_BIT
+                    ).memory();
+                }
+                case PatioFunctionDeclaration patioFunctionDeclaration -> {
+                    var name = patioFunctionDeclaration.name();
+
+                    var wasmFuncType = new WasmFuncType(
+                        patioFunctionDeclaration.params().stream()
+                            .map(param -> patioTypeToWasmValueType(param.type()))
+                            .toList(),
+                        patioTypeToWasmResultTypes(patioFunctionDeclaration.returnType())
+                    );
+
+                    module = module.addImport(new WasmImport(
+                        "env",
+                        name,
+                        new WasmImportDescriptor.Func(name, wasmFuncType)
+                    ));
+                }
+                case PatioFunctionDefinition patioFunctionDefinition -> {
+                    var name = patioFunctionDefinition.name();
+
+                    var wasmParams = patioFunctionDefinition.params().stream()
+                        .map(param -> patioLocalToWasmLocal(param))
+                        .toList();
+
+                    var wasmLocals = patioFunctionDefinition.locals().stream()
+                        .map(local -> patioLocalToWasmLocal(local))
+                        .toList();
+
+                    var wasmResultTypes = patioTypeToWasmResultTypes(patioFunctionDefinition.returnType());
+
+                    var context = Context.forFunction(Lists.concat(wasmParams, wasmLocals));
+                    var wasmBody = generateInstructions(
+                        patioFunctionDefinition.body(),
+                        context.clearInstructions()
+                    );
+
+                    module = module.addFunction(new WasmFunction(
+                        name,
+                        wasmParams,
+                        wasmLocals,
+                        wasmResultTypes,
+                        new WasmExpression(wasmBody)
+                    ));
+                }
+            }
+        }
+
+        module = wasmMemoryBuilder.toModule(module);
+
+        return module;
+    }
+
+    private static PVector<WasmInstruction> generateInstructions(
+        List<PatioInstruction> patioInstructions,
+        Context context
+    ) {
+        for (var patioInstruction : patioInstructions) {
+            context = generateInstruction(patioInstruction, context);
+        }
+
+        return context.instructions;
+    }
+
+    private static Context generateInstruction(
+        PatioInstruction patioInstruction,
+        Context context
+    ) {
+        return switch (patioInstruction) {
+            case PatioAlloc patioAlloc -> {
+                context = context
+                    .addInstruction(new WasmI32Const(sizeOf(patioAlloc.layout())))
+                    .addInstruction(new WasmCall(new WasmFuncName(ALLOC_NAME)))
+                    .addInstruction(new WasmLocalSet(context.wasmLocalIndex(patioAlloc.target())));
+
+                yield context;
+            }
+            case PatioCall patioCall -> {
+                for (var patioArg : patioCall.args()) {
+                    context = generateOperand(patioArg, context);
+                }
+
+                context = context.addInstruction(new WasmCall(new WasmFuncName(patioCall.functionName())));
+
+                if (patioCall.target().isPresent()) {
+                    var targetWasmLocalIndex = context.wasmLocalIndex(patioCall.target().get());
+                    context = context.addInstruction(new WasmLocalSet(targetWasmLocalIndex));
+                }
+
+                yield context;
+            }
+            case PatioCallPointer patioCallPointer -> {
+                for (var patioArg : patioCallPointer.args()) {
+                    context = generateOperand(patioArg, context);
+                }
+
+                context = generateOperand(patioCallPointer.functionPointer(), context);
+
+                var wasmFuncType = new WasmFuncType(
+                    patioCallPointer.args().stream()
+                        // TODO: get type from arg
+                        .map(arg -> WasmValueType.I32)
+                        .toList(),
+                    List.of(WasmValueType.I32)
+                );
+
+                context = context.addInstruction(new WasmCallIndirect(
+                    // TODO: remove hard-coded zero
+                    new WasmTableIndex(0),
+                    wasmFuncType
+                ));
+
+                var targetWasmLocalIndex = context.wasmLocalIndex(patioCallPointer.target());
+                context = context.addInstruction(new WasmLocalSet(targetWasmLocalIndex));
+
+                yield context;
+            }
+            case PatioDataPtrToInt patioDataPtrToInt -> {
+                context = generateOperand(patioDataPtrToInt.value(), context);
+
+                var targetWasmLocalIndex = context.wasmLocalIndex(patioDataPtrToInt.target());
+                context = context.addInstruction(new WasmLocalSet(targetWasmLocalIndex));
+
+                yield context;
+            }
+            case PatioIf patioIf -> {
+                context = generateOperand(patioIf.condition(), context);
+
+                context = context.addInstruction(new WasmIf(
+                    new WasmBlockType(List.of()),
+                    generateInstructions(patioIf.trueBody(), context.clearInstructions()),
+                    generateInstructions(patioIf.falseBody(), context.clearInstructions())
+                ));
+
+                yield context;
+            }
+            case PatioIntAdd patioIntAdd -> {
+                context = generateOperand(patioIntAdd.left(), context);
+
+                context = generateOperand(patioIntAdd.right(), context);
+
+                context = context.addInstruction(new WasmI32Add());
+
+                var targetWasmLocalIndex = context.wasmLocalIndex(patioIntAdd.target());
+                context = context.addInstruction(new WasmLocalSet(targetWasmLocalIndex));
+
+                yield context;
+            }
+            case PatioIntToDataPtr patioIntToDataPtr -> {
+                context = generateOperand(patioIntToDataPtr.value(), context);
+
+                var targetWasmLocalIndex = context.wasmLocalIndex(patioIntToDataPtr.target());
+                context = context.addInstruction(new WasmLocalSet(targetWasmLocalIndex));
+
+                yield context;
+            }
+            case PatioLocalStore patioLocalStore -> {
+                context = generateOperand(patioLocalStore.value(), context);
+
+                var targetWasmLocalIndex = context.wasmLocalIndex(patioLocalStore.target());
+                context = context.addInstruction(new WasmLocalSet(targetWasmLocalIndex));
+
+                yield context;
+            }
+            case PatioMemoryLoad patioMemoryLoad -> {
+                context = generateOperand(patioMemoryLoad.pointer(), context);
+                
+                context = context.addInstruction(new WasmI32Load(
+                    wasmMemoryArg(patioMemoryLoad.layout(), patioMemoryLoad.offset())
+                ));
+
+                var targetWasmLocalIndex = context.wasmLocalIndex(patioMemoryLoad.target());
+                context = context.addInstruction(new WasmLocalSet(targetWasmLocalIndex));
+                
+                yield context;
+            }
+            case PatioMemoryStore patioMemoryStore -> {
+                context = generateOperand(patioMemoryStore.pointer(), context);
+
+                context = generateOperand(patioMemoryStore.value(), context);
+
+                context = context.addInstruction(new WasmI32Store(
+                    wasmMemoryArg(patioMemoryStore.layout(), patioMemoryStore.offset())
+                ));
+
+                yield context;
+            }
+            case PatioReturn patioReturn -> {
+                context = generateOperand(patioReturn.value(), context);
+
+                context = context.addInstruction(new WasmReturn());
+
+                yield context;
+            }
+        };
+    }
+
+    private static WasmMemoryArg wasmMemoryArg(
+        PatioDataLayout patioDataLayout,
+        List<Integer> patioOffsets
+    ) {
+        var wasmOffset = 0;
+
+        for (var patioOffset : patioOffsets) {
+            switch (patioDataLayout) {
+                case PatioScalar patioScalar -> {
+                    throw new RuntimeException("offset not allowed for scalars");
+                }
+                case PatioStruct patioStruct -> {
+                    for (var i = 0; i < patioOffset; i++) {
+                        wasmOffset += sizeOf(patioStruct.elements().get(i));
+                    }
+                    patioDataLayout = patioStruct.elements().get(patioOffset);
+                }
+            }
+        }
+
+        // TODO: alignment
+        return new WasmMemoryArg(wasmOffset, WasmAlignment.NONE);
+    }
+
+    private static Context generateOperand(
+        PatioOperand patioOperand,
+        Context context
+    ) {
+        return switch (patioOperand) {
+            case PatioBool patioBool ->
+                context.addInstruction(new WasmI32Const(patioBool.value() ? 1 : 0));
+
+            case PatioDataPtr patioDataPtr -> {
+                var wasmMemoryAddress = new WasmMemoryAddress(patioDataPtr.name());
+                yield context.addInstruction(new WasmMemoryAddressConst(wasmMemoryAddress));
+            }
+
+            case PatioFunctionPtr patioFunctionPtr -> {
+                var wasmFuncName = new WasmFuncName(patioFunctionPtr.name());
+                yield context.addInstruction(new WasmTableElementIndexConst(wasmFuncName));
+            }
+
+            case PatioInt patioInt ->
+                context.addInstruction(new WasmI32Const(patioInt.value()));
+
+            case PatioLocalRef patioLocalRef -> {
+                var wasmLocalIndex = context.patioLocalNameToWasmLocalIndex.get(patioLocalRef.name());
+                yield context.addInstruction(new WasmLocalGet(wasmLocalIndex));
+            }
+        };
+    }
+
+    private static int sizeOf(PatioDataLayout layout) {
+        return switch (layout) {
+            case PatioScalar patioScalar ->
+                4;
+
+            case PatioStruct patioStruct ->
+                patioStruct.elements().stream()
+                    .map(element -> sizeOf(element))
+                    .reduce(0, Integer::sum);
+        };
+    }
+
+    private static WasmLocal patioLocalToWasmLocal(PatioLocal local) {
+        return new WasmLocal(local.name(), patioTypeToWasmValueType(local.type()));
+    }
+
+    private static WasmValueType patioTypeToWasmValueType(PatioType type) {
+        return WasmValueType.I32;
+    }
+
+    private static List<WasmValueType> patioTypeToWasmResultTypes(PatioType type) {
+        return type instanceof PatioVoidType ? List.of() : List.of(WasmValueType.I32);
+    }
+
+    private static class Context {
+        public static Context forFunction(List<WasmLocal> locals) {
+            var patioLocalNameToWasmLocalIndex = P.<String, WasmLocalIndex>map();
+
+            for (var local : locals) {
+                patioLocalNameToWasmLocalIndex = patioLocalNameToWasmLocalIndex.plus(
+                    local.name(),
+                    new WasmLocalIndex(patioLocalNameToWasmLocalIndex.size())
+                );
+            }
+
+            return new Context(
+                P.vector(),
+                patioLocalNameToWasmLocalIndex
+            );
+        }
+
+        private final PVector<WasmInstruction> instructions;
+        private final PMap<String, WasmLocalIndex> patioLocalNameToWasmLocalIndex;
+
+        private Context(
+            PVector<WasmInstruction> instructions,
+            PMap<String, WasmLocalIndex> patioLocalNameToWasmLocalIndex
+        ) {
+            this.instructions = instructions;
+            this.patioLocalNameToWasmLocalIndex = patioLocalNameToWasmLocalIndex;
+        }
+
+        private Context addInstruction(WasmInstruction instruction) {
+            return new Context(this.instructions.plus(instruction), this.patioLocalNameToWasmLocalIndex);
+        }
+
+        private WasmLocalIndex wasmLocalIndex(PatioLocalRef patioLocalRef) {
+            return this.patioLocalNameToWasmLocalIndex.get(patioLocalRef.name());
+        }
+
+        public Context clearInstructions() {
+            return new Context(P.vector(), this.patioLocalNameToWasmLocalIndex);
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/util/Lists.java b/src/main/java/org/shedlang/compiler/util/Lists.java
index c89534fc..527b5eb1 100644
--- a/src/main/java/org/shedlang/compiler/util/Lists.java
+++ b/src/main/java/org/shedlang/compiler/util/Lists.java
@@ -9,6 +9,13 @@ public class Lists {
     private Lists() {
     }
 
+    public static <T> List<T> concat(List<T> left, List<T> right) {
+        var result = new ArrayList<T>(left.size() + right.size());
+        result.addAll(left);
+        result.addAll(right);
+        return result;
+    }
+
     public static <T> Optional<T> last(List<T> list) {
         if (list.isEmpty()) {
             return Optional.empty();