Fix array indexing

2 files changed, 26 insertions, 12 deletions

diff --git a/builtins/array.h b/builtins/array.h
index e4540340..8f56d1c5 100644
--- a/builtins/array.h
+++ b/builtins/array.h
@@ -13,24 +13,30 @@
 
 // Convert negative indices to back-indexed without branching: index0 = index + (index < 0)*(len+1)) - 1
 #define Array_get(item_type, arr_expr, index_expr, filename, start, end) *({ \
-    const array_t arr = arr_expr; int64_t index = Int_to_Int64(index_expr, false); \
+    const array_t arr = arr_expr; int64_t index = 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", Int64$as_text(&index, no, NULL), arr.length); \
     (item_type*)(arr.data + arr.stride * off);})
+#define Array_get_unchecked(type, x, i) *({ const array_t arr = x; int64_t index = i; \
+                                          int64_t off = index + (index < 0) * (arr.length + 1) - 1; \
+                                          (type*)(arr.data + arr.stride * off);})
 #define Array_lvalue(item_type, arr_expr, index_expr, padded_item_size, filename, start, end) *({ \
-    array_t *arr = arr_expr; int64_t index = Int_to_Int64(index_expr, false); \
+    array_t *arr = arr_expr; int64_t index = 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", Int64$as_text(&index, no, NULL), arr->length); \
     if (arr->data_refcount > 0) \
         Array$compact(arr, padded_item_size); \
     (item_type*)(arr->data + arr->stride * off); })
+#define Array_lvalue_unchecked(item_type, arr_expr, index_expr, padded_item_size) *({ \
+    array_t *arr = arr_expr; int64_t index = index_expr; \
+    int64_t off = index + (index < 0) * (arr->length + 1) - 1; \
+    if (arr->data_refcount > 0) \
+        Array$compact(arr, padded_item_size); \
+    (item_type*)(arr->data + arr->stride * off); })
 #define Array_set(item_type, arr, index, value, padded_item_size, filename, start, end) \
     Array_lvalue(item_type, arr_expr, index, padded_item_size, filename, start, end) = value
-#define Array_get_unchecked(type, x, i) *({ const array_t arr = x; int64_t index = I(i); \
-                                          int64_t off = index + (index < 0) * (arr.length + 1) - 1; \
-                                          (type*)(arr.data + arr.stride * off);})
 #define is_atomic(x) _Generic(x, bool: true, int8_t: true, int16_t: true, int32_t: true, int64_t: true, float: true, double: true, default: false)
 #define TypedArray(t, ...) ({ t items[] = {__VA_ARGS__}; \
                          (array_t){.length=sizeof(items)/sizeof(items[0]), \
diff --git a/compile.c b/compile.c
index dffa145e..cc06b92d 100644
--- a/compile.c
+++ b/compile.c
@@ -216,10 +216,17 @@ static CORD compile_lvalue(env_t *env, ast_t *ast)
         if (container_t->tag == 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_lvalue(", compile_type(item_type), ", ", target_code, ", ", 
-                            compile(env, index->index), ", ", CORD_asprintf("%ld", padded_type_size(item_type)),
-                            ", ", Text$quoted(ast->file->filename, false), ", ", heap_strf("%ld", ast->start - ast->file->text),
-                            ", ", heap_strf("%ld", ast->end - ast->file->text), ")");
+            if (index->unchecked) {
+                return CORD_all("Array_lvalue_unchecked(", compile_type(item_type), ", ", target_code, ", ", 
+                                compile_int_to_type(env, index->index, Type(IntType, .bits=TYPE_IBITS64)),
+                                ", ", CORD_asprintf("%ld", padded_type_size(item_type)), ")");
+            } else {
+                return CORD_all("Array_lvalue(", compile_type(item_type), ", ", target_code, ", ", 
+                                compile_int_to_type(env, index->index, Type(IntType, .bits=TYPE_IBITS64)),
+                                ", ", CORD_asprintf("%ld", padded_type_size(item_type)),
+                                ", ", Text$quoted(ast->file->filename, false), ", ", heap_strf("%ld", ast->start - ast->file->text),
+                                ", ", heap_strf("%ld", ast->end - ast->file->text), ")");
+            }
         } else {
             code_err(ast, "I don't know how to assign to this target");
         }
@@ -2681,12 +2688,13 @@ CORD compile(env_t *env, ast_t *ast)
                 code_err(indexing->index, "Arrays can only be indexed by integers, not %T", index_t);
             type_t *item_type = Match(container_t, ArrayType)->item_type;
             CORD arr = compile_to_pointer_depth(env, indexing->indexed, 0, false);
-            CORD index = compile(env, indexing->index);
             file_t *f = indexing->index->file;
             if (indexing->unchecked)
-                return CORD_all("Array_get_unchecked(", compile_type(item_type), ", ", arr, ", ", index, ")");
+                return CORD_all("Array_get_unchecked(", compile_type(item_type), ", ", arr, ", ",
+                                compile_int_to_type(env, indexing->index, Type(IntType, .bits=TYPE_IBITS64)), ")");
             else
-                return CORD_all("Array_get(", compile_type(item_type), ", ", arr, ", ", index, ", ",
+                return CORD_all("Array_get(", compile_type(item_type), ", ", arr, ", ",
+                                compile_int_to_type(env, indexing->index, Type(IntType, .bits=TYPE_IBITS64)), ", ",
                                 Text$quoted(f->filename, false), ", ", CORD_asprintf("%ld", (int64_t)(indexing->index->start - f->text)), ", ",
                                 CORD_asprintf("%ld", (int64_t)(indexing->index->end - f->text)),
                                 ")");