Support Patio code gen for ifs

2 files changed, 123 insertions, 5 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index 63d5cd5e..475831c6 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -251,6 +251,59 @@ public class PatioCodeGen {
                 yield context;
             }
 
+            case IfInstruction if_ -> {
+                var conditionIntResult = context.allocateLocal(
+                    PatioCodeGenLocalSpace.TEMPORARY,
+                    "conditionInt",
+                    PatioTypes.INT
+                );
+                var conditionInt = conditionIntResult.value();
+                context = conditionIntResult.context();
+
+                context = context.temporaryPopIntoLocal(conditionInt, if_.source());
+
+                var conditionBoolResult = context.allocateLocal(
+                    PatioCodeGenLocalSpace.TEMPORARY,
+                    "conditionBool",
+                    PatioTypes.BOOL
+                );
+                var conditionBool = conditionBoolResult.value();
+                context = conditionBoolResult.context();
+
+                context = context.addInstruction(new PatioIntToBool(
+                    conditionBool,
+                    conditionInt,
+                    if_.source()
+                ));
+
+                var valueResult = context.allocateLocal(
+                    PatioCodeGenLocalSpace.TEMPORARY,
+                    PatioCodeGenTypes.VALUE_TYPE
+                );
+                var value = valueResult.value();
+                context = valueResult.context();
+
+                var trueBodyResult = generateBlock(if_.trueBody(), if_.source(), context);
+                var trueBody = trueBodyResult.value();
+                context = trueBodyResult.context();
+
+                var falseBodyResult = generateBlock(if_.falseBody(), if_.source(), context);
+                var falseBody = falseBodyResult.value();
+                context = falseBodyResult.context();
+
+                context = context.addInstruction(new PatioIf(
+                    List.of(value),
+                    conditionBool,
+                    trueBody,
+                    falseBody,
+                    if_.source()
+                ));
+
+                context = context.temporaryPush(value);
+
+                yield context;
+            }
+
             case IntAddInstruction intAdd ->
                 generateBinaryOperation(
                     (target, left, right, operationContext) ->
@@ -423,6 +476,29 @@ public class PatioCodeGen {
         ).temporaryPush(target);
     }
 
+    private static PatioCodeGenResult<PatioBlock> generateBlock(
+        List<Instruction> wasmInstructions,
+        Source source,
+        PatioCodeGenContext context
+    ) {
+        context = context.enterBlock();
+        context = generateInstructions(wasmInstructions, context);
+
+        var targetResult = context.temporaryPop();
+        var target = targetResult.value();
+        context = targetResult.context();
+
+        var instructionsResult = context.exitBlock();
+        var instructions = instructionsResult.value();
+        context = instructionsResult.context();
+
+        return context.toResult(new PatioBlock(
+            instructions,
+            List.of(target),
+            source
+        ));
+    }
+
     private static PatioCodeGenContext pushValue(int value, Source source, PatioCodeGenContext context) {
         var localResult = context.allocateLocal(PatioCodeGenLocalSpace.TEMPORARY, PatioCodeGenTypes.VALUE_TYPE);
         var local = localResult.value();
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index f19f92f2..59f307ce 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -18,11 +18,13 @@ import org.shedlang.compiler.types.NamespaceTypeInfo;
 import org.shedlang.compiler.util.P;
 
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import java.util.function.Function;
 
 public record PatioCodeGenContext(
     CompilationUnitContext compilationUnitContext,
+    PStack<FunctionContext> blockStack,
     FunctionContext functionContext
 ) {
     public record CompilationUnitContext(
@@ -205,12 +207,14 @@ public record PatioCodeGenContext(
 
     public static final PatioCodeGenContext EMPTY = new PatioCodeGenContext(
         CompilationUnitContext.EMPTY,
+        P.stack(),
         FunctionContext.EMPTY
     );
 
     public static PatioCodeGenContext forTypeInfo(TypeInfo typeInfo) {
         return new PatioCodeGenContext(
             CompilationUnitContext.forTypeInfo(typeInfo),
+            P.stack(),
             FunctionContext.EMPTY
         );
     }
@@ -264,7 +268,7 @@ public record PatioCodeGenContext(
     public PatioCodeGenContext temporaryPopIntoLocal(PatioLocalRef target, Source source) {
         var result = this.functionContext.temporaryPop();
 
-        return new PatioCodeGenContext(this.compilationUnitContext, result.context)
+        return new PatioCodeGenContext(this.compilationUnitContext, this.blockStack, result.context)
             .addInstruction(new PatioLocalStore(target, result.value, source));
     }
     
@@ -320,6 +324,7 @@ public record PatioCodeGenContext(
     ) {
         var bodyContext = new PatioCodeGenContext(
             this.compilationUnitContext,
+            P.stack(),
             FunctionContext.EMPTY
         );
 
@@ -346,24 +351,61 @@ public record PatioCodeGenContext(
                 bodyContext.functionContext.instructions,
                 source
             )),
+            this.blockStack,
             this.functionContext
         );
     }
 
+    public PatioCodeGenContext enterBlock() {
+        return new PatioCodeGenContext(
+            this.compilationUnitContext,
+            this.blockStack.plus(this.functionContext),
+            new FunctionContext(
+                P.stack(),
+                P.vector(),
+                this.functionContext.localsByName,
+                this.functionContext.stackIrNameToLocal,
+                P.vector()
+            )
+        );
+    }
+
+    public PatioCodeGenResult<List<PatioInstruction>> exitBlock() {
+        if (!this.functionContext.args.isEmpty()) {
+            throw new RuntimeException("cannot add args in block");
+        }
+
+        if (!this.functionContext.temporaryStack.isEmpty()) {
+            throw new RuntimeException("temporary stack must be empty at the end of a block");
+        }
+
+        return new PatioCodeGenContext(
+            this.compilationUnitContext,
+            this.blockStack.minus(0),
+            new FunctionContext(
+                this.blockStack.get(0).temporaryStack,
+                this.blockStack.get(0).instructions,
+                this.functionContext.localsByName,
+                this.functionContext.stackIrNameToLocal,
+                this.blockStack.get(0).args
+            )
+        ).toResult(this.functionContext.instructions);
+    }
+
     private PatioCodeGenContext withCompilationUnitContext(CompilationUnitContext compilationUnitContext) {
-        return new PatioCodeGenContext(compilationUnitContext, this.functionContext);
+        return new PatioCodeGenContext(compilationUnitContext, this.blockStack, this.functionContext);
     }
 
     private <T> PatioCodeGenResult<T> withCompilationUnitContextResult(CompilationUnitContext.Result<T> result) {
-        return new PatioCodeGenContext(result.context, this.functionContext).toResult(result.value);
+        return withCompilationUnitContext(result.context).toResult(result.value);
     }
 
     private PatioCodeGenContext withFunctionContext(FunctionContext functionContext) {
-        return new PatioCodeGenContext(this.compilationUnitContext, functionContext);
+        return new PatioCodeGenContext(this.compilationUnitContext, this.blockStack, functionContext);
     }
 
     private <T> PatioCodeGenResult<T> withFunctionContextResult(FunctionContext.Result<T> result) {
-        return new PatioCodeGenContext(this.compilationUnitContext, result.context).toResult(result.value);
+        return withFunctionContext(result.context).toResult(result.value);
     }
 
     public <T> PatioCodeGenResult<T> toResult(T value) {