Compile if expression with single conditional branch

1 files changed, 35 insertions, 9 deletions

diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 7365a8b4..e062ab78 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -5,8 +5,7 @@ import org.shedlang.compiler.ShedBuiltins;
 import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.util.P;
-
-import java.util.stream.Collectors;
+import org.shedlang.compiler.util.PCollectors;
 
 public class StackCompiler {
     private StackCompiler() {
@@ -29,6 +28,23 @@ public class StackCompiler {
             });
     }
 
+    private static PVector<Instruction> compileBlock(TypedBlockNode node) {
+        var body = node.statements().stream()
+            .flatMap(statement -> compileFunctionStatement(statement).stream())
+            .collect(PCollectors.toVector());
+
+        switch (node.behaviour()) {
+            case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
+            }
+            case TypedFunctionStatementBehaviour.Continue continue_ -> {
+                body = body.plus(new PushUnitInstruction(node.source()));
+            }
+            case TypedFunctionStatementBehaviour.Terminate terminate -> {
+            }
+        }
+        return body;
+    }
+
     private static PVector<Instruction> compileBlockReturn(TypedBlockReturnNode blockReturn) {
         return compileExpression(blockReturn.expression());
     }
@@ -73,7 +89,7 @@ public class StackCompiler {
                 compileCall(typedCallNode);
 
             case TypedIfNode typedIfNode ->
-                throw new UnsupportedOperationException("TODO");
+                compileIf(typedIfNode);
 
             case TypedIntLiteralNode typedIntLiteralNode ->
                 compileIntLiteral(typedIntLiteralNode);
@@ -99,17 +115,14 @@ public class StackCompiler {
             .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
             .toList();
 
-        var body = node.body().statements().stream()
-            .flatMap(statement -> compileFunctionStatement(statement).stream())
-            .collect(Collectors.toList());
+        var body = compileBlock(node.body());
 
         switch (node.body().behaviour()) {
             case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
-                body.add(new ReturnInstruction(node.source()));
+                body = body.plus(new ReturnInstruction(node.source()));
             }
             case TypedFunctionStatementBehaviour.Continue continue_ -> {
-                body.add(new PushUnitInstruction(node.source()));
-                body.add(new ReturnInstruction(node.source()));
+                body = body.plus(new ReturnInstruction(node.source()));
             }
             case TypedFunctionStatementBehaviour.Terminate terminate -> {
             }
@@ -150,6 +163,19 @@ public class StackCompiler {
         };
     }
 
+    private static PVector<Instruction> compileIf(TypedIfNode node) {
+        var conditionalBranch = node.conditionalBranches().get(0);
+        var conditionInstructions = compileExpression(conditionalBranch.condition());
+
+        var ifInstruction = new IfInstruction(
+            compileBlock(conditionalBranch.body()),
+            compileBlock(node.elseBody()),
+            node.source()
+        );
+
+        return conditionInstructions.plus(ifInstruction);
+    }
+
     private static PVector<Instruction> compileIntLiteral(TypedIntLiteralNode node) {
         return P.vector(new PushIntInstruction(node.value(), node.source()));
     }