Add for-else block

4 files changed, 58 insertions, 30 deletions

diff --git a/ast.c b/ast.c
index df12a456..60cd27b6 100644
--- a/ast.c
+++ b/ast.c
@@ -121,8 +121,8 @@ CORD ast_to_cord(ast_t *ast)
     T(Lambda, "(args=%r, body=%r)", arg_list_to_cord(data.args), ast_to_cord(data.body))
     T(FunctionCall, "(fn=%r, args=%r)", ast_to_cord(data.fn), arg_list_to_cord(data.args))
     T(Block, "(%r)", ast_list_to_cord(data.statements))
-    T(For, "(index=%r, value=%r, iter=%r, body=%r)", ast_to_cord(data.index), ast_to_cord(data.value),
-      ast_to_cord(data.iter), ast_to_cord(data.body))
+    T(For, "(index=%r, value=%r, iter=%r, body=%r, empty=%r)", ast_to_cord(data.index), ast_to_cord(data.value),
+      ast_to_cord(data.iter), ast_to_cord(data.body), ast_to_cord(data.empty))
     T(While, "(condition=%r, body=%r)", ast_to_cord(data.condition), ast_to_cord(data.body))
     T(If, "(condition=%r, body=%r, else=%r)", ast_to_cord(data.condition), ast_to_cord(data.body), ast_to_cord(data.else_body))
     T(When, "(subject=%r, clauses=%r, else=%r)", ast_to_cord(data.subject), when_clauses_to_cord(data.clauses), ast_to_cord(data.else_body))
@@ -137,8 +137,9 @@ CORD ast_to_cord(ast_t *ast)
     T(Index, "(indexed=%r, index=%r)", ast_to_cord(data.indexed), ast_to_cord(data.index))
     T(FieldAccess, "(fielded=%r, field=%s)", ast_to_cord(data.fielded), data.field)
     T(DocTest, "(expr=%r, output=%r)", ast_to_cord(data.expr), Str__quoted(data.output, true))
-    T(Use, "(%s)", Str__quoted(data.path, true))
-    T(LinkerDirective, "(%s)", Str__quoted(data.directive, true))
+    T(Use, "(%r)", Str__quoted(data.path, true))
+    T(LinkerDirective, "(%r)", Str__quoted(data.directive, true))
+    T(InlineCCode, "(%r)", Str__quoted(data.code, true))
 #undef T
     }
     return "???";
diff --git a/ast.h b/ast.h
index 254ac53a..ded91f90 100644
--- a/ast.h
+++ b/ast.h
@@ -111,6 +111,7 @@ typedef enum {
     DocTest,
     Use,
     LinkerDirective,
+    InlineCCode,
 } ast_e;
 
 struct ast_s {
@@ -193,7 +194,7 @@ struct ast_s {
             ast_list_t *statements;
         } Block;
         struct {
-            ast_t *index, *value, *iter, *body;
+            ast_t *index, *value, *iter, *body, *empty;
         } For;
         struct {
             ast_t *condition, *body;
@@ -253,6 +254,9 @@ struct ast_s {
         struct {
             const char *directive;
         } LinkerDirective;
+        struct {
+            CORD code;
+        } InlineCCode;
     } __data;
 };
 
diff --git a/compile.c b/compile.c
index ea576ae1..62da8af9 100644
--- a/compile.c
+++ b/compile.c
@@ -711,13 +711,17 @@ CORD compile(env_t *env, ast_t *ast)
                 set_binding(scope, CORD_to_const_char_star(index), new(binding_t, .type=Type(IntType, .bits=64)));
             CORD value = compile(env, for_->value);
             set_binding(scope, CORD_to_const_char_star(value), new(binding_t, .type=item_t));
-            return CORD_all(
-                "{ // For loop:\n"
-                "array_t $iter = ", compile(env, for_->iter), ";\n"
+            CORD loop = CORD_all(
                 "for (Int64_t ", index, " = 1; ", index, " <= $iter.length; ++", index, ") {\n"
                 "\t", compile_type(item_t), " ", value, " = *(", compile_type(item_t), "*)($iter.data + (", index, "-1)*$iter.stride);\n"
                 "\t", compile(scope, for_->body), "\n"
-                "}\n"
+                "}\n");
+            if (for_->empty)
+                loop = CORD_all("if ($iter.length == 0)\n", compile(env, for_->empty), "\nelse\n", loop);
+            return CORD_all(
+                "{ // For loop:\n"
+                "array_t $iter = ", compile(env, for_->iter), ";\n",
+                loop,
                 "}\n");
         }
         case TableType: {
@@ -734,25 +738,33 @@ CORD compile(env_t *env, ast_t *ast)
                 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(
-                    "{ // For loop:\n"
-                    "array_t $entries = (", compile(env, for_->iter), ").entries;\n"
+                CORD loop = CORD_all(
                     "for (Int64_t $offset = 0; $offset < $entries.length; ++$offset) {\n"
                     "\t", compile_type(key_t), " ", key, " = *(", compile_type(key_t), "*)($entries.data + $offset*$entries.stride);\n"
                     "\t", compile_type(value_t), " ", value, " = *(", compile_type(value_t), "*)($entries.data + $offset*$entries.stride + ", CORD_asprintf("%zu", value_offset), ");\n"
-                    "\t", compile(scope, for_->body), "\n"
+                    "\t", compile(scope, for_->body), "\n");
+                if (for_->empty)
+                    loop = CORD_all("if ($iter.length == 0)\n", compile(env, for_->empty), "\nelse\n", loop);
+                return CORD_all(
+                    "{ // For loop:\n"
+                    "array_t $entries = (", compile(env, for_->iter), ").entries;\n",
+                    loop,
                     "}\n"
                     "}\n");
             } else {
                 key = compile(env, for_->value);
                 set_binding(scope, CORD_to_const_char_star(key), new(binding_t, .type=key_t));
-                return CORD_all(
-                    "{ // For loop:\n"
-                    "array_t $entries = (", compile(env, for_->iter), ").entries;\n"
+                CORD loop = CORD_all(
                     "for (Int64_t $offset = 0; $offset < $entries.length; ++$offset) {\n"
                     "\t", compile_type(key_t), " ", key, " = *(", compile_type(key_t), "*)($entries.data + $offset*$entries.stride);\n"
                     "\t", compile(scope, for_->body), "\n"
-                    "}\n"
+                    "}\n");
+                if (for_->empty)
+                    loop = CORD_all("if ($iter.length == 0)\n", compile(env, for_->empty), "\nelse\n", loop);
+                return CORD_all(
+                    "{ // For loop:\n"
+                    "array_t $entries = (", compile(env, for_->iter), ").entries;\n",
+                    loop,
                     "}\n");
             }
         }
@@ -763,6 +775,8 @@ CORD compile(env_t *env, ast_t *ast)
                 code_err(for_->index, "It's redundant to have a separate iteration index");
             CORD value = compile(env, for_->value);
             set_binding(scope, CORD_to_const_char_star(value), new(binding_t, .type=item_t));
+            if (for_->empty)
+                code_err(for_->empty, "'else' is not implemented for loops over integers");
             return CORD_all(
                 "for (Int64_t ", value, " = 1, $n = ", compile(env, for_->iter), "; ", value, " <= $n; ++", value, ")\n"
                 "\t", compile(scope, for_->body), "\n");
@@ -775,34 +789,33 @@ CORD compile(env_t *env, ast_t *ast)
         type_t *t = get_type(env, ast);
         CORD code = CORD_all(
             "({ // Reduction:\n",
-            "array_t $reduction_iter = ", compile(env, reduction->iter), ";\n",
             compile_type(t), " $lhs;\n"
             );
-        ast_t *iter = FakeAST(Var, "$reduction_iter");
         env_t *scope = fresh_scope(env);
-        set_binding(scope, "$reduction_iter", new(binding_t, .type=get_type(env, reduction->iter)));
-        CORD is_empty = compile(scope, FakeAST(BinaryOp, .lhs=FakeAST(Length, iter), .op=BINOP_EQ, .rhs=FakeAST(Int, .i=0, .bits=64)));
+        ast_t *result = FakeAST(Var, "$lhs");
+        set_binding(scope, "$lhs", new(binding_t, .type=t));
+        ast_t *empty = NULL;
         if (reduction->fallback) {
             type_t *fallback_type = get_type(scope, reduction->fallback);
             if (fallback_type->tag == AbortType) {
-                code = CORD_all(code, "if (", is_empty, ")\n\t{", compile_statement(scope, reduction->fallback), "}\n");
+                empty = reduction->fallback;
             } else {
-                code = CORD_all(code, "if (", is_empty, ")\n\t{$lhs = ", compile(scope, reduction->fallback), ";}\n");
+                empty = FakeAST(Assign, .targets=new(ast_list_t, .ast=result), .values=new(ast_list_t, .ast=reduction->fallback));
             }
         } else {
-            CORD_appendf(&code, "if (%r)\n\t{fail_source(%s, %ld, %ld, \"This collection was empty!\");}\n",
-                         is_empty, Str__quoted(ast->file->filename, false), (long)(reduction->iter->start - reduction->iter->file->text),
-                         (long)(reduction->iter->end - reduction->iter->file->text));
+            empty = FakeAST(
+                InlineCCode, 
+                CORD_asprintf("fail_source(%s, %ld, %ld, \"This collection was empty!\");\n",
+                              Str__quoted(ast->file->filename, false), (long)(reduction->iter->start - reduction->iter->file->text),
+                              (long)(reduction->iter->end - reduction->iter->file->text)));
         }
         ast_t *i = FakeAST(Var, "$i");
         ast_t *item = FakeAST(Var, "$rhs");
-        ast_t *result = FakeAST(Var, "$lhs");
         ast_t *body = FakeAST(
             If, .condition=FakeAST(BinaryOp, .lhs=i, .op=BINOP_EQ, .rhs=FakeAST(Int, .i=1, .bits=64)),
             .body=FakeAST(Assign, .targets=new(ast_list_t, .ast=result), .values=new(ast_list_t, .ast=item)),
             .else_body=FakeAST(Assign, .targets=new(ast_list_t, .ast=result), .values=new(ast_list_t, .ast=reduction->combination)));
-        ast_t *loop = FakeAST(For, .index=i, .value=item, .iter=iter, .body=body);
-        set_binding(scope, "$lhs", new(binding_t, .type=t));
+        ast_t *loop = FakeAST(For, .index=i, .value=item, .iter=reduction->iter, .body=body, .empty=empty);
         set_binding(scope, "$rhs", new(binding_t, .type=t));
         code = CORD_all(code, compile(scope, loop), "\n$lhs;})");
         return code;
@@ -989,6 +1002,7 @@ CORD compile(env_t *env, ast_t *ast)
     }
     // Use,
     // LinkerDirective,
+    case InlineCCode: return Match(ast, InlineCCode)->code;
     case Unknown: code_err(ast, "Unknown AST");
     default: break;
     }
diff --git a/parse.c b/parse.c
index 685461da..5d1880c9 100644
--- a/parse.c
+++ b/parse.c
@@ -847,6 +847,7 @@ PARSER(parse_for) {
     // for [k,] v in iter [<indent>] body
     const char *start = pos;
     if (!match_word(&pos, "for")) return NULL;
+    int64_t starting_indent = get_indent(ctx->file, pos);
     ast_t *index = expect(ctx, start, &pos, parse_var, "I expected an iteration variable for this 'for'");
     spaces(&pos);
     ast_t *value = NULL;
@@ -857,7 +858,15 @@ PARSER(parse_for) {
     ast_t *iter = expect(ctx, start, &pos, parse_expr, "I expected an iterable value for this 'for'");
     match(&pos, "do"); // optional
     ast_t *body = expect(ctx, start, &pos, parse_opt_indented_block, "I expected a body for this 'for'"); 
-    return NewAST(ctx->file, start, pos, For, .index=value ? index : NULL, .value=value ? value : index, .iter=iter, .body=body);
+
+    const char *else_start = pos;
+    whitespace(&else_start);
+    ast_t *empty = NULL;
+    if (match_word(&else_start, "else") && get_indent(ctx->file, else_start) == starting_indent) {
+        pos = else_start;
+        empty = expect(ctx, pos, &pos, parse_opt_indented_block, "I expected a body for this 'else'");
+    }
+    return NewAST(ctx->file, start, pos, For, .index=value ? index : NULL, .value=value ? value : index, .iter=iter, .body=body, .empty=empty);
 }
 
 PARSER(parse_while) {