Use context in StackCompiler

1 files changed, 200 insertions, 98 deletions

diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 537e9322..7d1ed4b0 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -1,69 +1,78 @@
 package org.shedlang.compiler.stackir;
 
+import org.pcollections.PStack;
 import org.pcollections.PVector;
 import org.shedlang.compiler.ShedBuiltins;
 import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.stackir.instructions.*;
-import org.shedlang.compiler.util.Lists;
 import org.shedlang.compiler.util.P;
-import org.shedlang.compiler.util.PCollectors;
+
+import java.util.List;
 
 public class StackCompiler {
     private StackCompiler() {
     }
 
-    private static PVector<Instruction> compileBinaryOperation(TypedBinaryOperationNode node) {
-        return compileExpression(node.left())
-            .plusAll(compileExpression(node.right()))
-            .plus(switch (node.operator()) {
-                case INT_ADD ->
-                    new IntAddInstruction(node.source());
-                case INT_EQUALS ->
-                    new IntEqualsInstruction(node.source());
-                case INT_NOT_EQUAL ->
-                    new IntNotEqualInstruction(node.source());
-                case INT_SUB ->
-                    new IntSubInstruction(node.source());
-                case STRING_ADD ->
-                    new StringAddInstruction(node.source());
-            });
-    }
-
-    private static PVector<Instruction> compileBlock(TypedBlockNode node) {
-        var body = node.statements().stream()
-            .flatMap(statement -> compileFunctionStatement(statement).stream())
-            .collect(PCollectors.toVector());
+    private static Context compileBinaryOperation(
+        TypedBinaryOperationNode node,
+        Context context
+    ) {
+        context = compileExpression(node.left(), context);
+
+        context = compileExpression(node.right(), context);
+
+        context = context.addInstruction(switch (node.operator()) {
+            case INT_ADD ->
+                new IntAddInstruction(node.source());
+            case INT_EQUALS ->
+                new IntEqualsInstruction(node.source());
+            case INT_NOT_EQUAL ->
+                new IntNotEqualInstruction(node.source());
+            case INT_SUB ->
+                new IntSubInstruction(node.source());
+            case STRING_ADD ->
+                new StringAddInstruction(node.source());
+        });
+
+        return context;
+    }
+
+    private static Context compileBlock(TypedBlockNode node, Context context) {
+        for (var statement : node.statements()) {
+            context = compileFunctionStatement(statement, context);
+        }
 
         switch (node.behaviour()) {
             case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
             }
             case TypedFunctionStatementBehaviour.Continue continue_ -> {
-                body = body.plus(new PushUnitInstruction(node.source()));
+                context = context.addInstruction(new PushUnitInstruction(node.source()));
             }
             case TypedFunctionStatementBehaviour.Terminate terminate -> {
             }
         }
-        return body;
+
+        return context;
     }
 
-    private static PVector<Instruction> compileBlockReturn(TypedBlockReturnNode blockReturn) {
-        return compileExpression(blockReturn.expression());
+    private static Context compileBlockReturn(TypedBlockReturnNode blockReturn, Context context) {
+        return compileExpression(blockReturn.expression(), context);
     }
 
-    private static PVector<Instruction> compileBoolLiteral(TypedBoolLiteralNode node) {
-        return P.vector(new PushBoolInstruction(node.value(), node.source()));
+    private static Context compileBoolLiteral(TypedBoolLiteralNode node, Context context) {
+        return context.addInstruction(new PushBoolInstruction(node.value(), node.source()));
     }
 
-    private static PVector<Instruction> compileCall(TypedCallNode node) {
-        var instructions = compileExpression(node.receiver());
+    private static Context compileCall(TypedCallNode node, Context context) {
+        context = compileExpression(node.receiver(), context);
 
-        instructions = node.args().positional().stream()
-            .map(arg -> compileExpression(arg))
-            .reduce(instructions, PVector::plusAll);
+        for (var arg : node.args().positional()) {
+            context = compileExpression(arg, context);
+        }
 
-        instructions = node.args().named().stream()
-            .map(arg -> compileExpression(arg.expression()))
-            .reduce(instructions, PVector::plusAll);
+        for (var arg : node.args().named()) {
+            context = compileExpression(arg.expression(), context);
+        }
 
         var callInstruction = new CallInstruction(
             node.args().positional().size(),
@@ -73,41 +82,49 @@ public class StackCompiler {
             node.source()
         );
 
-        instructions = instructions.plus(callInstruction);
+        context = context.addInstruction(callInstruction);
 
-        return instructions;
+        return context;
     }
 
-    public static PVector<Instruction> compileExpression(TypedExpressionNode node) {
+    public static Context compileExpression(TypedExpressionNode node, Context context) {
         return switch (node) {
             case TypedBinaryOperationNode typedBinaryOperationNode ->
-                compileBinaryOperation(typedBinaryOperationNode);
+                compileBinaryOperation(typedBinaryOperationNode, context);
 
             case TypedBoolLiteralNode typedBoolLiteralNode ->
-                compileBoolLiteral(typedBoolLiteralNode);
+                compileBoolLiteral(typedBoolLiteralNode, context);
 
             case TypedCallNode typedCallNode ->
-                compileCall(typedCallNode);
+                compileCall(typedCallNode, context);
 
             case TypedIfNode typedIfNode ->
-                compileIf(typedIfNode);
+                compileIf(typedIfNode, context);
 
             case TypedIntLiteralNode typedIntLiteralNode ->
-                compileIntLiteral(typedIntLiteralNode);
+                compileIntLiteral(typedIntLiteralNode, context);
 
             case TypedReferenceNode typedReferenceNode ->
-                compileReference(typedReferenceNode);
+                compileReference(typedReferenceNode, context);
 
             case TypedStringLiteralNode typedStringLiteralNode ->
-                compileStringLiteral(typedStringLiteralNode);
+                compileStringLiteral(typedStringLiteralNode, context);
         };
     }
 
-    private static PVector<Instruction> compileExpressionStatement(TypedExpressionStatementNode node) {
-        return compileExpression(node.expression()).plus(new DiscardInstruction(node.source()));
+    private static Context compileExpressionStatement(
+        TypedExpressionStatementNode node,
+        Context context
+    ) {
+        context = compileExpression(node.expression(), context);
+        context = context.addInstruction(new DiscardInstruction(node.source()));
+        return context;
     }
 
-    private static PVector<Instruction> compileFunctionDefinition(TypedFunctionDefinitionNode node) {
+    private static Context compileFunctionDefinition(
+        TypedFunctionDefinitionNode node,
+        Context context
+    ) {
         var positionalParams = node.params().positional().stream()
             .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
             .toList();
@@ -116,19 +133,25 @@ public class StackCompiler {
             .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
             .toList();
 
-        var body = compileBlock(node.body());
+        context = context.enterBlock();
+
+        context = compileBlock(node.body(), context);
 
         switch (node.body().behaviour()) {
             case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
-                body = body.plus(new ReturnInstruction(node.source()));
+                context = context.addInstruction(new ReturnInstruction(node.source()));
             }
             case TypedFunctionStatementBehaviour.Continue continue_ -> {
-                body = body.plus(new ReturnInstruction(node.source()));
+                context = context.addInstruction(new ReturnInstruction(node.source()));
             }
             case TypedFunctionStatementBehaviour.Terminate terminate -> {
             }
         }
 
+        var body = context.instructions;
+
+        context = context.exitBlock();
+
         var defineFunctionInstruction = new DefineFunctionInstruction(
             node.name(),
             positionalParams,
@@ -140,104 +163,183 @@ public class StackCompiler {
 
         var storeInstruction = new LocalStoreInstruction(node.name(), node.source());
 
-        return P.vector(defineFunctionInstruction, storeInstruction);
+        return context
+            .addInstruction(defineFunctionInstruction)
+            .addInstruction(storeInstruction);
     }
 
-    private static PVector<Instruction> compileFunctionReturn(TypedFunctionReturnNode node) {
-        return compileExpression(node.expression())
-            .plus(new ReturnInstruction(node.source()));
+    private static Context compileFunctionReturn(
+        TypedFunctionReturnNode node,
+        Context context
+    ) {
+        context = compileExpression(node.expression(), context);
+        context = context.addInstruction(new ReturnInstruction(node.source()));
+        return context;
     }
     
-    public static PVector<Instruction> compileFunctionStatement(TypedFunctionStatementNode node) {
+    public static Context compileFunctionStatement(
+        TypedFunctionStatementNode node,
+        Context context
+    ) {
         return switch (node) {
             case TypedBlockReturnNode blockReturn ->
-                compileBlockReturn(blockReturn);
+                compileBlockReturn(blockReturn, context);
             
             case TypedExpressionStatementNode expressionStatement ->
-                compileExpressionStatement(expressionStatement);
+                compileExpressionStatement(expressionStatement, context);
 
             case TypedFunctionReturnNode functionReturn ->
-                compileFunctionReturn(functionReturn);
+                compileFunctionReturn(functionReturn, context);
 
             case TypedValNode val ->
-                compileVal(val);
+                compileVal(val, context);
         };
     }
 
-    private static PVector<Instruction> compileIf(TypedIfNode node) {
-        var instructions = compileBlock(node.elseBody());
-
-        for (var conditionalBranch : Lists.reversed(node.conditionalBranches())) {
-            var conditionInstructions = compileExpression(conditionalBranch.condition());
+    private static Context compileIf(TypedIfNode node, Context context) {
+        return compileIf(
+            node.conditionalBranches(),
+            node.elseBody(),
+            context
+        );
+    }
 
-            var ifInstruction = new IfInstruction(
-                compileBlock(conditionalBranch.body()),
-                instructions,
-                node.source()
+    private static Context compileIf(
+        List<TypedConditionalBranchNode> conditionalBranches,
+        TypedBlockNode elseBody,
+        Context context
+    ) {
+        if (conditionalBranches.isEmpty()) {
+            return compileBlock(elseBody, context);
+        } else {
+            var conditionalBranch = conditionalBranches.get(0);
+            context = compileExpression(conditionalBranch.condition(), context);
+
+            context = context.enterBlock();
+            context = compileBlock(conditionalBranch.body(), context);
+            var ifTrueInstructions = context.instructions;
+            context = context.exitBlock();
+
+            context = context.enterBlock();
+            context = compileIf(
+                conditionalBranches.subList(1, conditionalBranches.size()),
+                elseBody,
+                context
             );
-
-            instructions = conditionInstructions.plus(ifInstruction);
+            var ifFalseInstructions = context.instructions;
+            context = context.exitBlock();
+
+            context = context.addInstruction(new IfInstruction(
+                ifTrueInstructions,
+                ifFalseInstructions,
+                conditionalBranch.source()
+            ));
         }
 
-        return instructions;
+        return context;
     }
 
-    private static PVector<Instruction> compileIntLiteral(TypedIntLiteralNode node) {
-        return P.vector(new PushIntInstruction(node.value(), node.source()));
+    private static Context compileIntLiteral(TypedIntLiteralNode node, Context context) {
+        return context.addInstruction(new PushIntInstruction(node.value(), node.source()));
     }
 
     public static StackIrNamespace compileNamespace(TypedNamespaceNode node) {
-        var instructions = P.<Instruction>vector();
+        var context = Context.EMPTY;
 
-        instructions = instructions
-            .plus(new NamespaceImportInstruction(ShedBuiltins.NAMESPACE_NAME, node.source()))
-            .plus(new FieldLoadInstruction(ShedBuiltins.NAMESPACE_NAME, ShedBuiltins.PRINT_NAME, node.source()))
-            .plus(new LocalStoreInstruction(ShedBuiltins.PRINT_NAME, node.source()));
+        context = context
+            .addInstruction(new NamespaceImportInstruction(ShedBuiltins.NAMESPACE_NAME, node.source()))
+            .addInstruction(new FieldLoadInstruction(ShedBuiltins.NAMESPACE_NAME, ShedBuiltins.PRINT_NAME, node.source()))
+            .addInstruction(new LocalStoreInstruction(ShedBuiltins.PRINT_NAME, node.source()));
 
         for (var statement : node.body()) {
-            instructions = instructions.plusAll(compileNamespaceStatement(statement));
+            context = compileNamespaceStatement(statement, context);
         }
 
         for (var fieldName : node.fieldNames()) {
-            instructions = instructions
-                .plus(new LocalLoadInstruction(fieldName, node.source()))
-                .plus(new NamespaceFieldStoreInstruction(
+            context = context
+                .addInstruction(new LocalLoadInstruction(fieldName, node.source()))
+                .addInstruction(new NamespaceFieldStoreInstruction(
                     node.namespaceName(),
                     fieldName,
                     node.source()
                 ));
         }
 
-        return new StackIrNamespace(node.namespaceName(), instructions);
+        return new StackIrNamespace(node.namespaceName(), context.instructions);
     }
 
-    public static PVector<Instruction> compileNamespaceStatement(TypedNamespaceStatementNode node) {
+    public static Context compileNamespaceStatement(TypedNamespaceStatementNode node, Context context) {
         return switch (node) {
             case TypedFunctionDefinitionNode functionDefinition ->
-                compileFunctionDefinition(functionDefinition);
+                compileFunctionDefinition(functionDefinition, context);
 
             case TypedValNode val ->
-                compileVal(val);
+                compileVal(val, context);
         };
     }
 
-    private static PVector<Instruction> compileReference(TypedReferenceNode node) {
-        return P.vector(new LocalLoadInstruction(node.name(), node.source()));
+    private static Context compileReference(TypedReferenceNode node, Context context) {
+        return context.addInstruction(new LocalLoadInstruction(node.name(), node.source()));
     }
 
-    private static PVector<Instruction> compileStringLiteral(TypedStringLiteralNode node) {
-        return P.vector(new PushStringInstruction(node.value(), node.source()));
+    private static Context compileStringLiteral(
+        TypedStringLiteralNode node,
+        Context context
+    ) {
+        var instruction = new PushStringInstruction(node.value(), node.source());
+        return context.addInstruction(instruction);
     }
 
-    private static PVector<Instruction> compileTarget(TypedTargetNode target) {
+    private static Context compileTarget(TypedTargetNode target, Context context) {
         return switch (target) {
             case TypedTargetNode.Variable variable ->
-                P.vector(new LocalStoreInstruction(variable.name(), variable.source()));
+                context.addInstruction(new LocalStoreInstruction(variable.name(), variable.source()));
         };
     }
 
-    private static PVector<Instruction> compileVal(TypedValNode val) {
-        return compileExpression(val.expression())
-            .plusAll(compileTarget(val.target()));
+    private static Context compileVal(TypedValNode val, Context context) {
+        context = compileExpression(val.expression(), context);
+        context = compileTarget(val.target(), context);
+        return context;
+    }
+
+    public static class Context {
+        public static final Context EMPTY = new Context(
+            P.stack(),
+            P.vector()
+        );
+
+        private final PStack<PVector<Instruction>> instructionsStack;
+        private final PVector<Instruction> instructions;
+
+        private Context(PStack<PVector<Instruction>> instructionsStack, PVector<Instruction> instructions) {
+            this.instructionsStack = instructionsStack;
+            this.instructions = instructions;
+        }
+
+        public Context enterBlock() {
+            return new Context(
+                this.instructionsStack.plus(this.instructions),
+                P.vector()
+            );
+        }
+
+        public Context exitBlock() {
+            return new Context(
+                this.instructionsStack.minus(0),
+                this.instructionsStack.get(0)
+            );
+        }
+
+        private Context addInstruction(Instruction instruction) {
+            return new Context(
+                this.instructionsStack,
+                this.instructions.plus(instruction)
+            );
+        }
+
+        public List<Instruction> instructions() {
+            return this.instructions;
+        }
     }
 }