Fix/improve assignment to arrays and tables

2 files changed, 48 insertions, 13 deletions

diff --git a/builtins/array.h b/builtins/array.h
index 209a72e2..6375f6d4 100644
--- a/builtins/array.h
+++ b/builtins/array.h
@@ -11,12 +11,20 @@
 #include "types.h"
 
 // Convert negative indices to back-indexed without branching: index0 = index + (index < 0)*(len+1)) - 1
-#define $Array_get(type, x, i, filename, start, end) *({ \
-    const array_t $arr = x; int64_t $index = (int64_t)(i); \
+#define $Array_get(item_type, arr_expr, index_expr, filename, start, end) *({ \
+    const array_t $arr = arr_expr; int64_t $index = (int64_t)(index_expr); \
     int64_t $off = $index + ($index < 0) * ($arr.length + 1) - 1; \
     if (__builtin_expect($off < 0 || $off >= $arr.length, 0)) \
         fail_source(filename, start, end, "Invalid array index: %r (array has length %ld)\n", Int__as_text(&$index, USE_COLOR, NULL), $arr.length); \
-    (type*)($arr.data + $arr.stride * $off);})
+    (item_type*)($arr.data + $arr.stride * $off);})
+#define $Array_set(item_type, arr_expr, index_expr, value_expr, typeinfo, filename, start, end) { \
+    array_t *$arr = arr_expr; int64_t $index = (int64_t)(index_expr); \
+    int64_t $off = $index + ($index < 0) * ($arr->length + 1) - 1; \
+    if (__builtin_expect($off < 0 || $off >= $arr->length, 0)) \
+        fail_source(filename, start, end, "Invalid array index: %r (array has length %ld)\n", Int__as_text(&$index, USE_COLOR, NULL), $arr->length); \
+    if ($arr->data_refcount > 0) \
+        Array__compact($arr, typeinfo); \
+    *(item_type*)($arr->data + $arr->stride * $off) = value_expr; }
 #define $Array_get_unchecked(type, x, i) *({ const array_t $arr = x; int64_t $index = (int64_t)(i); \
                                           int64_t $off = $index + ($index < 0) * ($arr.length + 1) - 1; \
                                           (type*)($arr.data + $arr.stride * $off);})
diff --git a/compile.c b/compile.c
index 82c302d8..e08a13f7 100644
--- a/compile.c
+++ b/compile.c
@@ -14,6 +14,8 @@
 #include "typecheck.h"
 #include "builtins/util.h"
 
+static CORD compile_to_pointer_depth(env_t *env, ast_t *ast, int64_t target_depth, bool allow_optional);
+
 CORD compile_type_ast(env_t *env, type_ast_t *t)
 {
     switch (t->tag) {
@@ -116,13 +118,41 @@ static void check_assignable(env_t *env, ast_t *ast)
     if (!can_be_mutated(env, ast)) {
         if (ast->tag == Index || ast->tag == FieldAccess) {
             ast_t *subject = ast->tag == Index ? Match(ast, Index)->indexed : Match(ast, FieldAccess)->fielded;
-            code_err(subject, "This is a readonly pointer, which can't be assigned to");
+            code_err(subject, "This is an immutable value, you can't assign to it");
         } else {
             code_err(ast, "This is a value of type %T and can't be assigned to", get_type(env, ast));
         }
     }
 }
 
+static CORD compile_assignment(env_t *env, ast_t *target, CORD value)
+{
+    check_assignable(env, target);
+    if (target->tag == Index) {
+        auto index = Match(target, Index);
+        type_t *container_t = get_type(env, index->indexed);
+        container_t = value_type(container_t);
+        switch (container_t->tag) {
+        case ArrayType: {
+            CORD target_code = compile_to_pointer_depth(env, index->indexed, 1, false);
+            type_t *item_type = Match(container_t, ArrayType)->item_type;
+            return CORD_all("$Array_set(", compile_type(env, item_type), ", ", target_code, ", ", 
+                            compile(env, index->index), ", ", value, ", ", compile_type_info(env, container_t),
+                            ", ", Text__quoted(target->file->filename, false), ", ", heap_strf("%ld", target->start - target->file->text),
+                            ", ", heap_strf("%ld", target->end - target->file->text), ");\n");
+        }
+        case TableType: {
+            CORD target_code = compile_to_pointer_depth(env, index->indexed, 1, false);
+            return CORD_all("Table_set_value(", target_code, ", ", compile(env, index->index), ", ",
+                            value, ", ", compile_type_info(env, container_t), ");\n");
+        }
+        case PointerType: break;
+        default: code_err(target, "I don't know how to assign to this target");
+        }
+    }
+    return CORD_all(compile(env, target), " = ", value, ";\n");
+}
+
 CORD compile_statement(env_t *env, ast_t *ast)
 {
     switch (ast->tag) {
@@ -194,7 +224,7 @@ CORD compile_statement(env_t *env, ast_t *ast)
                 // Common case: assigning to one variable:
                 check_assignable(env, assign->targets->ast);
                 CORD var = compile(env, assign->targets->ast);
-                CORD code = CORD_all(var, " = ", compile(env, assign->values->ast), ";");
+                CORD code = compile_assignment(env, assign->targets->ast, compile(env, assign->values->ast));
                 CORD_appendf(&code, "$test(&%r, %r, %r, %r, %ld, %ld);",
                     var, compile_type_info(env, get_type(env, assign->targets->ast)),
                     compile(env, WrapAST(test->expr, TextLiteral, .cord=test->output)),
@@ -212,10 +242,8 @@ CORD compile_statement(env_t *env, ast_t *ast)
                 for (ast_list_t *value = assign->values; value; value = value->next)
                     CORD_appendf(&code, "%r $%ld = %r;\n", compile_type(env, get_type(env, value->ast)), i++, compile(env, value->ast));
                 i = 1;
-                for (ast_list_t *target = assign->targets; target; target = target->next) {
-                    check_assignable(env, target->ast);
-                    CORD_appendf(&code, "%r = $%ld;\n", compile(env, target->ast), i++);
-                }
+                for (ast_list_t *target = assign->targets; target; target = target->next)
+                    code = CORD_all(code, compile_assignment(env, target->ast, CORD_asprintf("$%ld", i++)));
 
                 CORD_appendf(&code, "$test(&$1, %r, %r, %r, %ld, %ld);",
                     compile_type_info(env, get_type(env, assign->targets->ast)),
@@ -273,7 +301,7 @@ CORD compile_statement(env_t *env, ast_t *ast)
         auto assign = Match(ast, Assign);
         // Single assignment:
         if (assign->targets && !assign->targets->next)
-            return CORD_asprintf("%r = %r;", compile(env, assign->targets->ast), compile(env, assign->values->ast));
+            return compile_assignment(env, assign->targets->ast, compile(env, assign->values->ast));
 
         CORD code = "{ // Assignment\n";
         int64_t i = 1;
@@ -281,8 +309,7 @@ CORD compile_statement(env_t *env, ast_t *ast)
             CORD_appendf(&code, "%r $%ld = %r;\n", compile_type(env, get_type(env, value->ast)), i++, compile(env, value->ast));
         i = 1;
         for (ast_list_t *target = assign->targets; target; target = target->next) {
-            check_assignable(env, target->ast);
-            CORD_appendf(&code, "%r = $%ld;\n", compile(env, target->ast), i++);
+            code = CORD_cat(code, compile_assignment(env, target->ast, CORD_asprintf("$%ld", i++)));
         }
         return CORD_cat(code, "\n}");
     }
@@ -745,7 +772,7 @@ CORD compile_string(env_t *env, ast_t *ast, CORD color)
     return expr_as_text(env, expr, t, color);
 }
 
-static CORD compile_to_pointer_depth(env_t *env, ast_t *ast, int64_t target_depth, bool allow_optional)
+CORD compile_to_pointer_depth(env_t *env, ast_t *ast, int64_t target_depth, bool allow_optional)
 {
     CORD val = compile(env, ast);
     type_t *t = get_type(env, ast);