fix(compiler): hand last fork readers counted ownership together

bend2/comp.ts

@@ -2088,7 +2088,7 @@ function emit_jump(sc: Scope, args: string[], k: Name, bang = false): void {
   file_push(sc, `WL_AGAIN(${sc.seg.fid});`);
 }
 
-function emit_args(sc: Scope, ck: Spine, jump = false, fork = false): string[] {
+function emit_args(sc: Scope, ck: Spine, jump = false): string[] {
   const k = ck.k!;
   const brw = brw_of(k);
   ck.all.forEach((a, q) => {
@@ -2118,7 +2118,7 @@ function emit_args(sc: Scope, ck: Spine, jump = false, fork = false): string[] {
       }
       if (dead && !twin && val_brw(sc, bd.val)) {
         sc.uses.delete(p);
-      } else if (!fork) {
+      } else {
         sc.uses.set(p, { ...bd, n: Math.max(bd.n - 1, 1) });
       }
       b = bd.val;
@@ -2543,12 +2543,18 @@ function emit_fork(sc: Scope, x: Of<"Let">, ers: HTerm[]): void {
   const name = seg_name(sc, "j");
   let hold: Of<"Var">[] = [];
   if (fork) {
+    for (const c of calls) c.xs.forEach((a, i) => {
+      const x = term_strip(a);
+      if (x.$ !== "Var") return;
+      const p = probe_of(x), b = sc.uses.get(p);
+      if (b && rest_use([o.b], p) === 0) val_own(sc, b.val, c.k! + "~" + i);
+    });
     spare_flush(sc);
     sc.seg.fork = true;
     const pl = { ...sc, uses: new Map(sc.uses) };
     block(pl, "if (!seq) {", () => {
       const margs = calls.map((c, j) => emit_args({ ...pl,
-        rest: [...x.v.filter((_, i) => i !== j), o.b] }, c, false, true));
+        rest: [...x.v.slice(j + 1), o.b] }, c));
       const live = [...pl.uses].filter(([p, b]) =>
         !val_brw(pl, b.val) || rest_use([o.b], p) > 0);
       hold = live.map(([p]) => p);

guide/SHADERS.md

@@ -192,6 +192,12 @@ borrows: `Blk.go` only matches, `Tile.fin` consumes last. So the scene is
 host-owned, borrowed by the bang, and returned (`Frame.show`). Read through
 counts, the frame costs 8.0 -> 16.5 ms (draw 6.0 -> 14.1).
 
+When a fork is the last use of a locally owned box, its readers take counted
+ownership together, so recursive readers free nodes as they walk rather than
+leaving one whole-value sink at the join. All sibling readers must own a
+reference before that handoff; giving ownership to only one could invalidate
+another's borrow. A value needed after the fork retains its original owner.
+
 - Two consumers of one list share it: 45 ms frames and a 448 s Metal compile. A
   `+` where one use would do heats its type (symreg 2.83 -> 3.30 s).
 - A def that returns its argument shares it: `Bool.pick(Cloth, c, a, b)` on a

tests/reg/borrow_fork_hold.bend

@@ -1,9 +1,6 @@
-# regression witness: two forked readers of one list whose join does not
-# mention it. The join now holds the list (a captured variable that dies
-# at the fork) and sinks it after both tasks, so the readers borrow it.
-# Before, each reader owned its copy: a term_keep before the fork, the
-# type hot, and each match a ctr_take. In the C, WL_CASE(FID_TWICE) has
-# no term_keep and WL_CASE(FID_XS_LEN) reads through term_peek.
+# Two different last fork readers of one list must both observe every
+# element while safely releasing their references. The join has no later
+# use of the list; neither reader may free beneath its sibling.
 import Base
 
 type Xs is Data:

tests/reg/brw_fork_cross.bend

@@ -1,8 +1,6 @@
-# borrow witness: a fork whose kids read (cnt~0 borrowed) and consume
-# (tail~0 owned: it returns a field) one list. The reader's use still
-# counts when the consumer pops, so tail gets a kept copy (a term_keep on
-# each path, parallel and sequential; Xs hot), both tasks see a live list,
-# and the join holds it and sinks it after both. Correct at 8 threads.
+# A forked reader and a tail consumer must share a live list safely.
+# Returning a field from one task must not invalidate the other task's
+# traversal; the returned tail is still readable after the join.
 import Base
 
 type Xs is Data:

tests/reg/brw_fork_dead.bend

@@ -1,8 +1,6 @@
-# borrow witness: a fork with a dead parameter. Xs.seven ignores its list,
-# so seven~0 is borrowed (an unused rooted word is let go, never sunk);
-# cnt~0 is borrowed (its sibling holds the list). In the C, FID_TWICE and
-# FID_THRICE have no term_keep, the join holds xs and the joiner sinks it
-# once, and cnt reads through term_peek.
+# A fork with an unused list argument must release that task's reference
+# without invalidating its sibling. Thrice also retains the list for a
+# later traversal, so that original owner must survive the join.
 import Base
 
 type Xs is Data:

tests/reg/fork_last_readers_tree.bend

@@ -0,0 +1,54 @@
+# Different fork readers must retain independent ownership, while a returned
+# owner remains live through the join. Each reader recursively forks its tree.
+import Base
+
+type Tree is Data:
+  Leaf{v: U32}
+  Node{l: Tree, r: Tree}
+
+def gen(+d: Nat) -> Tree:
+  match d:
+    case 0n:
+      Leaf{5}
+    case 1n+p:
+      a b = gen(p) gen(p)
+      Node{a, b}
+
+def sum(t: Tree) -> U32:
+  match t:
+    case Leaf{v}:
+      v
+    case Node{l, r}:
+      a b = sum(l) sum(r)
+      (a + b : U32)
+
+def count(t: Tree) -> U32:
+  match t:
+    case Leaf{v}:
+      1
+    case Node{l, r}:
+      a b = count(l) count(r)
+      (a + b : U32)
+
+def consume(+t: Tree) -> U32:
+  a b = sum(t) count(t)
+  (a + b : U32)
+
+def retain(+t: Tree) -> U32 & Tree:
+  a b = sum(t) count(t)
+  ((a + b : U32), t)
+
+def retained(p: U32 & Tree) -> U32:
+  (value, tree) = p
+  (value + sum(tree) : U32)
+
+def main() -> IO(Unit):
+  kept = retained(retain(gen(10n)))
+  do IO<Unit>:
+    Unit <- IO.print(U32.show(consume(gen(10n))))
+    Unit <- IO.print(U32.show(kept))
+    IO.print(U32.show(consume!(gen(10n))))
+
+#|6144
+#|11264
+#|6144