added: #stack_alloc modifier to closures

Author: brendanfh

compiler/include/astnodes.h

@@ -1512,6 +1512,7 @@ struct AstCaptureBlock {
     bh_arr(AstCaptureLocal *) captures;
 
     u32 total_size_in_bytes;
+    b32 alloc_on_stack : 1;
 };
 
 struct AstCaptureLocal {

compiler/include/wasm_emit.h

@@ -683,8 +683,16 @@ typedef struct DeferredStmt {
     };
 } DeferredStmt;
 
+typedef enum AllocatedSpaceKind {
+    Allocated_Space_Kind_Local,
+    Allocated_Space_Kind_Raw
+} AllocatedSpaceKind;
+
 typedef struct AllocatedSpace {
+    AllocatedSpaceKind kind;
     u64 depth;
+    u32 size;
+    u32 alignment;
     AstTyped *expr;
 } AllocatedSpace;
 

compiler/src/parser.c

@@ -3479,8 +3479,11 @@ static AstFunction* parse_function_definition(OnyxParser* parser, OnyxToken* tok
     }
 
     if (consume_token_if_next(parser, Token_Type_Keyword_Use)) {
+        b32 stack_alloc = parse_possible_directive(parser, "stack_alloc");
+
         expect_token(parser, '(');
         func_def->captures = parse_capture_list(parser, ')');
+        if (func_def->captures) func_def->captures->alloc_on_stack = stack_alloc;
         consume_token_if_next(parser, ',');
 
         if (bh_arr_length(parser->current_function_stack) > 1) {
@@ -3683,8 +3686,11 @@ static b32 parse_possible_quick_function_definition(OnyxParser* parser, AstTyped
         }
 
         if (consume_token_if_next(parser, Token_Type_Keyword_Use)) {
+            b32 stack_alloc = parse_possible_directive(parser, "stack_alloc");
+
             expect_token(parser, '(');
             captures = parse_capture_list(parser, ')');
+            if (captures) captures->alloc_on_stack = stack_alloc;
         }
     }
 

compiler/src/wasm_emit.c

@@ -150,10 +150,7 @@ static void local_raw_free(LocalAllocator* la, WasmType wt) {
     la->freed[idx]++;
 }
 
-static u64 local_allocate_type_in_memory(LocalAllocator* la, Type *type) {
-    u32 size = type_size_of(type);
-    u32 alignment = type_alignment_of(type);
-
+static u64 local_allocate_bytes_in_memory(LocalAllocator *la, u32 size, u32 alignment) {
     bh_align(la->curr_stack, alignment);
 
     if (la->max_stack < la->curr_stack)
@@ -172,12 +169,24 @@ static u64 local_allocate_type_in_memory(LocalAllocator* la, Type *type) {
     return la->curr_stack - size;
 }
 
+static void local_free_bytes_in_memory(LocalAllocator* la, u32 size, u32 alignment) {
+    bh_align(size, alignment);
+
+    la->curr_stack -= size;
+}
+
+static u64 local_allocate_type_in_memory(LocalAllocator* la, Type *type) {
+    u32 size = type_size_of(type);
+    u32 alignment = type_alignment_of(type);
+
+    return local_allocate_bytes_in_memory(la, size, alignment);
+}
+
 static void local_free_type_in_memory(LocalAllocator* la, Type* type) {
     u32 size = type_size_of(type);
     u32 alignment = type_alignment_of(type);
-    bh_align(size, alignment);
 
-    la->curr_stack -= size;
+    local_free_bytes_in_memory(la, size, alignment);
 }
 
 static u64 local_allocate(LocalAllocator* la, AstTyped* local) {
@@ -769,6 +778,19 @@ EMIT_FUNC(statement, AstNode* stmt) {
     *pcode = code;
 }
 
+EMIT_FUNC_RETURNING(u64, stack_alloc, u32 size, u32 alignment) {
+    u32 local_idx = local_allocate_bytes_in_memory(mod->local_alloc, size, alignment);
+
+    bh_arr_push(mod->local_allocations, ((AllocatedSpace) {
+        .kind = Allocated_Space_Kind_Raw,
+        .depth = bh_arr_length(mod->structured_jump_target),
+        .size = size,
+        .alignment = alignment
+    }));
+
+    return local_idx;
+}
+
 EMIT_FUNC_RETURNING(u64, local_allocation, AstTyped* stmt) {
     //
     // If the statement does not have a type, it should not
@@ -809,6 +831,7 @@ EMIT_FUNC_RETURNING(u64, local_allocation, AstTyped* stmt) {
     }
 
     bh_arr_push(mod->local_allocations, ((AllocatedSpace) {
+        .kind = Allocated_Space_Kind_Local,
         .depth = bh_arr_length(mod->structured_jump_target),
         .expr  = stmt,
     }));
@@ -821,10 +844,19 @@ EMIT_FUNC_NO_ARGS(free_local_allocations) {
 
     u64 depth = bh_arr_length(mod->structured_jump_target);
     while (bh_arr_length(mod->local_allocations) > 0 && bh_arr_last(mod->local_allocations).depth >= depth) {
-        // CHECK: Not sure this next line is okay to be here...
-        bh_imap_delete(&mod->local_map, (u64) bh_arr_last(mod->local_allocations).expr);
+        AllocatedSpace alloc = bh_arr_last(mod->local_allocations);
+        switch (alloc.kind) {
+            case Allocated_Space_Kind_Local:
+                // CHECK: Not sure this next line is okay to be here...
+                bh_imap_delete(&mod->local_map, (u64) alloc.expr);
+                local_free(mod->local_alloc, alloc.expr);
+                break;
+
+            case Allocated_Space_Kind_Raw:
+                local_free_bytes_in_memory(mod->local_alloc, alloc.size, alloc.alignment);
+                break;
+        }
 
-        local_free(mod->local_alloc, bh_arr_last(mod->local_allocations).expr);
         bh_arr_pop(mod->local_allocations);
     }
 }
@@ -2085,7 +2117,7 @@ EMIT_FUNC(call, AstCall* call) {
         WIL(NULL, WI_GLOBAL_SET, stack_trace_pass_global);
     }
 
-    if (call->callee->kind == Ast_Kind_Function) {
+    if (call->callee->kind == Ast_Kind_Function && !((AstFunction *) call->callee)->captures) {
         CodePatchInfo code_patch;
         code_patch.kind = Code_Patch_Callee;
         code_patch.func_idx = mod->current_func_idx;
@@ -3349,17 +3381,23 @@ EMIT_FUNC(expression, AstTyped* expr) {
             }
 
             // Allocate the block
-            WIL(NULL, WI_I32_CONST, func->captures->total_size_in_bytes);
+            if (func->captures->alloc_on_stack) {
+                u32 address = emit_stack_alloc(mod, &code, func->captures->total_size_in_bytes, 16);
+                emit_stack_address(mod, &code, address, NULL);
 
-            CodePatchInfo code_patch;
-            code_patch.kind = Code_Patch_Callee;
-            code_patch.func_idx = mod->current_func_idx;
-            code_patch.instr = bh_arr_length(code);
-            code_patch.node_related_to_patch = (AstNode *) mod->context->builtins.closure_block_allocate;
-            bh_arr_push(mod->code_patches, code_patch);
-            WIL(NULL, WI_CALL, 0);
+            } else {
+                WIL(NULL, WI_I32_CONST, func->captures->total_size_in_bytes);
+
+                CodePatchInfo code_patch;
+                code_patch.kind = Code_Patch_Callee;
+                code_patch.func_idx = mod->current_func_idx;
+                code_patch.instr = bh_arr_length(code);
+                code_patch.node_related_to_patch = (AstNode *) mod->context->builtins.closure_block_allocate;
+                bh_arr_push(mod->code_patches, code_patch);
+                WIL(NULL, WI_CALL, 0);
 
-            ensure_node_has_been_submitted_for_emission(mod->context, (AstNode *) mod->context->builtins.closure_block_allocate);
+                ensure_node_has_been_submitted_for_emission(mod->context, (AstNode *) mod->context->builtins.closure_block_allocate);
+            }
 
             u64 capture_block_ptr = local_raw_allocate(mod->local_alloc, WASM_TYPE_PTR);
             WIL(NULL, WI_LOCAL_TEE, capture_block_ptr);

tests/closure_stack_alloc

@@ -0,0 +1,14 @@
+1
+4
+6
+8
+9
+10
+12
+13
+14
+15
+16
+17
+18
+19

tests/closure_stack_alloc.onyx

@@ -0,0 +1,27 @@
+use core {*}
+
+main :: () {
+    // Ensure that any use of closures will break the program
+    context.closure_allocate = n => {
+        panic("LAND MINE")
+        return cast(rawptr, 0)
+    }
+
+    a := Array.make(.[2, 3, 5, 7, 11])
+
+    res := Iterator.from(1 .. 20)
+        // While the Iterator itself will allocate in order to
+        // store its context, the closure will not allocate.
+        |> Iterator.filter((x) use #stack_alloc (a) => {
+            // This tests using a stack allocated closure in a macro,
+            // that is called directly (() use ...)(), which caused
+            // issues in the initial implementation.
+            return !Slice.some(a, (y) use #stack_alloc (x) => {
+                return x == y
+            })
+        })
+
+    for i in res {
+        println(i)
+    }
+}