From 1cb40e8b65f1c749eada2e4155085cb50174e705 Mon Sep 17 00:00:00 2001 From: Tim Whiting Date: Sat, 24 Jan 2026 19:51:37 -0700 Subject: [PATCH] Improve refcount for continuations --- lib/std/core/hnd.kk | 52 ++++++++++++++++++-------------- lib/std/core/inline/hnd.c | 44 +++++++++++++++++++-------- test/cgen/simple-refcount.kk | 22 ++++++++++++++ test/cgen/simple-refcount.kk.out | 2 ++ 4 files changed, 84 insertions(+), 36 deletions(-) create mode 100644 test/cgen/simple-refcount.kk create mode 100644 test/cgen/simple-refcount.kk.out diff --git a/lib/std/core/hnd.kk b/lib/std/core/hnd.kk index 1fc25ab2c..2213adace 100644 --- a/lib/std/core/hnd.kk +++ b/lib/std/core/hnd.kk @@ -603,31 +603,37 @@ fun get( ref: ref) : ,div> a inline extern unsafe-st(f : () -> |e> a ) : (() -> e a) inline "#1" -fun protect-prompt( resumed : ref, k : resume-result -> e r, res : r ) : e r - val did-resume : bool = (unsafe-st{ !resumed })() - if did-resume then - // if resumed, we no longer need to protect - res - elif !yielding() then - // otherwise, if we are not yielding, resume k with finalization (to run all finally clauses) - k(Finalize(res)) - elif yielding-non-final() then - // if we yield non-final to an operation, extend the continuation with this prompt (so we keep protecting after being resumed) - yield-cont( fn(cont,x) protect-prompt(pretend-decreasing(resumed),k,cont(x)) ) - else - // if we are in a final yield, capture it, resume k with finalization, and reyield - val yld = yield-capture() - k(Finalize(res)) - if yielding() return yield-extend( fn(_x) unsafe-reyield(yld) ) // yikes, a finally clause is itself yielding... - unsafe-reyield(yld) +type protect-state + NeedsFinalization(k : resume-result -> e r) + NoFinalization + +fun protect-prompt( resumed : ref>, res : r) : e r + val hdiv=@Hnodiv // Needed since resumed is in global state. + val did-resume : some protect-state = (unsafe-st{ !resumed })() + match did-resume + NoFinalization -> // resumed already, no need to protect + res + NeedsFinalization(k) -> + if !yielding() then + // otherwise, if we are not yielding, resume k with finalization (to run all finally clauses) + k(Finalize(res)) + elif yielding-non-final() then + // if we yield non-final to an operation, extend the continuation with this prompt (so we keep protecting after being resumed) + yield-cont( fn(cont,x) protect-prompt(pretend-decreasing(resumed),cont(x)) ) + else + // if we are in a final yield, capture it, resume k with finalization, and reyield + val yld = yield-capture() + k(Finalize(res)) + if yielding() return yield-extend( fn(_x) unsafe-reyield(yld) ) // yikes, a finally clause is itself yielding... + unsafe-reyield(yld) fun protect( x : a, clause : (x:a, k: b -> e r) -> e r, k : resume-result -> e r ) : e r - val resumed = (unsafe-st{ref(False)})() + val resumed = (unsafe-st{ref(NeedsFinalization(k))})() fun kprotect(ret) - (unsafe-st{resumed := True})() + (unsafe-st{resumed := NoFinalization})() k(Deep(ret)) val res = clause(x,kprotect) - protect-prompt(resumed,k,res) + protect-prompt(resumed,res) /* pub fun clause-control1( clause : (x:a, k: b -> e r) -> e r ) : clause1 @@ -715,12 +721,12 @@ fun under2( ev : ev, op : (a1,a2) -> e b, x1 : a1, x2 : a2 ) : e b z fun protect2( x1 : a1, x2:a2, clause : (x:a1,x:a2, k: b -> e r) -> e r, k : resume-result -> e r ) : e r - val resumed = (unsafe-st{ref(False)})() + val resumed = (unsafe-st{ref(NeedsFinalization(k))})() fun kprotect(ret) - (unsafe-st{ resumed := True })() + (unsafe-st{ resumed := NoFinalization })() k(Deep(ret)) val res = clause(x1,x2,kprotect) - protect-prompt(resumed,k,res) + protect-prompt(resumed,res) pub fun clause-control2( clause : (x1:a1, x2:a2, k: b -> e r) -> e r ) : clause2 Clause2(fn(m,_ev,x1,x2){ yield-to(m, fn(k){ protect2(x1,x2,clause,k) }) }) diff --git a/lib/std/core/inline/hnd.c b/lib/std/core/inline/hnd.c index c6336872b..4fb98aebb 100644 --- a/lib/std/core/inline/hnd.c +++ b/lib/std/core/inline/hnd.c @@ -268,22 +268,40 @@ static kk_box_t kcompose( kk_function_t fself, kk_box_t x, kk_context_t* ctx) { kk_intx_t count = kk_intf_unbox(self->count); kk_function_t* conts = &self->conts[0]; // call each continuation in order - for(kk_intx_t i = 0; i < count; i++) { - // todo: take uniqueness of fself into account to avoid dup_function - kk_function_t f = kk_function_dup(conts[i],ctx); - x = kk_function_call(kk_box_t, (kk_function_t, kk_box_t, kk_context_t*), f, (f, x, ctx), ctx); - if (kk_yielding(ctx)) { - // if yielding, `yield_next` all continuations that still need to be done - while(++i < count) { - // todo: if fself is unique, we could copy without dup? - kk_yield_extend(kk_function_dup(conts[i],ctx),ctx); + if kk_likely(kk_datatype_ptr_is_unique(fself, ctx)) { + // Special handling for unique continuation function to avoid dup/drop overhead for the continuation and captured variables. + for(kk_intx_t i = 0; i < count; i++) { + kk_function_t f = conts[i]; + x = kk_function_call(kk_box_t, (kk_function_t, kk_box_t, kk_context_t*), f, (f, x, ctx), ctx); + if (kk_yielding(ctx)) { + // if yielding, `yield_next` all continuations that still need to be done + while(++i < count) { + kk_yield_extend(conts[i],ctx); // just move the continuation (no dup needed since it's parent is unique and being dropped) + } + kk_free((void*)self, ctx); + kk_box_drop(x,ctx); // still drop even though we yield as it may release a boxed value type? + return kk_box_any(ctx); // return yielding } - kk_function_drop(fself,ctx); - kk_box_drop(x,ctx); // still drop even though we yield as it may release a boxed value type? - return kk_box_any(ctx); // return yielding } + kk_free((void*)self, ctx); + // kk_function_drop(self,ctx); Can't do this, since all of it's child functions are dropped! + } else { + for(kk_intx_t i = 0; i < count; i++) { + // todo: take uniqueness of fself into account to avoid dup_function + kk_function_t f = kk_function_dup(conts[i],ctx); + x = kk_function_call(kk_box_t, (kk_function_t, kk_box_t, kk_context_t*), f, (f, x, ctx), ctx); + if (kk_yielding(ctx)) { + // if yielding, `yield_next` all continuations that still need to be done + while(++i < count) { + kk_yield_extend(kk_function_dup(conts[i],ctx),ctx); + } + kk_function_drop(fself,ctx); + kk_box_drop(x,ctx); // still drop even though we yield as it may release a boxed value type? + return kk_box_any(ctx); // return yielding + } + } + kk_function_drop(fself,ctx); } - kk_function_drop(fself,ctx); return x; } diff --git a/test/cgen/simple-refcount.kk b/test/cgen/simple-refcount.kk new file mode 100644 index 000000000..437033af3 --- /dev/null +++ b/test/cgen/simple-refcount.kk @@ -0,0 +1,22 @@ +import std/num/int32 + +extern create-box(): io-noexn any + c inline "kk_cptr_raw_box(&kk_free_fun, kk_malloc(sizeof(int), _ctx), _ctx)" + +extern get-refcount(^box: any): io-noexn int32 + c inline "kk_block_refcount(kk_box_to_ptr(box, _ctx))" + +effect ctl catch(): () + +fun main() + with handler ctl catch() + resume(()) + resume(()) + val box = create-box() + val count0 = box.get-refcount() + catch() + val count1 = box.get-refcount() + println((count0.int, count1.int).show) + // Expected output: + // (1,1) + // (1,0) \ No newline at end of file diff --git a/test/cgen/simple-refcount.kk.out b/test/cgen/simple-refcount.kk.out new file mode 100644 index 000000000..c124f7a23 --- /dev/null +++ b/test/cgen/simple-refcount.kk.out @@ -0,0 +1,2 @@ +(1,1) +(1,0) \ No newline at end of file