Interpret vector element references

Author: mdbrnowski

interpreter.py

@@ -4,7 +4,7 @@
 from generated.MyParser import MyParser
 from generated.MyParserVisitor import MyParserVisitor
 from utils.memory import MemoryStack
-from utils.values import Int, Float, String, Vector
+from utils.values import Value, Int, Float, String, Vector
 
 
 class Break(Exception):
@@ -15,6 +15,13 @@ class Continue(Exception):
     pass
 
 
+def not_same_type(a: Value, b: Value):
+    return type(a) is not type(b) or (
+        isinstance(a, Vector)
+        and (a.dims != b.dims or a.primitive_type != b.primitive_type)
+    )
+
+
 class Interpreter(MyParserVisitor):
     def __init__(self):
         self.memory_stack = MemoryStack()
@@ -84,29 +91,44 @@ def visitComparison(self, ctx: MyParser.ComparisonContext):
                 return a >= b
 
     def visitSimpleAssignment(self, ctx: MyParser.SimpleAssignmentContext):
-        if ctx.id_():
-            self.visitChildren(ctx)
+        if ctx.id_():  # a = 1
             self.memory_stack.put(ctx.id_().getText(), self.visit(ctx.expression()))
-        else:
-            pass  # todo
+        else:  # a[0] = 1
+            ref_value = self.visit(ctx.elementReference())
+            new_value = self.visit(ctx.expression())
+            if not_same_type(ref_value, new_value):
+                raise TypeError
+            ref_value.value = new_value.value
 
     def visitCompoundAssignment(self, ctx: MyParser.CompoundAssignmentContext):
-        if ctx.id_():
-            self.visitChildren(ctx)
-            old_value = self.memory_stack.get(ctx.id_().getText())
+        if ctx.id_():  # a += 1
+            value = self.memory_stack.get(ctx.id_().getText())
             new_value = self.visit(ctx.expression())
             match ctx.getChild(1).symbol.type:
                 case MyParser.ASSIGN_PLUS:
-                    new_value = old_value + new_value
+                    new_value = value + new_value
                 case MyParser.ASSIGN_MINUS:
-                    new_value = old_value - new_value
+                    new_value = value - new_value
                 case MyParser.ASSIGN_MULTIPLY:
-                    new_value = old_value * new_value
+                    new_value = value * new_value
                 case MyParser.ASSIGN_DIVIDE:
-                    new_value = old_value / new_value
+                    new_value = value / new_value
             self.memory_stack.put(ctx.id_().getText(), new_value)
-        else:
-            pass  # todo
+        else:  # a[0] += 1
+            ref_value = self.visit(ctx.elementReference())
+            new_value = self.visit(ctx.expression())
+            if not_same_type(ref_value, new_value):
+                raise TypeError
+            match ctx.getChild(1).symbol.type:
+                case MyParser.ASSIGN_PLUS:
+                    new_value = ref_value + new_value
+                case MyParser.ASSIGN_MINUS:
+                    new_value = ref_value - new_value
+                case MyParser.ASSIGN_MULTIPLY:
+                    new_value = ref_value * new_value
+                case MyParser.ASSIGN_DIVIDE:
+                    new_value = ref_value / new_value
+            ref_value.value = new_value.value
 
     def visitPrint(self, ctx: MyParser.PrintContext):
         for i in range(ctx.getChildCount() // 2):
@@ -182,7 +204,17 @@ def visitVector(self, ctx: MyParser.VectorContext):
         return Vector(elements)
 
     def visitElementReference(self, ctx: MyParser.ElementReferenceContext):
-        return self.visitChildren(ctx)  # todo
+        indices = [
+            self.visit(ctx.expression(i)) for i in range(ctx.getChildCount() // 2 - 1)
+        ]
+        if {type(idx) for idx in indices} != {Int}:
+            raise TypeError
+        result = self.visit(ctx.id_())
+        for idx in indices:
+            if not isinstance(result, Vector):
+                raise TypeError
+            result = result.value[idx.value]
+        return result
 
     def visitId(self, ctx: MyParser.IdContext):
         return self.memory_stack.get(ctx.getText())

test_main.py

@@ -85,6 +85,16 @@ def test_sem_errors(name: str, line_numbers: list[int], additional: str):
         ("while", [4, 3, 2, 1, 0]),
         ("for", [1, 10, 2, 10, 3, 10, 4, 10]),
         ("break_continue", [1, 2, 1, 2, 4] * 2),
+        (
+            "element_reference",
+            [
+                [1, 0],
+                0,
+                [[1, 2], [0, 1]],
+                [[0, 2], [0, 1]],
+                [[0, 0], [0, 1]],
+            ],
+        ),
     ],
 )
 def test_interpreter(name: str, output: str):

tests/interpreter/element_reference.txt

@@ -0,0 +1,9 @@
+A = eye(2);
+print A[0];
+print A[0, 1];
+A[0, 1] = 2;
+print A;
+A[0] = [0, 2];
+print A;
+A[0, 1] -= 2;
+print A;

utils/values.py

@@ -108,7 +108,11 @@ def __init__(self, value: list):
         if (
             len(
                 {
-                    (type(elem), elem.dims if isinstance(elem, Vector) else None)
+                    (
+                        (elem.dims, elem.primitive_type)
+                        if isinstance(elem, Vector)
+                        else type(elem)
+                    )
                     for elem in value
                 }
             )
@@ -118,8 +122,10 @@ def __init__(self, value: list):
 
         if isinstance(value[0], Vector):
             self.dims = (len(value), *value[0].dims)
+            self.primitive_type = value[0].primitive_type
         else:
             self.dims = (len(value),)
+            self.primitive_type = type(value[0])
 
     def __str__(self):
         return "[" + ", ".join(str(elem) for elem in self.value) + "]"