🐐 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]
-    if isinstance(target, ValDict):
-      if not isinstance(index, (ValStr, ValInt)):
-        raise CentvrionError("Dict key must be a string or integer")
-      d = dict(target.value())
-      d[index.value()] = val
-      vtable[self.id.name] = ValDict(d)
-      return vtable, ValNul()
-    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)
+    root = vtable[self.id.name]
+    containers = [root]
+    for idx in evaluated_indices[:-1]:
+      containers.append(_index_get(containers[-1], idx))
+    new_val = _index_set(containers[-1], evaluated_indices[-1], val)
+    for i in range(len(containers) - 2, -1, -1):
+      new_val = _index_set(containers[i], evaluated_indices[i], new_val)
+    vtable[self.id.name] = new_val
     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()
-    return (
-      idx_lines + val_lines + [
-        f'CentValue {arr_tmp} = cent_scope_get(&_scope, "{node.id.name}");',
-        f"cent_list_index_set(&{arr_tmp}, {idx_var}, {val_var});",
-        f'cent_scope_set(&_scope, "{node.id.name}", {arr_tmp});',
-      ]
-    )
+    lines += val_lines
+    root_tmp = ctx.fresh_tmp()
+    lines.append(f'CentValue {root_tmp} = cent_scope_get(&_scope, "{node.id.name}");')
+    if len(idx_vars) == 1:
+      lines.append(f"cent_list_index_set(&{root_tmp}, {idx_vars[0]}, {val_var});")
+    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):