🐐 index assignment

centvrion/ast_nodes.py

@@ -307,6 +307,33 @@ class Designa(Node):
     return vtable, ValNul()
 
 
+class DesignaIndex(Node):
+  def __init__(self, variable: ID, index, value) -> None:
+    self.id = variable
+    self.index = index
+    self.value = value
+
+  def __eq__(self, other):
+    return type(self) == type(other) and self.id == other.id and self.index == other.index and self.value == other.value
+
+  def __repr__(self) -> str:
+    return f"DesignaIndex({self.id!r}, {self.index!r}, {self.value!r})"
+
+  def print(self):
+    return f"DESIGNA {self.id.print()}[{self.index.print()}] VT {self.value.print()}"
+
+  def _eval(self, vtable):
+    vtable, index = self.index.eval(vtable)
+    vtable, val = self.value.eval(vtable)
+    i = index.value()
+    lst = list(vtable[self.id.name].value())
+    if i < 1 or i > len(lst):
+      raise IndexError(f"Index {i} out of range for array of length {len(lst)}")
+    lst[i - 1] = val
+    vtable[self.id.name] = ValList(lst)
+    return vtable, ValNul()
+
+
 class Defini(Node):
   def __init__(self, name, parameters, statements) -> None:
     self.name = name

centvrion/parser.py

@@ -70,6 +70,10 @@ class Parser():
     def statement_designa(tokens):
       return ast_nodes.Designa(tokens[1], tokens[3])
 
+    @self.pg.production('statement : KEYWORD_DESIGNA id SYMBOL_LBRACKET expression SYMBOL_RBRACKET KEYWORD_VT expression')
+    def statement_designa_index(tokens):
+      return ast_nodes.DesignaIndex(tokens[1], tokens[3], tokens[6])
+
     @self.pg.production('statement : expression')
     def statement_expression(tokens):
       return ast_nodes.ExpressionStatement(tokens[0])

tests.py

@@ -6,7 +6,7 @@ from parameterized import parameterized
 
 from centvrion.ast_nodes import (
   ArrayIndex, Bool, BinOp, BuiltIn, DataArray, DataRangeArray, Defini,
-  Designa, DumStatement, Erumpe, ExpressionStatement, ID,
+  Designa, DesignaIndex, DumStatement, Erumpe, ExpressionStatement, ID,
   Invoca, ModuleCall, Nullus, Numeral, PerStatement,
   Program, Redi, SiStatement, String, UnaryMinus, UnaryNot,
   num_to_int, int_to_num, make_string,
@@ -924,6 +924,58 @@ class TestArrayIndex(unittest.TestCase):
     run_test(self, source, nodes, value)
 
 
+# --- Array index assignment ---
+
+array_index_assign_tests = [
+  # assign to middle element
+  ("DESIGNA a VT [I, II, III]\nDESIGNA a[II] VT X\na[II]",
+    Program([], [
+      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
+      DesignaIndex(ID("a"), Numeral("II"), Numeral("X")),
+      ExpressionStatement(ArrayIndex(ID("a"), Numeral("II"))),
+    ]),
+    ValInt(10)),
+  # assign to first element
+  ("DESIGNA a VT [I, II, III]\nDESIGNA a[I] VT V\na[I]",
+    Program([], [
+      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
+      DesignaIndex(ID("a"), Numeral("I"), Numeral("V")),
+      ExpressionStatement(ArrayIndex(ID("a"), Numeral("I"))),
+    ]),
+    ValInt(5)),
+  # assign to last element
+  ("DESIGNA a VT [I, II, III]\nDESIGNA a[III] VT L\na[III]",
+    Program([], [
+      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
+      DesignaIndex(ID("a"), Numeral("III"), Numeral("L")),
+      ExpressionStatement(ArrayIndex(ID("a"), Numeral("III"))),
+    ]),
+    ValInt(50)),
+  # other elements unaffected
+  ("DESIGNA a VT [I, II, III]\nDESIGNA a[II] VT X\na[I]",
+    Program([], [
+      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
+      DesignaIndex(ID("a"), Numeral("II"), Numeral("X")),
+      ExpressionStatement(ArrayIndex(ID("a"), Numeral("I"))),
+    ]),
+    ValInt(1)),
+  # expression as index
+  ("DESIGNA a VT [I, II, III]\nDESIGNA i VT II\nDESIGNA a[i] VT X\na[II]",
+    Program([], [
+      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
+      Designa(ID("i"), Numeral("II")),
+      DesignaIndex(ID("a"), ID("i"), Numeral("X")),
+      ExpressionStatement(ArrayIndex(ID("a"), Numeral("II"))),
+    ]),
+    ValInt(10)),
+]
+
+class TestArrayIndexAssign(unittest.TestCase):
+  @parameterized.expand(array_index_assign_tests)
+  def test_array_index_assign(self, source, nodes, value):
+    run_test(self, source, nodes, value)
+
+
 # --- Comments ---
 
 comment_tests = [