Support if instructions in WASM backend

2 files changed, 84 insertions, 16 deletions

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 cd5f3b24..ac8c8d62 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
@@ -155,9 +155,7 @@ public class WasmCodeGen {
                     defineFunction.namedParams(),
                     defineFunction.freeVariables(),
                     functionContext -> {
-                        for (var bodyInstruction : defineFunction.body()) {
-                            functionContext = generateInstruction(bodyInstruction, functionContext);
-                        }
+                        functionContext = generateInstructions(defineFunction.body(), functionContext);
                         return functionContext;
                     },
                     context
@@ -181,8 +179,25 @@ public class WasmCodeGen {
                 );
             }
 
-            case IfInstruction if_ ->
-                throw new UnsupportedOperationException("TODO");
+            case IfInstruction if_ -> {
+                context = context.enterBlock();
+                context = generateInstructions(if_.trueBody(), context);
+                var trueBodyResult = context.exitBlock();
+                var trueBody = trueBodyResult.value();
+                context = trueBodyResult.context();
+
+                context = context.enterBlock();
+                context = generateInstructions(if_.falseBody(), context);
+                var falseBodyResult = context.exitBlock();
+                var falseBody = falseBodyResult.value();
+                context = falseBodyResult.context();
+
+                yield context.addInstruction(new WasmIf(
+                    new WasmBlockType(List.of(WasmCodeGenTypes.VALUE_TYPE)),
+                    trueBody,
+                    falseBody
+                ));
+            }
 
             case IntAddInstruction intAdd ->
                 context.addInstruction(new WasmI32Add());
@@ -251,4 +266,14 @@ public class WasmCodeGen {
                 WasmCodeGenStrings.add(context);
         };
     }
+
+    private static WasmCodeGenContext generateInstructions(
+        List<Instruction> instructions,
+        WasmCodeGenContext context
+    ) {
+        for (var instruction : instructions) {
+            context = generateInstruction(instruction, context);
+        }
+        return context;
+    }
 }
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 8e147ad6..1cb0e2e2 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
@@ -11,6 +11,7 @@ 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 org.shedlang.compiler.typechecker.TypeInfo;
+import org.shedlang.compiler.util.P;
 
 import java.util.List;
 import java.util.Optional;
@@ -168,25 +169,54 @@ public record WasmCodeGenContext(
     record FunctionContext(
         WasmFuncType funcType,
         PVector<WasmInstruction> instructions,
+        PStack<PVector<WasmInstruction>> blockStack,
         OrderedPMap<String, WasmLocal> locals,
         PMap<String, WasmLocalIndex> shedLocalNameToLocalIndex
     ) {
-        public static final FunctionContext EMPTY = new FunctionContext(
-            new WasmFuncType(List.of(), List.of()),
-            Empty.vector(),
-            OrderedPMap.empty(),
-            Empty.map()
+        public static final FunctionContext EMPTY = forFuncType(
+            new WasmFuncType(List.of(), List.of())
         );
 
+        public static FunctionContext forFuncType(WasmFuncType funcType) {
+            return new FunctionContext(
+                funcType,
+                Empty.vector(),
+                Empty.stack(),
+                OrderedPMap.empty(),
+                Empty.map()
+            );
+        }
+
         public FunctionContext addInstruction(WasmInstruction instruction) {
             return new FunctionContext(
                 this.funcType,
                 this.instructions.plus(instruction),
+                this.blockStack,
                 this.locals,
                 this.shedLocalNameToLocalIndex
             );
         }
 
+        public FunctionContext enterBlock() {
+            return new FunctionContext(
+                this.funcType,
+                P.vector(),
+                this.blockStack.plus(this.instructions),
+                this.locals,
+                this.shedLocalNameToLocalIndex
+            );
+        }
+
+        public Result<PVector<WasmInstruction>> exitBlock() {
+            return new FunctionContext(
+                this.funcType,
+                this.blockStack.get(0),
+                this.blockStack.minus(0),
+                this.locals,
+                this.shedLocalNameToLocalIndex
+            ).toResult(this.instructions);
+        }
+
         public Result<WasmLocalIndex> generateLocal(WasmCodeGenLocalSpace space, Optional<String> name) {
             var localIndex = new WasmLocalIndex(locals.size());
             var localName = NameAllocator.generateName(
@@ -197,6 +227,7 @@ public record WasmCodeGenContext(
             return new FunctionContext(
                 this.funcType,
                 this.instructions,
+                this.blockStack,
                 this.locals.plus(
                     localName,
                     new WasmLocal(localName, WasmCodeGenTypes.VALUE_TYPE)
@@ -213,6 +244,7 @@ public record WasmCodeGenContext(
             return new FunctionContext(
                 context.funcType,
                 context.instructions,
+                context.blockStack,
                 context.locals,
                 context.shedLocalNameToLocalIndex.plus(name, localIndex)
             ).toResult(localIndex);
@@ -322,6 +354,22 @@ public record WasmCodeGenContext(
         );
     }
 
+    public WasmCodeGenContext enterBlock() {
+        return new WasmCodeGenContext(
+            this.moduleContext,
+            this.functionContext.enterBlock()
+        );
+    }
+
+    public WasmCodeGenResult<List<WasmInstruction>> exitBlock() {
+        var exitResult = this.functionContext.exitBlock();
+
+        return new WasmCodeGenContext(
+            this.moduleContext,
+            exitResult.context
+        ).toResult(exitResult.value);
+    }
+
     public WasmFunction toFunction(String name) {
         return this.functionContext.toFunction(name);
     }
@@ -359,12 +407,7 @@ public record WasmCodeGenContext(
     ) {
         var bodyContext = new WasmCodeGenContext(
             this.moduleContext,
-            new FunctionContext(
-                funcType,
-                Empty.vector(),
-                OrderedPMap.empty(),
-                Empty.map()
-            )
+            FunctionContext.forFuncType(funcType)
         );
 
         bodyContext = create.apply(bodyContext);