Merge pull request #182 from onyx-lang/experiment/for-macros

Experiment: For-loops as macro expansions

Author: brendanfh

compiler/include/astnodes.h

@@ -866,9 +866,7 @@ struct AstFor           {
     // NOTE: Stores the iteration variable
     Scope *scope;
 
-    // NOTE: Local defining the iteration variable
-    AstLocal* var;
-    AstLocal* index_var;
+    bh_arr(AstLocal *) indexing_variables;
 
     // NOTE: This can be any expression, but it is checked that
     // it is of a type that we know how to iterate over.
@@ -877,16 +875,18 @@ struct AstFor           {
 
     AstBlock *stmt;
 
+    // NOTE: This is set when a for-loop isn't over a primitive type
+    // and instead is over a custom type, such as `Iterator` or `Map`.
+    // To properly invoke `__for_expansion`, we need to store the prepared
+    // call, as there could be overloads we have to wait for by yielding.
+    AstCall *intermediate_macro_expansion;
+
     // ROBUSTNESS: This should be able to be set by a compile time variable at some point.
     // But for now, this will do.
     b32 by_pointer : 1;
     b32 no_close   : 1;
-    b32 has_first  : 1;
-
-    // NOTE: This is used by the AstDirectiveFirst node for this
-    // for node to know which local variable to use.
-    u64 first_local;
 };
+
 struct AstIfWhile {
     AstNode_base;
 
@@ -907,7 +907,14 @@ struct AstIfWhile {
         };
 
         // Used by While
-        b32 bottom_test;
+        struct {
+            // NOTE: This is used by the AstDirectiveFirst node for this
+            // for node to know which local variable to use.
+            u64 first_local;
+            b32 has_first;
+
+            b32 bottom_test;
+        };
     };
 };
 typedef struct AstIfWhile AstIf;
@@ -1572,7 +1579,8 @@ struct AstDirectiveRemove {
 
 struct AstDirectiveFirst {
     AstTyped_base;
-    AstFor *for_node;
+
+    AstIfWhile *while_node;
 };
 
 struct AstDirectiveExportName {
@@ -1607,12 +1615,17 @@ struct AstCallSite {
     b32 collapsed : 1;
 };
 
+typedef struct CodeBlockBindingSymbol {
+    OnyxToken *symbol;
+    AstType   *type_node; // This can be NULL if no type was given.
+} CodeBlockBindingSymbol;
+
 // Represents a "pastable" block of code.
 struct AstCodeBlock {
     AstTyped_base;
 
     AstNode *code;
-    bh_arr(OnyxToken *) binding_symbols;
+    bh_arr(CodeBlockBindingSymbol) binding_symbols;
 
     b32 is_expression: 1;
 };
@@ -1622,6 +1635,9 @@ struct AstDirectiveInsert {
 
     AstTyped *code_expr;
     bh_arr(AstTyped *) binding_exprs;
+
+    // Set when using #skip_scope
+    AstTyped *skip_scope_index;
 };
 
 struct AstDirectiveInit {
@@ -1921,8 +1937,7 @@ typedef enum CheckerMode {
 typedef struct CheckerData {
     b32 expression_types_must_be_known;
     b32 all_checks_are_final;
-    b32 inside_for_iterator;
-    bh_arr(AstFor *) for_node_stack;
+    bh_arr(AstIfWhile *) while_node_stack;
     bh_imap __binop_impossible_cache[Binary_Op_Count];
     AstCall __op_maybe_overloaded;
     Entity *current_entity;
@@ -2109,6 +2124,9 @@ struct CompilerBuiltins {
     bh_arr(AstFunction *) init_procedures;
     AstOverloadedFunction *implicit_bool_cast;
     AstOverloadedFunction *dispose_used_local;
+
+    AstType *for_expansion_flag_type;
+    AstOverloadedFunction *for_expansion;
 };
 
 typedef struct TypeStore TypeStore;
@@ -2361,6 +2379,8 @@ AstPolyCallType* convert_call_to_polycall(Context *context, AstCall* call);
 
 void insert_auto_dispose_call(Context *context, AstLocal *local);
 
+AstCall * create_implicit_for_expansion_call(Context *context, AstFor *fornode);
+
 typedef struct OverloadReturnTypeCheck {
     Type *expected_type;
     AstTyped *node;

compiler/src/astnodes.c

@@ -1948,6 +1948,61 @@ void insert_auto_dispose_call(Context *context, AstLocal *local) {
 }
 
 
+AstCall * create_implicit_for_expansion_call(Context *context, AstFor *fornode) {
+    AstCall *call = onyx_ast_node_new(context->ast_alloc, sizeof(AstCall), Ast_Kind_Call);
+    call->token = fornode->token;
+    call->callee = (AstTyped *) context->builtins.for_expansion;
+    call->next = fornode->next;
+
+    arguments_initialize(context, &call->args);
+
+    //
+    // Create the code block that will represent the body of the for-loop.
+    // This code block is given up to 2 inputs, depending on if the index variable
+    // was set in the for-loop. The implementer of a __for_expansion overload must
+    // always provide 2 values when `#unquote`-ing, as they cannot currently know if
+    // the index variable was asked for.
+    //
+    // Maybe we could pass `void` as the `index_type` if the index variable is not necessary?
+    // I would like to keep the implementations of __for_expansion simple and easy
+    // to read, but sometimes there are extra complications you cannot avoid...
+    //
+    AstCodeBlock *body_code_block = onyx_ast_node_new(context->ast_alloc, sizeof(AstCodeBlock), Ast_Kind_Code_Block);
+    body_code_block->token = fornode->token;
+    body_code_block->type_node = context->builtins.code_type;
+    body_code_block->code = (AstNode *) fornode->stmt;
+    ((AstBlock *) body_code_block->code)->rules = Block_Rule_Code_Block;
+
+    bh_arr_new(context->ast_alloc, body_code_block->binding_symbols, bh_arr_length(fornode->indexing_variables));
+    bh_arr_each(AstLocal *, indexing_variable, fornode->indexing_variables) {
+        CodeBlockBindingSymbol sym;
+        sym.symbol    = (*indexing_variable)->token;
+        sym.type_node = (*indexing_variable)->type_node;
+        bh_arr_push(body_code_block->binding_symbols, sym);
+    }
+
+    i32 flags = 0;
+    if (fornode->by_pointer) flags |= 1; // BY_POINTER
+    if (fornode->no_close)   flags |= 2; // NO_CLOSE
+    
+    AstNumLit *flag_node = make_int_literal(context, flags);
+    flag_node->type_node = context->builtins.for_expansion_flag_type;
+    // flag_node->type = type_build_from_ast(context, context->builtins.for_expansion_flag_type);
+    // assert(flag_node->type);
+
+    // Arguments are: 
+    //    Iterator
+    //    Code block with 2 inputs (value, index)
+    //    Flags
+    bh_arr_push(call->args.values, (AstTyped *) make_argument(context, (AstTyped *) fornode->iter));
+    bh_arr_push(call->args.values, (AstTyped *) make_argument(context, (AstTyped *) flag_node));
+    bh_arr_push(call->args.values, (AstTyped *) make_argument(context, (AstTyped *) body_code_block));
+
+    return call;
+}
+
+
+
 b32 resolve_intrinsic_interface_constraint(Context *context, AstConstraint *constraint) {
     AstInterface *interface = constraint->interface;
     Type* type = type_build_from_ast(context, (AstType *) constraint->args[0]);

compiler/src/builtins.c

@@ -511,6 +511,18 @@ void initialize_builtins(Context *context) {
         return;
     }
 
+    context->builtins.for_expansion = (AstOverloadedFunction *) symbol_raw_resolve(context, p->scope, "__for_expansion");
+    if (context->builtins.for_expansion == NULL || context->builtins.for_expansion->kind != Ast_Kind_Overloaded_Function) {
+        ONYX_ERROR((OnyxFilePos) { 0 }, Error_Critical, "'__for_expansion' #match procedure not found.");
+        return;
+    }
+
+    context->builtins.for_expansion_flag_type = (AstType *) symbol_raw_resolve(context, p->scope, "__For_Expansion_Flags");
+    if (context->builtins.for_expansion_flag_type == NULL || context->builtins.for_expansion_flag_type->kind != Ast_Kind_Enum_Type) {
+        ONYX_ERROR((OnyxFilePos) { 0 }, Error_Critical, "'__For_Expansion_Flags' enum procedure not found.");
+        return;
+    }
+
     context->builtins.closure_block_allocate = (AstFunction *) symbol_raw_resolve(context, p->scope, "__closure_block_allocate");
     if (context->builtins.closure_block_allocate == NULL || context->builtins.closure_block_allocate->kind != Ast_Kind_Function) {
         ONYX_ERROR((OnyxFilePos) { 0 }, Error_Critical, "'__closure_block_allocate' procedure not found.");