🐐 Index assignment

centvrion/ast_nodes.py

@@ -598,42 +598,97 @@ class Designa(Node):
     return vtable, ValNul()
 
 
+def _index_get(container, index):
+  if isinstance(container, ValDict):
+    if not isinstance(index, (ValStr, ValInt)):
+      raise CentvrionError("Dict key must be a string or integer")
+    d = container.value()
+    k = index.value()
+    if k not in d:
+      raise CentvrionError("Key not found in dict")
+    return d[k]
+  if isinstance(container, ValList):
+    i = index.value()
+    lst = container.value()
+    if i < 1 or i > len(lst):
+      raise CentvrionError(f"Index {i} out of range for array of length {len(lst)}")
+    return lst[i - 1]
+  if isinstance(container, ValStr):
+    if isinstance(index, ValInt):
+      i = index.value()
+    elif isinstance(index, ValFrac) and index.value().denominator == 1:
+      i = index.value().numerator
+    else:
+      raise CentvrionError("String index must be a number")
+    s = container.value()
+    if i < 1 or i > len(s):
+      raise CentvrionError(f"Index {i} out of range for string of length {len(s)}")
+    return ValStr(s[i - 1])
+  raise CentvrionError("Cannot index into a non-array, non-dict value")
+
+
+def _index_set(container, index, value):
+  if isinstance(container, ValDict):
+    if not isinstance(index, (ValStr, ValInt)):
+      raise CentvrionError("Dict key must be a string or integer")
+    d = dict(container.value())
+    d[index.value()] = value
+    return ValDict(d)
+  if isinstance(container, ValList):
+    i = index.value()
+    lst = list(container.value())
+    if i < 1 or i > len(lst):
+      raise CentvrionError(f"Index {i} out of range for array of length {len(lst)}")
+    lst[i - 1] = value
+    return ValList(lst)
+  if isinstance(container, ValStr):
+    if isinstance(index, ValInt):
+      i = index.value()
+    elif isinstance(index, ValFrac) and index.value().denominator == 1:
+      i = index.value().numerator
+    else:
+      raise CentvrionError("String index must be a number")
+    if not isinstance(value, ValStr) or len(value.value()) != 1:
+      raise CentvrionError("String index assignment requires a single character")
+    s = container.value()
+    if i < 1 or i > len(s):
+      raise CentvrionError(f"Index {i} out of range for string of length {len(s)}")
+    return ValStr(s[:i - 1] + value.value() + s[i:])
+  raise CentvrionError("Cannot assign to index of a non-array, non-dict value")
+
+
 class DesignaIndex(Node):
-  def __init__(self, variable: ID, index, value) -> None:
+  def __init__(self, variable: ID, indices, value) -> None:
     self.id = variable
-    self.index = index
+    self.indices = indices if isinstance(indices, list) else [indices]
     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
+    return type(self) == type(other) and self.id == other.id and self.indices == other.indices and self.value == other.value
 
   def __repr__(self) -> str:
-    return f"DesignaIndex({self.id!r}, {self.index!r}, {self.value!r})"
+    return f"DesignaIndex({self.id!r}, {self.indices!r}, {self.value!r})"
 
   def print(self):
-    return f"DESIGNA {self.id.print()}[{self.index.print()}] VT {self.value.print()}"
+    idx_str = ''.join(f'[{idx.print()}]' for idx in self.indices)
+    return f"DESIGNA {self.id.print()}{idx_str} VT {self.value.print()}"
 
   def _eval(self, vtable):
-    vtable, index = self.index.eval(vtable)
+    evaluated_indices = []
+    for idx_expr in self.indices:
+      vtable, idx_val = idx_expr.eval(vtable)
+      evaluated_indices.append(idx_val)
     vtable, val = self.value.eval(vtable)
     if self.id.name not in vtable:
       raise CentvrionError(f"Undefined variable: {self.id.name}")
-    target = vtable[self.id.name]
+    root = vtable[self.id.name]
-    if isinstance(target, ValDict):
+    containers = [root]
-      if not isinstance(index, (ValStr, ValInt)):
+    for idx in evaluated_indices[:-1]:
-        raise CentvrionError("Dict key must be a string or integer")
+      containers.append(_index_get(containers[-1], idx))
-      d = dict(target.value())
+    new_val = _index_set(containers[-1], evaluated_indices[-1], val)
-      d[index.value()] = val
+    for i in range(len(containers) - 2, -1, -1):
-      vtable[self.id.name] = ValDict(d)
+      new_val = _index_set(containers[i], evaluated_indices[i], new_val)
-      return vtable, ValNul()
+    vtable[self.id.name] = new_val
-    if not isinstance(target, ValList):
-      raise CentvrionError(f"{self.id.name} is not an array or dict")
-    i = index.value()
-    lst = list(target.value())
-    if i < 1 or i > len(lst):
-      raise CentvrionError(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()
 
 

centvrion/compiler/emit_stmt.py

@@ -16,16 +16,32 @@ def emit_stmt(node, ctx):
     return val_lines + [f'cent_scope_set(&_scope, "{node.id.name}", {val_var});']
 
   if isinstance(node, DesignaIndex):
-    idx_lines, idx_var = emit_expr(node.index, ctx)
+    lines = []
+    idx_vars = []
+    for idx_expr in node.indices:
+      idx_lines, idx_var = emit_expr(idx_expr, ctx)
+      lines += idx_lines
+      idx_vars.append(idx_var)
     val_lines, val_var = emit_expr(node.value, ctx)
-    arr_tmp = ctx.fresh_tmp()
+    lines += val_lines
-    return (
+    root_tmp = ctx.fresh_tmp()
-      idx_lines + val_lines + [
+    lines.append(f'CentValue {root_tmp} = cent_scope_get(&_scope, "{node.id.name}");')
-        f'CentValue {arr_tmp} = cent_scope_get(&_scope, "{node.id.name}");',
+    if len(idx_vars) == 1:
-        f"cent_list_index_set(&{arr_tmp}, {idx_var}, {val_var});",
+      lines.append(f"cent_list_index_set(&{root_tmp}, {idx_vars[0]}, {val_var});")
-        f'cent_scope_set(&_scope, "{node.id.name}", {arr_tmp});',
+    else:
-      ]
+      # Walk down to collect intermediate containers
-    )
+      container_tmps = [root_tmp]
+      for idx_var in idx_vars[:-1]:
+        tmp = ctx.fresh_tmp()
+        lines.append(f"CentValue {tmp} = cent_list_index({container_tmps[-1]}, {idx_var});")
+        container_tmps.append(tmp)
+      # Set at deepest level
+      lines.append(f"cent_list_index_set(&{container_tmps[-1]}, {idx_vars[-1]}, {val_var});")
+      # Rebuild up the chain
+      for i in range(len(container_tmps) - 2, -1, -1):
+        lines.append(f"cent_list_index_set(&{container_tmps[i]}, {idx_vars[i]}, {container_tmps[i + 1]});")
+    lines.append(f'cent_scope_set(&_scope, "{node.id.name}", {root_tmp});')
+    return lines
 
   if isinstance(node, DesignaDestructure):
     n = len(node.ids)

centvrion/compiler/runtime/cent_runtime.c

@@ -864,8 +864,23 @@ void cent_list_index_set(CentValue *lst, CentValue idx, CentValue v) {
     cent_dict_set(lst, idx, v);
     return;
   }
+  if (lst->type == CENT_STR) {
+    if (idx.type != CENT_INT)
+      cent_type_error("string index must be an integer");
+    if (v.type != CENT_STR || strlen(v.sval) != 1)
+      cent_type_error("string index assignment requires a single character");
+    long slen = (long)strlen(lst->sval);
+    long i = idx.ival;
+    if (i < 1 || i > slen)
+      cent_runtime_error("string index out of range");
+    char *buf = cent_arena_alloc(cent_arena, slen + 1);
+    memcpy(buf, lst->sval, slen + 1);
+    buf[i - 1] = v.sval[0];
+    lst->sval = buf;
+    return;
+  }
   if (lst->type != CENT_LIST)
-    cent_type_error("index-assign requires a list or dict");
+    cent_type_error("index-assign requires a list, dict, or string");
   if (idx.type != CENT_INT)
     cent_type_error("list index must be an integer");
   long i = idx.ival;

centvrion/parser.py

@@ -173,9 +173,17 @@ 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')
+    @self.pg.production('index_chain : SYMBOL_LBRACKET expression SYMBOL_RBRACKET')
+    def index_chain_single(tokens):
+      return [tokens[1]]
+
+    @self.pg.production('index_chain : SYMBOL_LBRACKET expression SYMBOL_RBRACKET index_chain')
+    def index_chain_multi(tokens):
+      return [tokens[1]] + tokens[3]
+
+    @self.pg.production('statement : KEYWORD_DESIGNA id index_chain KEYWORD_VT expression')
     def statement_designa_index(tokens):
-      return ast_nodes.DesignaIndex(tokens[1], tokens[3], tokens[6])
+      return ast_nodes.DesignaIndex(tokens[1], tokens[2], tokens[4])
 
     @self.pg.production('statement : KEYWORD_DESIGNA id SYMBOL_COMMA id_list_rest KEYWORD_VT expression')
     def statement_designa_destructure(tokens):

tests.py

@@ -1774,7 +1774,7 @@ array_index_assign_tests = [
   ("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")),
+      DesignaIndex(ID("a"), [Numeral("II")], Numeral("X")),
       ExpressionStatement(ArrayIndex(ID("a"), Numeral("II"))),
     ]),
     ValInt(10)),
@@ -1782,7 +1782,7 @@ array_index_assign_tests = [
   ("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")),
+      DesignaIndex(ID("a"), [Numeral("I")], Numeral("V")),
       ExpressionStatement(ArrayIndex(ID("a"), Numeral("I"))),
     ]),
     ValInt(5)),
@@ -1790,7 +1790,7 @@ array_index_assign_tests = [
   ("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")),
+      DesignaIndex(ID("a"), [Numeral("III")], Numeral("L")),
       ExpressionStatement(ArrayIndex(ID("a"), Numeral("III"))),
     ]),
     ValInt(50)),
@@ -1798,7 +1798,7 @@ array_index_assign_tests = [
   ("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")),
+      DesignaIndex(ID("a"), [Numeral("II")], Numeral("X")),
       ExpressionStatement(ArrayIndex(ID("a"), Numeral("I"))),
     ]),
     ValInt(1)),
@@ -1807,7 +1807,7 @@ array_index_assign_tests = [
     Program([], [
       Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
       Designa(ID("i"), Numeral("II")),
-      DesignaIndex(ID("a"), ID("i"), Numeral("X")),
+      DesignaIndex(ID("a"), [ID("i")], Numeral("X")),
       ExpressionStatement(ArrayIndex(ID("a"), Numeral("II"))),
     ]),
     ValInt(10)),
@@ -1819,6 +1819,109 @@ class TestArrayIndexAssign(unittest.TestCase):
     run_test(self, source, nodes, value)
 
 
+# --- Multi-dimensional array index assignment ---
+
+multidim_assign_tests = [
+  # 2D array assignment
+  ("DESIGNA a VT [[I, II], [III, IV]]\nDESIGNA a[I][II] VT X\na[I][II]",
+    Program([], [
+      Designa(ID("a"), DataArray([DataArray([Numeral("I"), Numeral("II")]), DataArray([Numeral("III"), Numeral("IV")])])),
+      DesignaIndex(ID("a"), [Numeral("I"), Numeral("II")], Numeral("X")),
+      ExpressionStatement(ArrayIndex(ArrayIndex(ID("a"), Numeral("I")), Numeral("II"))),
+    ]),
+    ValInt(10)),
+  # other elements unaffected
+  ("DESIGNA a VT [[I, II], [III, IV]]\nDESIGNA a[I][II] VT X\na[II][I]",
+    Program([], [
+      Designa(ID("a"), DataArray([DataArray([Numeral("I"), Numeral("II")]), DataArray([Numeral("III"), Numeral("IV")])])),
+      DesignaIndex(ID("a"), [Numeral("I"), Numeral("II")], Numeral("X")),
+      ExpressionStatement(ArrayIndex(ArrayIndex(ID("a"), Numeral("II")), Numeral("I"))),
+    ]),
+    ValInt(3)),
+  # dict inside array
+  ('DESIGNA a VT [TABVLA {"x" VT I}]\nDESIGNA a[I]["x"] VT X\na[I]["x"]',
+    Program([], [
+      Designa(ID("a"), DataArray([DataDict([(String("x"), Numeral("I"))])])),
+      DesignaIndex(ID("a"), [Numeral("I"), String("x")], Numeral("X")),
+      ExpressionStatement(ArrayIndex(ArrayIndex(ID("a"), Numeral("I")), String("x"))),
+    ]),
+    ValInt(10)),
+  # array inside dict
+  ('DESIGNA d VT TABVLA {"a" VT [I, II]}\nDESIGNA d["a"][I] VT X\nd["a"][I]',
+    Program([], [
+      Designa(ID("d"), DataDict([(String("a"), DataArray([Numeral("I"), Numeral("II")]))])),
+      DesignaIndex(ID("d"), [String("a"), Numeral("I")], Numeral("X")),
+      ExpressionStatement(ArrayIndex(ArrayIndex(ID("d"), String("a")), Numeral("I"))),
+    ]),
+    ValInt(10)),
+  # 3 levels deep
+  ("DESIGNA a VT [[[I]]]\nDESIGNA a[I][I][I] VT X\na[I][I][I]",
+    Program([], [
+      Designa(ID("a"), DataArray([DataArray([DataArray([Numeral("I")])])])),
+      DesignaIndex(ID("a"), [Numeral("I"), Numeral("I"), Numeral("I")], Numeral("X")),
+      ExpressionStatement(ArrayIndex(ArrayIndex(ArrayIndex(ID("a"), Numeral("I")), Numeral("I")), Numeral("I"))),
+    ]),
+    ValInt(10)),
+]
+
+class TestMultidimAssign(unittest.TestCase):
+  @parameterized.expand(multidim_assign_tests)
+  def test_multidim_assign(self, source, nodes, value):
+    run_test(self, source, nodes, value)
+
+
+# --- String index assignment ---
+
+string_index_assign_tests = [
+  # assign to middle character
+  ('DESIGNA s VT "ABCDE"\nDESIGNA s[III] VT "X"\ns',
+    Program([], [
+      Designa(ID("s"), String("ABCDE")),
+      DesignaIndex(ID("s"), [Numeral("III")], String("X")),
+      ExpressionStatement(ID("s")),
+    ]),
+    ValStr("ABXDE")),
+  # assign to first character
+  ('DESIGNA s VT "ABCDE"\nDESIGNA s[I] VT "Z"\ns',
+    Program([], [
+      Designa(ID("s"), String("ABCDE")),
+      DesignaIndex(ID("s"), [Numeral("I")], String("Z")),
+      ExpressionStatement(ID("s")),
+    ]),
+    ValStr("ZBCDE")),
+  # assign to last character
+  ('DESIGNA s VT "ABCDE"\nDESIGNA s[V] VT "Z"\ns',
+    Program([], [
+      Designa(ID("s"), String("ABCDE")),
+      DesignaIndex(ID("s"), [Numeral("V")], String("Z")),
+      ExpressionStatement(ID("s")),
+    ]),
+    ValStr("ABCDZ")),
+  # variable as index
+  ('DESIGNA s VT "ABCDE"\nDESIGNA i VT II\nDESIGNA s[i] VT "X"\ns',
+    Program([], [
+      Designa(ID("s"), String("ABCDE")),
+      Designa(ID("i"), Numeral("II")),
+      DesignaIndex(ID("s"), [ID("i")], String("X")),
+      ExpressionStatement(ID("s")),
+    ]),
+    ValStr("AXCDE")),
+  # string inside array
+  ('DESIGNA a VT ["ABC", "DEF"]\nDESIGNA a[I][II] VT "X"\na[I]',
+    Program([], [
+      Designa(ID("a"), DataArray([String("ABC"), String("DEF")])),
+      DesignaIndex(ID("a"), [Numeral("I"), Numeral("II")], String("X")),
+      ExpressionStatement(ArrayIndex(ID("a"), Numeral("I"))),
+    ]),
+    ValStr("AXC")),
+]
+
+class TestStringIndexAssign(unittest.TestCase):
+  @parameterized.expand(string_index_assign_tests)
+  def test_string_index_assign(self, source, nodes, value):
+    run_test(self, source, nodes, value)
+
+
 # --- Array slicing ---
 
 array_slice_tests = [
@@ -2495,7 +2598,7 @@ dict_assign_tests = [
   ('DESIGNA d VT TABVLA {"a" VT I}\nDESIGNA d["a"] VT X\nd["a"]',
     Program([], [
       Designa(ID("d"), DataDict([(String("a"), Numeral("I"))])),
-      DesignaIndex(ID("d"), String("a"), Numeral("X")),
+      DesignaIndex(ID("d"), [String("a")], Numeral("X")),
       ExpressionStatement(ArrayIndex(ID("d"), String("a"))),
     ]),
     ValInt(10)),
@@ -2503,7 +2606,7 @@ dict_assign_tests = [
   ('DESIGNA d VT TABVLA {"a" VT I}\nDESIGNA d["b"] VT II\nd["b"]',
     Program([], [
       Designa(ID("d"), DataDict([(String("a"), Numeral("I"))])),
-      DesignaIndex(ID("d"), String("b"), Numeral("II")),
+      DesignaIndex(ID("d"), [String("b")], Numeral("II")),
       ExpressionStatement(ArrayIndex(ID("d"), String("b"))),
     ]),
     ValInt(2)),
@@ -2511,7 +2614,7 @@ dict_assign_tests = [
   ('DESIGNA d VT TABVLA {"a" VT I}\nDESIGNA d["b"] VT II\nd["a"]',
     Program([], [
       Designa(ID("d"), DataDict([(String("a"), Numeral("I"))])),
-      DesignaIndex(ID("d"), String("b"), Numeral("II")),
+      DesignaIndex(ID("d"), [String("b")], Numeral("II")),
       ExpressionStatement(ArrayIndex(ID("d"), String("a"))),
     ]),
     ValInt(1)),