Introduce BlockContext in ShidCodeGen

1 files changed, 68 insertions, 47 deletions

diff --git a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
index e26dd39c..f7014dc6 100644
--- a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
@@ -5,6 +5,7 @@ import org.shedlang.compiler.ShedBuiltins;
 import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.shid.instructions.*;
+import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.Iterables;
 import org.shedlang.compiler.util.P;
@@ -224,7 +225,7 @@ public class ShidCodeGen {
         Source source,
         Context context
     ) {
-        context = context.enterBlock(bodyNode.boundVariables());
+        context = context.enterBlock(bodyNode.boundVariables(), bodyNode.source());
 
         var c = context;
 
@@ -255,7 +256,7 @@ public class ShidCodeGen {
             }
         }
 
-        var bodyResult = context.exitBlock(bodyNode.source());
+        var bodyResult = context.exitBlock();
         var body = bodyResult.value;
         context = bodyResult.context;
 
@@ -373,10 +374,10 @@ public class ShidCodeGen {
         Context context
     ) {
         if (conditionalBranches.isEmpty()) {
-            context = context.enterBlock(elseBody.boundVariables());
+            context = context.enterBlock(elseBody.boundVariables(), elseBody.source());
             context = compileBlockWithYield(elseBody, context);
 
-            var bodyResult = context.exitBlock(elseBody.source());
+            var bodyResult = context.exitBlock();
             var body = bodyResult.value;
             context = bodyResult.context;
 
@@ -385,21 +386,21 @@ public class ShidCodeGen {
             var conditionalBranch = conditionalBranches.get(0);
             context = compileExpression(conditionalBranch.condition(), context);
 
-            context = context.enterBlock(conditionalBranch.body().boundVariables());
+            context = context.enterBlock(conditionalBranch.body().boundVariables(), conditionalBranch.body().source());
             context = compileBlockWithYield(conditionalBranch.body(), context);
 
-            var ifTrueResult = context.exitBlock(conditionalBranch.body().source());
+            var ifTrueResult = context.exitBlock();
             var ifTrue = ifTrueResult.value;
             context = ifTrueResult.context;
 
-            context = context.enterBlock(List.of());
+            context = context.enterBlock(List.of(), conditionalBranch.source());
             context = compileIf(
                 conditionalBranches.subList(1, conditionalBranches.size()),
                 elseBody,
                 context
             );
 
-            var ifFalseResult = context.exitBlock(conditionalBranch.source());
+            var ifFalseResult = context.exitBlock();
             var ifFalse = ifFalseResult.value;
             context = ifFalseResult.context;
 
@@ -419,7 +420,7 @@ public class ShidCodeGen {
     }
 
     public static ShidNamespace compileNamespace(TypedNamespaceNode node) {
-        var context = Context.EMPTY.enterBlock(node.boundVariables());
+        var context = Context.EMPTY.enterBlock(node.boundVariables(), node.source());
 
         context = context.addShedVal(ShedBuiltins.PRINT_NAME);
         var printLocalName = context.shedVariableToLocal(ShedBuiltins.PRINT_NAME);
@@ -479,7 +480,7 @@ public class ShidCodeGen {
                 ));
         }
 
-        var bodyResult = context.exitBlock(node.source());
+        var bodyResult = context.exitBlock();
         var body = bodyResult.value;
         context = bodyResult.context;
 
@@ -642,10 +643,10 @@ public class ShidCodeGen {
             context = compileTryHandlerMethod(handlerMethod, context);
         }
 
-        context = context.enterBlock(node.body().boundVariables());
+        context = context.enterBlock(node.body().boundVariables(), node.body().source());
         context = compileBlockWithYield(node.body(), context);
 
-        var bodyResult = context.exitBlock(node.body().source());
+        var bodyResult = context.exitBlock();
         var body = bodyResult.value;
         context = bodyResult.context;
 
@@ -697,10 +698,10 @@ public class ShidCodeGen {
         var switchCases = new ArrayList<ShidSwitch.Case>();
 
         for (var whenCase : node.cases()) {
-            context = context.enterBlock(whenCase.body().boundVariables());
+            context = context.enterBlock(whenCase.body().boundVariables(), whenCase.body().source());
             context = compileBlockWithYield(whenCase.body(), context);
 
-            var caseBodyResult = context.exitBlock(whenCase.body().source());
+            var caseBodyResult = context.exitBlock();
             var caseBody = caseBodyResult.value;
             context = caseBodyResult.context;
 
@@ -721,14 +722,14 @@ public class ShidCodeGen {
     }
 
     private static Context compileWhile(TypedWhileNode node, Context context) {
-        context = context.enterBlock(node.body().boundVariables());
+        context = context.enterBlock(node.body().boundVariables(), node.body().source());
         context = compileBlockWithoutYield(node.body(), context);
 
-        var ifTrueResult = context.exitBlock(node.body().source());
+        var ifTrueResult = context.exitBlock();
         var ifTrue = ifTrueResult.value;
         context = ifTrueResult.context;
 
-        context = context.enterBlock(List.of());
+        context = context.enterBlock(List.of(), node.source());
 
         context = compileExpression(node.condition(), context);
         context = context.addInstruction(new ShidIf(
@@ -747,7 +748,7 @@ public class ShidCodeGen {
 
         context = context.addInstruction(new ShidLoopContinue(node.source()));
 
-        var bodyResult = context.exitBlock(node.source());
+        var bodyResult = context.exitBlock();
         var body = bodyResult.value;
         context = bodyResult.context;
 
@@ -756,39 +757,69 @@ public class ShidCodeGen {
         return context;
     }
 
+    private static class BlockContext {
+        public static final BlockContext EMPTY = new BlockContext(
+            P.vector(),
+            Sources.NULL
+        );
+
+        private final PVector<ShidInstruction> instructions;
+        private final Source source;
+
+        public BlockContext(
+            PVector<ShidInstruction> instructions,
+            Source source
+        ) {
+            this.instructions = instructions;
+            this.source = source;
+        }
+
+        public ShidBlock toBlock() {
+            return new ShidBlock(this.instructions, this.source);
+        }
+
+        public BlockContext addInstruction(ShidInstruction instruction) {
+            return new BlockContext(
+                this.instructions.plus(instruction),
+                this.source
+            );
+        }
+    }
+
     public static class Context {
         public static final Context EMPTY = new Context(
             P.stack(),
-            P.vector(),
+            // TODO: no block context by default?
+            BlockContext.EMPTY,
             Empty.set(),
             P.stack(),
             P.map()
         );
 
-        private final PStack<PVector<ShidInstruction>> instructionsStack;
-        private final PVector<ShidInstruction> instructions;
+        private final PStack<BlockContext> blockStack;
+        private final BlockContext block;
         private final PSet<String> localNames;
         private final PStack<PMap<String, ShedVariable>> shedVariablesByNameStack;
         private final PMap<String, ShedVariable> shedVariablesByName;
 
         private Context(
-            PStack<PVector<ShidInstruction>> instructionsStack,
-            PVector<ShidInstruction> instructions,
+            PStack<BlockContext> blockStack,
+            BlockContext block,
             PSet<String> localNames,
             PStack<PMap<String, ShedVariable>> shedVariablesByNameStack,
             PMap<String, ShedVariable> shedVariablesByName
         ) {
-            this.instructionsStack = instructionsStack;
-            this.instructions = instructions;
+            this.blockStack = blockStack;
+            this.block = block;
             this.localNames = localNames;
             this.shedVariablesByNameStack = shedVariablesByNameStack;
             this.shedVariablesByName = shedVariablesByName;
         }
 
-        public Context enterBlock(List<BoundVariable> shedVariables) {
+        public Context enterBlock(List<BoundVariable> shedVariables, Source source) {
             var context = new Context(
-                this.instructionsStack.plus(this.instructions),
-                P.vector(),
+                this.blockStack.plus(this.block),
+                new BlockContext(P.vector(), source),
                 this.localNames,
                 this.shedVariablesByNameStack.plus(this.shedVariablesByName),
                 this.shedVariablesByName
@@ -801,30 +832,20 @@ public class ShidCodeGen {
             return context;
         }
 
-        public Result<ShidBlock> exitBlock(Source source) {
+        public Result<ShidBlock> exitBlock() {
             return new Context(
-                this.instructionsStack.minus(0),
-                this.instructionsStack.get(0),
+                this.blockStack.minus(0),
+                this.blockStack.get(0),
                 this.localNames,
                 this.shedVariablesByNameStack.minus(0),
                 this.shedVariablesByNameStack.get(0)
-            ).toResult(new ShidBlock(this.instructions, source));
+            ).toResult(this.block.toBlock());
         }
 
         private Context addInstruction(ShidInstruction instruction) {
             return new Context(
-                this.instructionsStack,
-                this.instructions.plus(instruction),
-                this.localNames,
-                this.shedVariablesByNameStack,
-                this.shedVariablesByName
-            );
-        }
-
-        public Context addInstructions(PVector<ShidInstruction> instructions) {
-            return new Context(
-                this.instructionsStack,
-                this.instructions.plusAll(instructions),
+                this.blockStack,
+                this.block.addInstruction(instruction),
                 this.localNames,
                 this.shedVariablesByNameStack,
                 this.shedVariablesByName
@@ -855,8 +876,8 @@ public class ShidCodeGen {
             var localName = generateLocalNameForVariable(boundVariable.name());
 
             return new Context(
-                this.instructionsStack,
-                this.instructions,
+                this.blockStack,
+                this.block,
                 this.localNames.plus(localName),
                 this.shedVariablesByNameStack,
                 this.shedVariablesByName.plus(boundVariable.name(), new ShedVariable(localName, boundVariable.isMutable()))
@@ -876,7 +897,7 @@ public class ShidCodeGen {
         }
 
         public List<ShidInstruction> instructions() {
-            return this.instructions;
+            return this.block.instructions;
         }
 
         public <T> Result<T> toResult(T value) {