Push shed variable map into BlockContext

1 files changed, 32 insertions, 31 deletions

diff --git a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
index f7014dc6..2888446d 100644
--- a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
@@ -760,17 +760,21 @@ public class ShidCodeGen {
     private static class BlockContext {
         public static final BlockContext EMPTY = new BlockContext(
             P.vector(),
+            P.map(),
             Sources.NULL
         );
 
         private final PVector<ShidInstruction> instructions;
+        private final PMap<String, ShedVariable> shedVariablesByName;
         private final Source source;
 
         public BlockContext(
             PVector<ShidInstruction> instructions,
+            PMap<String, ShedVariable> shedVariablesByName,
             Source source
         ) {
             this.instructions = instructions;
+            this.shedVariablesByName = shedVariablesByName;
             this.source = source;
         }
 
@@ -781,6 +785,25 @@ public class ShidCodeGen {
         public BlockContext addInstruction(ShidInstruction instruction) {
             return new BlockContext(
                 this.instructions.plus(instruction),
+                this.shedVariablesByName,
+                this.source
+            );
+        }
+
+        public ShedVariable shedVariable(String variableName) {
+            var shedVariable = this.shedVariablesByName.get(variableName);
+
+            if (shedVariable == null) {
+                throw new RuntimeException("Unknown Shed variable: " + variableName);
+            }
+
+            return shedVariable;
+        }
+
+        public BlockContext addShedVariable(String name, ShedVariable variable) {
+            return new BlockContext(
+                this.instructions,
+                this.shedVariablesByName.plus(name, variable),
                 this.source
             );
         }
@@ -791,38 +814,28 @@ public class ShidCodeGen {
             P.stack(),
             // TODO: no block context by default?
             BlockContext.EMPTY,
-            Empty.set(),
-            P.stack(),
-            P.map()
+            Empty.set()
         );
 
         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<BlockContext> blockStack,
             BlockContext block,
-            PSet<String> localNames,
-            PStack<PMap<String, ShedVariable>> shedVariablesByNameStack,
-            PMap<String, ShedVariable> shedVariablesByName
+            PSet<String> localNames
         ) {
             this.blockStack = blockStack;
             this.block = block;
             this.localNames = localNames;
-            this.shedVariablesByNameStack = shedVariablesByNameStack;
-            this.shedVariablesByName = shedVariablesByName;
         }
 
         public Context enterBlock(List<BoundVariable> shedVariables, Source source) {
             var context = new Context(
                 this.blockStack.plus(this.block),
-                new BlockContext(P.vector(), source),
-                this.localNames,
-                this.shedVariablesByNameStack.plus(this.shedVariablesByName),
-                this.shedVariablesByName
+                new BlockContext(P.vector(), this.block.shedVariablesByName, source),
+                this.localNames
             );
 
             for (var shedVariable : shedVariables) {
@@ -836,9 +849,7 @@ public class ShidCodeGen {
             return new Context(
                 this.blockStack.minus(0),
                 this.blockStack.get(0),
-                this.localNames,
-                this.shedVariablesByNameStack.minus(0),
-                this.shedVariablesByNameStack.get(0)
+                this.localNames
             ).toResult(this.block.toBlock());
         }
 
@@ -846,9 +857,7 @@ public class ShidCodeGen {
             return new Context(
                 this.blockStack,
                 this.block.addInstruction(instruction),
-                this.localNames,
-                this.shedVariablesByNameStack,
-                this.shedVariablesByName
+                this.localNames
             );
         }
 
@@ -857,13 +866,7 @@ public class ShidCodeGen {
         }
 
         public ShedVariable shedVariable(String variableName) {
-            var shedVariable = this.shedVariablesByName.get(variableName);
-
-            if (shedVariable == null) {
-                throw new RuntimeException("Unknown Shed variable: " + variableName);
-            }
-
-            return shedVariable;
+            return this.block.shedVariable(variableName);
         }
 
         // TODO: make private / remove? Variables should be declared using enterBlock().
@@ -877,10 +880,8 @@ public class ShidCodeGen {
 
             return new Context(
                 this.blockStack,
-                this.block,
-                this.localNames.plus(localName),
-                this.shedVariablesByNameStack,
-                this.shedVariablesByName.plus(boundVariable.name(), new ShedVariable(localName, boundVariable.isMutable()))
+                this.block.addShedVariable(boundVariable.name(), new ShedVariable(localName, boundVariable.isMutable())),
+                this.localNames.plus(localName)
             );
         }