Throw exception if assigning value of wrong type

5 files changed, 46 insertions, 13 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/VariableNode.java b/src/main/java/org/zwobble/couscous/ast/VariableNode.java
index 4c689c5..acfa29b 100644
--- a/src/main/java/org/zwobble/couscous/ast/VariableNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/VariableNode.java
@@ -1,5 +1,8 @@
 package org.zwobble.couscous.ast;
 
+import org.zwobble.couscous.values.TypeReference;
+
 public interface VariableNode {
     int getId();
+    TypeReference getType();
 }
diff --git a/src/main/java/org/zwobble/couscous/interpreter/Environment.java b/src/main/java/org/zwobble/couscous/interpreter/Environment.java
index a24bf1c..27e337a 100644
--- a/src/main/java/org/zwobble/couscous/interpreter/Environment.java
+++ b/src/main/java/org/zwobble/couscous/interpreter/Environment.java
@@ -1,6 +1,5 @@
 package org.zwobble.couscous.interpreter;
 
-import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 
@@ -8,27 +7,36 @@ import org.zwobble.couscous.Project;
 import org.zwobble.couscous.ast.VariableNode;
 import org.zwobble.couscous.values.ConcreteType;
 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.val;
 
 public class Environment {
-    private final Map<Integer, Optional<InterpreterValue>> stackFrame;
     private final Project project;
+    private final Map<Integer, Optional<InterpreterValue>> stackFrameValues;
+    private final Map<Integer, TypeReference> stackFrameTypes;
 
-    public Environment(Project project, Map<Integer, Optional<InterpreterValue>> stackFrame) {
+    public Environment(Project project, Map<VariableNode, Optional<InterpreterValue>> stackFrame) {
         this.project = project;
-        this.stackFrame = new HashMap<>(stackFrame);
+        this.stackFrameValues = mapKeys(stackFrame, variable -> variable.getId());
+        this.stackFrameTypes = stackFrame.keySet()
+            .stream()
+            .collect(toMap(variable -> variable.getId(), variable -> variable.getType()));
     }
 
     public InterpreterValue get(int variableId) {
         checkVariableIsInScope(variableId);
-        val value = stackFrame.get(variableId);
+        val value = stackFrameValues.get(variableId);
         return value.orElseThrow(() -> new UnboundVariable(variableId));
     }
 
     public void put(int variableId, InterpreterValue value) {
         checkVariableIsInScope(variableId);
-        stackFrame.put(variableId, Optional.of(value));
+        checkVariableType(variableId, value);
+        stackFrameValues.put(variableId, Optional.of(value));
     }
 
     public void put(VariableNode variable, InterpreterValue value) {
@@ -40,13 +48,21 @@ public class Environment {
         return project.findClass(className);
     }
 
-    public Environment withStackFrame(Map<Integer, Optional<InterpreterValue>> stackFrame) {
+    public Environment withStackFrame(Map<VariableNode, Optional<InterpreterValue>> stackFrame) {
         return new Environment(project, stackFrame);
     }
 
     private void checkVariableIsInScope(int variableId) {
-        if (!stackFrame.containsKey(variableId)) {
+        if (!stackFrameValues.containsKey(variableId)) {
             throw new VariableNotInScope(variableId);
         }
     }
+
+    private void checkVariableType(int variableId, InterpreterValue value) {
+        val variableType = stackFrameTypes.get(variableId);
+        val valueType = value.getType().getReference();
+        if (!variableType.equals(valueType)) {
+            throw new UnexpectedValueType(variableType, valueType);
+        }
+    }
 }
diff --git a/src/main/java/org/zwobble/couscous/interpreter/StackFrameBuilder.java b/src/main/java/org/zwobble/couscous/interpreter/StackFrameBuilder.java
index ac425cd..658bf42 100644
--- a/src/main/java/org/zwobble/couscous/interpreter/StackFrameBuilder.java
+++ b/src/main/java/org/zwobble/couscous/interpreter/StackFrameBuilder.java
@@ -10,20 +10,20 @@ import com.google.common.collect.ImmutableMap;
 import java.util.Optional;
 
 public class StackFrameBuilder {
-    private final ImmutableMap.Builder<Integer, Optional<InterpreterValue>> variables
+    private final ImmutableMap.Builder<VariableNode, Optional<InterpreterValue>> variables
         = ImmutableMap.builder();
     
     public StackFrameBuilder declare(VariableNode variable, InterpreterValue value) {
-        variables.put(variable.getId(), Optional.of(value));
+        variables.put(variable, Optional.of(value));
         return this;
     }
     
     public StackFrameBuilder declare(VariableNode variable) {
-        variables.put(variable.getId(), Optional.empty());
+        variables.put(variable, Optional.empty());
         return this;
     }
     
-    public Map<Integer, Optional<InterpreterValue>> build() {
+    public Map<VariableNode, Optional<InterpreterValue>> build() {
         return variables.build();
     }
 }
diff --git a/src/main/java/org/zwobble/couscous/util/CousCousMaps.java b/src/main/java/org/zwobble/couscous/util/CousCousMaps.java
new file mode 100644
index 0000000..5fb4dca
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/util/CousCousMaps.java
@@ -0,0 +1,14 @@
+package org.zwobble.couscous.util;
+
+import java.util.Map;
+import java.util.function.Function;
+
+import static java.util.stream.Collectors.toMap;
+
+public class CousCousMaps {
+    public static <K1, K2, V> Map<K2, V> mapKeys(Map<K1, V> map, Function<K1, K2> func) {
+        return map.entrySet()
+            .stream()
+            .collect(toMap(entry -> func.apply(entry.getKey()), entry -> entry.getValue()));
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/values/ConcreteType.java b/src/main/java/org/zwobble/couscous/values/ConcreteType.java
index 520b504..f3c226c 100644
--- a/src/main/java/org/zwobble/couscous/values/ConcreteType.java
+++ b/src/main/java/org/zwobble/couscous/values/ConcreteType.java
@@ -129,7 +129,7 @@ public class ConcreteType<T> {
         return method;
     }
     
-    private TypeReference getReference() {
+    public TypeReference getReference() {
         return typeRef(name);
     }