Improve refcount for unique continuations

lib/std/core/hnd.kk

@@ -603,31 +603,37 @@ fun get( ref: ref<h,a>) : <read<h>,div> a
 inline extern unsafe-st(f : () -> <st<global>|e> a ) : (() -> e a)
   inline "#1"
 
-fun protect-prompt( resumed : ref<global,bool>,  k : resume-result<b,r> -> 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<b,r,e>
+  NeedsFinalization(k : resume-result<b,r> -> e r)
+  NoFinalization
+
+fun protect-prompt( resumed : ref<global,protect-state<b,r,e>>, res : r) : e r
+  val hdiv=@Hnodiv // Needed since resumed is in global state.
+  val did-resume : some<b,r,e> protect-state<b,r,e> = (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<b,r> -> 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<a,b,e,r>
@@ -715,12 +721,12 @@ fun under2( ev : ev<h>, 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<b,r> -> 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<a1,a2,b,h,e,r>
   Clause2(fn(m,_ev,x1,x2){ yield-to(m, fn(k){ protect2(x1,x2,clause,k) }) })

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;
 }
 

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) 

test/cgen/simple-refcount.kk.out

@@ -0,0 +1,2 @@
+(1,1)
+(1,0)