Improve codegen for table/array iteration by inlining the iteration

macros

7 files changed, 54 insertions, 40 deletions

diff --git a/ast.c b/ast.c
index 316bb958..e0354f53 100644
--- a/ast.c
+++ b/ast.c
@@ -179,4 +179,20 @@ int printf_ast(FILE *stream, const struct printf_info *info, const void *const a
     }
 }
 
+bool is_idempotent(ast_t *ast)
+{
+    switch (ast->tag) {
+    case Int: case Bool: case Num: case Var: case Nil: case TextLiteral: return true;
+    case Index: {
+        auto index = Match(ast, Index);
+        return (index->index == NULL) && is_idempotent(index->indexed);
+    }
+    case FieldAccess: {
+        auto access = Match(ast, FieldAccess);
+        return is_idempotent(access->fielded);
+    }
+    default: return false;
+    }
+}
+
 // vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/ast.h b/ast.h
index c43c423c..2fd55250 100644
--- a/ast.h
+++ b/ast.h
@@ -277,5 +277,6 @@ CORD ast_to_xml(ast_t *ast);
 CORD type_ast_to_xml(type_ast_t *ast);
 int printf_ast(FILE *stream, const struct printf_info *info, const void *const args[]);
 ast_list_t *get_ast_children(ast_t *ast);
+bool is_idempotent(ast_t *ast);
 
 // vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/builtins/array.h b/builtins/array.h
index df66e636..b538be30 100644
--- a/builtins/array.h
+++ b/builtins/array.h
@@ -51,13 +51,6 @@
                          .data_refcount=1}; })
 #define $ARRAY_INCREF(arr) (arr).data_refcount |= ((arr).data_refcount << 1) | 1
 #define $ARRAY_DECREF(arr) (arr).data_refcount &= 2
-#define $ARRAY_FOREACH(arr_expr, i, item_type, x, body, else_body) {\
-        array_t $arr = arr_expr; \
-        $ARRAY_INCREF($arr); \
-        if ($arr.length == 0) else_body \
-        else for (int64_t i = 1; i <= $arr.length; i++) { item_type x = *(item_type*)($arr.data + (i-1)*$arr.stride); body } \
-        $ARRAY_DECREF($arr); \
-    }
 
 #define Array__insert_value(arr, item_expr, index, type) ({ __typeof(item_expr) $item = item_expr; Array__insert(arr, &$item, index, type); })
 void Array__insert(array_t *arr, const void *item, int64_t index, const TypeInfo *type);
diff --git a/builtins/table.h b/builtins/table.h
index ab0004ec..38c66427 100644
--- a/builtins/table.h
+++ b/builtins/table.h
@@ -26,19 +26,6 @@
     if (__builtin_expect($v == NULL, 0)) \
         fail_source(filename, start, end, "The key %r is not in this table\n", generic_as_text(&$k, USE_COLOR, $info->TableInfo.key)); \
     *$v; })
-#define $TABLE_FOREACH(table_expr, key_type, k, value_type, v, value_offset, body, else_body) {\
-        array_t $entries = (table_expr).entries; \
-        if ($entries.length == 0) else_body \
-        else { \
-            $ARRAY_INCREF($entries); \
-            for (int64_t $i = 0; $i < $entries.length; $i++) { \
-                key_type k = *(key_type*)($entries.data + $i*$entries.stride); \
-                value_type v = *(value_type*)($entries.data + $i*$entries.stride + value_offset); \
-                body \
-            } \
-            $ARRAY_DECREF($entries); \
-        } \
-    }
 
 table_t Table_from_entries(array_t entries, const TypeInfo *type);
 void *Table_get(table_t t, const void *key, const TypeInfo *type);
diff --git a/compile.c b/compile.c
index 75943004..3bb8151b 100644
--- a/compile.c
+++ b/compile.c
@@ -593,27 +593,46 @@ CORD compile_statement(env_t *env, ast_t *ast)
             type_t *item_t = Match(iter_t, ArrayType)->item_type;
             CORD index = for_->index ? compile(env, for_->index) : "$i";
             CORD value = compile(env, for_->value);
-            return CORD_all("$ARRAY_FOREACH(", compile(env, for_->iter), ", ", index, ", ", compile_type(env, item_t), ", ", value, ", ",
-                            body, ", ", for_->empty ? compile_statement(env, for_->empty) : "{}", ")", stop);
+            CORD array = is_idempotent(for_->iter) ? compile(env, for_->iter) : "$arr";
+            CORD loop = CORD_all("$ARRAY_INCREF(", array, ");\n"
+                                 "for (int64_t ", index, " = 1; ", index, " <= ", array, ".length; ++", index, ") {\n",
+                                 compile_type(env, item_t), " ", value,
+                                 " = *(", compile_type(env, item_t), "*)(", array, ".data + (",index,"-1)*", array, ".stride);\n",
+                                 body, "\n}");
+            if (for_->empty)
+                loop = CORD_all("if (", array, ".length > 0) {\n", loop, "\n} else ", compile_statement(env, for_->empty));
+            loop = CORD_all(loop, stop, "\n$ARRAY_DECREF(", array, ");\n");
+            if (!is_idempotent(for_->iter))
+                loop = CORD_all("{\narray_t ",array," = ", compile(env, for_->iter), ";\n", loop, "\n}");
+            return loop;
         }
         case TableType: {
             type_t *key_t = Match(iter_t, TableType)->key_type;
             type_t *value_t = Match(iter_t, TableType)->value_type;
+
+            CORD table = is_idempotent(for_->iter) ? compile(env, for_->iter) : "$table";
+            CORD loop = CORD_all("$ARRAY_INCREF(", table, ".entries);\n"
+                                 "for (int64_t $i = 0; $i < ",table,".entries.length; ++$i) {\n");
             if (for_->index) {
-                CORD key = compile(env, for_->index);
-                CORD value = compile(env, for_->value);
+                loop = CORD_all(loop, compile_type(env, key_t), " ", compile(env, for_->index), " = *(", compile_type(env, key_t), "*)(",
+                                table,".entries.data + $i*", table, ".entries.stride);\n");
 
                 size_t value_offset = type_size(key_t);
                 if (type_align(value_t) > 1 && value_offset % type_align(value_t))
                     value_offset += type_align(value_t) - (value_offset % type_align(value_t)); // padding
-                return CORD_all("$TABLE_FOREACH(", compile(env, for_->iter), ", ", compile_type(env, key_t), ", ", key, ", ",
-                                compile_type(env, value_t), ", ", value, ", ", heap_strf("%zu", value_offset),
-                                ", ", body, ", ", for_->empty ? compile_statement(env, for_->empty) : "{}", ")", stop);
+                loop = CORD_all(loop, compile_type(env, value_t), " ", compile(env, for_->value), " = *(", compile_type(env, value_t), "*)(",
+                                table,".entries.data + $i*", table, ".entries.stride + ", heap_strf("%zu", value_offset), ");\n");
             } else {
-                CORD key = compile(env, for_->value);
-                return CORD_all("$ARRAY_FOREACH((", compile(env, for_->iter), ").entries, $i, ", compile_type(env, key_t), ", ", key, ", ",
-                                body, ", ", for_->empty ? compile_statement(env, for_->empty) : "{}", ")", stop);
+                loop = CORD_all(loop, compile_type(env, key_t), " ", compile(env, for_->value), " = *(", compile_type(env, key_t), "*)(",
+                                table,".entries.data + $i*", table, ".entries.stride);\n");
             }
+            loop = CORD_all(loop, body, "\n}");
+            if (for_->empty)
+                loop = CORD_all("if (", table, ".entries.length > 0) {\n", loop, "\n} else ", compile_statement(env, for_->empty));
+            loop = CORD_all(loop, stop, "\n$ARRAY_DECREF(", table, ".entries);\n");
+            if (!is_idempotent(for_->iter))
+                loop = CORD_all("{\ntable_t ",table," = ", compile(env, for_->iter), ";\n", loop, "\n}");
+            return loop;
         }
         case IntType: {
             CORD value = compile(env, for_->value);
diff --git a/environment.c b/environment.c
index 806bc019..b96d51eb 100644
--- a/environment.c
+++ b/environment.c
@@ -176,12 +176,13 @@ env_t *new_compilation_unit(void)
         binding_t *type_binding = Table_str_get(*env->globals, global_types[i].name);
         assert(type_binding);
         env_t *ns_env = Match(type_binding->type, TypeInfoType)->env;
-        $ARRAY_FOREACH(global_types[i].namespace, j, ns_entry_t, entry, {
-            type_t *type = parse_type_string(ns_env, entry.type_str);
+        for (int64_t j = 0; j < global_types[i].namespace.length; j++) {
+            ns_entry_t *entry = global_types[i].namespace.data + j*global_types[i].namespace.stride;
+            type_t *type = parse_type_string(ns_env, entry->type_str);
             if (type->tag == ClosureType) type = Match(type, ClosureType)->fn;
-            binding_t *b = new(binding_t, .code=entry.code, .type=type);
-            set_binding(ns_env, entry.name, b);
-        }, {})
+            binding_t *b = new(binding_t, .code=entry->code, .type=type);
+            set_binding(ns_env, entry->name, b);
+        }
     }
 
     return env;
diff --git a/parse.c b/parse.c
index a8ef5cd3..5355f91e 100644
--- a/parse.c
+++ b/parse.c
@@ -1112,8 +1112,7 @@ PARSER(parse_pass) {
 PARSER(parse_skip) {
     const char *start = pos;
     if (!match_word(&pos, "skip")) return NULL;
-    spaces(&pos);
-    const char* target;
+    const char *target;
     if (match_word(&pos, "for")) target = "for";
     else if (match_word(&pos, "while")) target = "while";
     else target = get_id(&pos);
@@ -1124,8 +1123,7 @@ PARSER(parse_skip) {
 PARSER(parse_stop) {
     const char *start = pos;
     if (!match_word(&pos, "stop")) return NULL;
-    spaces(&pos);
-    const char* target;
+    const char *target;
     if (match_word(&pos, "for")) target = "for";
     else if (match_word(&pos, "while")) target = "while";
     else target = get_id(&pos);
@@ -1136,7 +1134,6 @@ PARSER(parse_stop) {
 PARSER(parse_return) {
     const char *start = pos;
     if (!match_word(&pos, "return")) return NULL;
-    spaces(&pos);
     ast_t *value = optional(ctx, &pos, parse_expr);
     ast_t *ret = NewAST(ctx->file, start, pos, Return, .value=value);
     return ret;