diff --git a/examples/133-closure-env-routes-through-context-allocator.sx b/examples/133-closure-env-routes-through-context-allocator.sx new file mode 100644 index 0000000..032a403 --- /dev/null +++ b/examples/133-closure-env-routes-through-context-allocator.sx @@ -0,0 +1,47 @@ +// Phase 1.3 — the closure env-buffer heap-copy in `lowerLambda` must +// dispatch through `context.allocator`, not `.heap_alloc` directly. +// So when a `push Context.{ allocator = tracer }` block is active, a +// capturing closure created inside it MUST allocate its env through +// the tracker. +// +// Mirrors the shape of `130-xx-value-routes-through-context-allocator.sx` +// for the protocol-erasure heap path — same Tracer, same install via +// `push Context`, same `Tracer.count = 1` assertion. Different +// allocation site (closure env vs xx-value heap copy). + +#import "modules/std.sx"; + +Tracer :: struct { + count: s64; + + init :: () -> *Tracer { + t : *Tracer = xx libc_malloc(size_of(Tracer)); + t.count = 0; + t; + } +} + +impl Allocator for Tracer { + alloc :: (self: *Tracer, size: s64) -> *void { + self.count += 1; + return libc_malloc(size); + } + dealloc :: (self: *Tracer, ptr: *void) { + libc_free(ptr); + } +} + +main :: () -> s32 { + tracer := Tracer.init(); + push Context.{ allocator = xx tracer, data = null } { + // Capturing closure. lowerLambda allocates an env struct on the + // stack, copies the captures in, then heap-copies the env via + // `allocViaContext` — which dispatches through the installed + // tracer's `alloc`. + captured : s64 = 100; + add_capture := closure((y: s64) -> s64 => y + captured); + _ = add_capture(1); + } + print("Tracer.count = {}\n", tracer.count); + 0; +} diff --git a/src/ir/lower.zig b/src/ir/lower.zig index f81dab6..abe9562 100644 --- a/src/ir/lower.zig +++ b/src/ir/lower.zig @@ -5828,6 +5828,15 @@ pub const Lowering = struct { self.builder.func = saved_func; self.builder.current_block = saved_block; self.builder.inst_counter = saved_counter; + // Restore the caller's `current_ctx_ref` BEFORE we emit the env + // alloc/memcpy below — those run in the caller's scope, and + // `allocViaContext` reads `current_ctx_ref` to find the + // installed allocator. Without this, the env_heap dispatch + // would still see `Ref.fromIndex(0)` (the lambda's own ctx + // param), which doesn't exist in the caller's frame and + // silently routes through the default context instead of any + // surrounding `push Context.{ allocator = ... }`. + self.current_ctx_ref = saved_ctx_ref_lam; // Create proper closure type (user-visible params only — skip ctx + env). const skip_count: usize = if (lambda_wants_ctx) 2 else 1; @@ -5852,11 +5861,15 @@ pub const Lowering = struct { self.builder.store(gep, val); } - // Copy env to heap (so it outlives the stack frame) + // Copy env to heap (so it outlives the stack frame). + // Route through `context.allocator.alloc` rather than calling + // libc malloc directly so closures respect a surrounding + // `push Context.{ allocator = ... }` and a tracker / arena + // counts the env allocation alongside everything else. const env_byte_size = self.computeEnvSize(capture_list); const env_size = self.builder.constInt(@intCast(env_byte_size), .s64); const ptr_void = self.module.types.ptrTo(.void); - const env_heap = self.builder.emit(.{ .heap_alloc = .{ .operand = env_size } }, ptr_void); + const env_heap = self.allocViaContext(env_size, ptr_void); // memcpy(heap, stack_alloca, size) _ = self.callForeign("memcpy", &.{ env_heap, env_local, env_size }, ptr_void); diff --git a/tests/expected/133-closure-env-routes-through-context-allocator.exit b/tests/expected/133-closure-env-routes-through-context-allocator.exit new file mode 100644 index 0000000..573541a --- /dev/null +++ b/tests/expected/133-closure-env-routes-through-context-allocator.exit @@ -0,0 +1 @@ +0 diff --git a/tests/expected/133-closure-env-routes-through-context-allocator.txt b/tests/expected/133-closure-env-routes-through-context-allocator.txt new file mode 100644 index 0000000..06be9f3 --- /dev/null +++ b/tests/expected/133-closure-env-routes-through-context-allocator.txt @@ -0,0 +1 @@ +Tracer.count = 1 diff --git a/tests/expected/50-smoke.txt b/tests/expected/50-smoke.txt index 4c86871..4ab0c8c 100644 --- a/tests/expected/50-smoke.txt +++ b/tests/expected/50-smoke.txt @@ -517,7 +517,7 @@ closure-form: no cancel closure-null-env: true closure-slice: 10 20 30 closure-arena: 15 -closure-gpa: 17 allocs=-1 +closure-gpa: 17 allocs=0 closure-opt: 42 closure-ropt: 50 closure-ropt: none