Use single map for stack frame

1 files changed, 21 insertions, 11 deletions

diff --git a/src/main/java/org/zwobble/couscous/interpreter/Environment.java b/src/main/java/org/zwobble/couscous/interpreter/Environment.java
index 27e337a..21cfe04 100644
--- a/src/main/java/org/zwobble/couscous/interpreter/Environment.java
+++ b/src/main/java/org/zwobble/couscous/interpreter/Environment.java
@@ -10,33 +10,35 @@ import org.zwobble.couscous.values.InterpreterValue;
 import org.zwobble.couscous.values.TypeReference;
 
 import static java.util.stream.Collectors.toMap;
-import static org.zwobble.couscous.util.CousCousMaps.mapKeys;
 
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import lombok.Setter;
 import lombok.val;
 
 public class Environment {
     private final Project project;
-    private final Map<Integer, Optional<InterpreterValue>> stackFrameValues;
-    private final Map<Integer, TypeReference> stackFrameTypes;
+    private final Map<Integer, VariableEntry> stackFrame;
 
     public Environment(Project project, Map<VariableNode, Optional<InterpreterValue>> stackFrame) {
         this.project = project;
-        this.stackFrameValues = mapKeys(stackFrame, variable -> variable.getId());
-        this.stackFrameTypes = stackFrame.keySet()
+        this.stackFrame = stackFrame.entrySet()
             .stream()
-            .collect(toMap(variable -> variable.getId(), variable -> variable.getType()));
+            .collect(toMap(
+                entry -> entry.getKey().getId(),
+                entry -> VariableEntry.of(entry.getKey().getType(), entry.getValue())));
     }
 
     public InterpreterValue get(int variableId) {
         checkVariableIsInScope(variableId);
-        val value = stackFrameValues.get(variableId);
-        return value.orElseThrow(() -> new UnboundVariable(variableId));
+        val entry = stackFrame.get(variableId);
+        return entry.getValue().orElseThrow(() -> new UnboundVariable(variableId));
     }
 
     public void put(int variableId, InterpreterValue value) {
         checkVariableIsInScope(variableId);
         checkVariableType(variableId, value);
-        stackFrameValues.put(variableId, Optional.of(value));
+        stackFrame.get(variableId).setValue(Optional.of(value));
     }
 
     public void put(VariableNode variable, InterpreterValue value) {
@@ -53,16 +55,24 @@ public class Environment {
     }
 
     private void checkVariableIsInScope(int variableId) {
-        if (!stackFrameValues.containsKey(variableId)) {
+        if (!stackFrame.containsKey(variableId)) {
             throw new VariableNotInScope(variableId);
         }
     }
 
     private void checkVariableType(int variableId, InterpreterValue value) {
-        val variableType = stackFrameTypes.get(variableId);
+        val variableType = stackFrame.get(variableId).getType();
         val valueType = value.getType().getReference();
         if (!variableType.equals(valueType)) {
             throw new UnexpectedValueType(variableType, valueType);
         }
     }
+    
+    @AllArgsConstructor(staticName="of")
+    @Getter
+    @Setter
+    private static class VariableEntry {
+        private final TypeReference type;
+        private Optional<InterpreterValue> value;
+    }
 }