Compile if expression with single conditional branch

2 files changed, 62 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()));
     }
diff --git a/src/main/java/org/shedlang/compiler/util/PCollectors.java b/src/main/java/org/shedlang/compiler/util/PCollectors.java
index 26033670..a0019eef 100644
--- a/src/main/java/org/shedlang/compiler/util/PCollectors.java
+++ b/src/main/java/org/shedlang/compiler/util/PCollectors.java
@@ -2,6 +2,7 @@ package org.shedlang.compiler.util;
 
 import org.pcollections.Empty;
 import org.pcollections.PMap;
+import org.pcollections.PVector;
 
 import java.util.Map;
 import java.util.stream.Collector;
@@ -61,4 +62,30 @@ public class PCollectors {
             return map;
         }
     }
+
+    public static <T> Collector<T, PVectorCollector<T>, PVector<T>> toVector() {
+        return Collector.of(
+            PVectorCollector::new,
+            PVectorCollector::add,
+            PVectorCollector::addAll,
+            PVectorCollector::toVector
+        );
+    }
+
+    private static class PVectorCollector<T> {
+        private PVector<T> vector = Empty.vector();
+
+        void add(T value) {
+            vector = vector.plus(value);
+        }
+
+        PVectorCollector<T> addAll(PVectorCollector<T> other) {
+            vector = vector.plusAll(other.vector);
+            return this;
+        }
+
+        PVector<T> toVector() {
+            return vector;
+        }
+    }
 }