perf(sorted_set): compute union size from split drops instead of recounting

Resolves the TODO in union: after merging, the result size was
recomputed by a full each() traversal of the merged tree. An element
common to both sets is dropped exactly once, at the found pivot of the
split, so union now counts those drops and derives the size
arithmetically: size = |self| + |src| - dups.

Rebased reimplementation: when this PR was first written, union's
`split` was the only splitter and gained an Int drop-count component.
main has since grown `split_member` (the split-based difference/
intersection/symmetric_difference rewrite), whose found flag carries
exactly the same information - so union now uses `split_member`, and
the now-unused `split` is deleted along with its whitebox test, which
is replaced by an equivalent `split_member` test pinning found for
present, absent, below-range, above-range and empty splits.

`split_member`'s result also becomes a `#valtype` struct
(`SplitResult`: left/found/right) instead of a 3-tuple, keeping the
per-node result stack-allocated on native (same pattern as
`JsonNumberScan` in json/lex_number.mbt). This benefits every
split-based operation, not just union.

Benchmarks (moon bench --release, n=10K sets, means of 10x batches):

  union, main -> this branch (recount removed + valtype):
    50% overlap  native 518 -> 429us (-17%), js 297 -> 222us (-25%),
                 wasm-gc 197 -> 158us (-20%)
    disjoint     native 382 -> 304us (-20%), js 261 -> 153us (-41%),
                 wasm-gc 154 -> 96us (-38%)
    identical    native 592 -> 520us (-12%), js 323 -> 277us (-14%),
                 wasm-gc 232 -> 193us (-17%)
  valtype struct vs tuple (isolated, all split-based ops):
    native -5% to -9% across union/difference/intersection;
    js flat; wasm-gc within noise (+/-2%).

Tests: union exact-size assertions for overlapping, disjoint and
identical operands cross-checked against an element count; the
split_member whitebox test above; the existing invariant and
quickcheck suites. Mutation-verified: double-counting dups fails 8
tests including the exact-size test and the whitebox invariant
scenarios.

Partially addresses #3824 (the defensive copy_tree of both operands is
inherent to the mutable node design and stays).

Codex CLI review (xhigh, first-round sign-off on the rebased
reimplementation): the dups induction re-verified on split_member
(disjoint fragments, pivot counted exactly once, empty-side arms cannot
hide common elements), sizes read from unmutated copies, SplitResult
preserves tuple semantics at all four call sites, no split callers
remain, interface byte-identical.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Codex CLI <codex@openai.com>

Author: bobzhang

sorted_set/set.mbt

@@ -177,12 +177,19 @@ pub fn[V : Compare] SortedSet::union(
   self : SortedSet[V],
   src : SortedSet[V],
 ) -> SortedSet[V] {
+  // An element common to both sets is dropped exactly once, at the found
+  // pivot of `split_member`; counting those drops gives the result size
+  // arithmetically instead of re-traversing the merged tree.
+  let mut dups = 0
   fn aux(a : Node[V]?, b : Node[V]?) -> Node[V]? {
     match (a, b) {
       (Some(_), None) => a
       (None, Some(_)) => b
       (Some({ value: va, left: la, right: ra, .. }), Some(_)) => {
-        let (l, r) = split(b, va)
+        let { left: l, found, right: r } = split_member(b, va)
+        if found {
+          dups += 1
+        }
         Some(join(aux(la, l), va, aux(ra, r)))
       }
       (None, None) => None
@@ -191,15 +198,8 @@ pub fn[V : Compare] SortedSet::union(
 
   match (self.root, src.root) {
     (Some(_), Some(_)) => {
-      let t1 = copy_tree(self.root)
-      let t2 = copy_tree(src.root)
-      let t = aux(t1, t2)
-      let mut ct = 0
-      let ret = { root: t, size: 0 }
-      // TODO: optimize this. Avoid counting the size of the set.
-      ret.each(_x => ct += 1)
-      ret.size = ct
-      ret
+      let t = aux(copy_tree(self.root), copy_tree(src.root))
+      { root: t, size: self.size + src.size - dups }
     }
     (Some(_), None) => { root: copy_tree(self.root), size: self.size }
     (None, Some(_)) => { root: copy_tree(src.root), size: src.size }
@@ -208,42 +208,32 @@ pub fn[V : Compare] SortedSet::union(
 }
 
 ///|
-fn[V : Compare] split(root : Node[V]?, value : V) -> (Node[V]?, Node[V]?) {
-  match root {
-    None => (None, None)
-    Some(node) => {
-      let comp = value.compare(node.value)
-      if comp == 0 {
-        (node.left, node.right)
-      } else if comp < 0 {
-        let (l, r) = split(node.left, value)
-        (l, Some(join(r, node.value, node.right)))
-      } else {
-        let (l, r) = split(node.right, value)
-        (Some(join(node.left, node.value, l)), r)
-      }
-    }
-  }
+// `#valtype` keeps the split result stack-allocated on the native target:
+// `split_member` builds and returns one per visited node on every
+// set-operation pivot path.
+#valtype
+priv struct SplitResult[V] {
+  left : Node[V]?
+  found : Bool
+  right : Node[V]?
 }
 
 ///|
-/// Splits a tree by a value, returning (left, found, right).
-fn[V : Compare] split_member(
-  root : Node[V]?,
-  value : V,
-) -> (Node[V]?, Bool, Node[V]?) {
+/// Splits a tree by a value into the elements less than it, a flag for
+/// whether it was present, and the elements greater than it.
+fn[V : Compare] split_member(root : Node[V]?, value : V) -> SplitResult[V] {
   match root {
-    None => (None, false, None)
+    None => { left: None, found: false, right: None }
     Some(node) => {
       let comp = value.compare(node.value)
       if comp == 0 {
-        (node.left, true, node.right)
+        { left: node.left, found: true, right: node.right }
       } else if comp < 0 {
-        let (l, found, r) = split_member(node.left, value)
-        (l, found, Some(join(r, node.value, node.right)))
+        let { left, found, right } = split_member(node.left, value)
+        { left, found, right: Some(join(right, node.value, node.right)) }
       } else {
-        let (l, found, r) = split_member(node.right, value)
-        (Some(join(node.left, node.value, l)), found, r)
+        let { left, found, right } = split_member(node.right, value)
+        { left: Some(join(node.left, node.value, left)), found, right }
       }
     }
   }
@@ -369,7 +359,7 @@ pub fn[V : Compare] SortedSet::difference(
           (None, _) => None
           (_, None) => a
           (Some({ value: va, left: la, right: ra, .. }), _) => {
-            let (lb, found, rb) = split_member(b, va)
+            let { left: lb, found, right: rb } = split_member(b, va)
             if found {
               found_count += 1
               concat(aux(la, lb), aux(ra, rb))
@@ -425,7 +415,7 @@ pub fn[V : Compare] SortedSet::symmetric_difference(
       (None, _) => b
       (_, None) => a
       (Some({ value: va, left: la, right: ra, .. }), _) => {
-        let (lb, found, rb) = split_member(b, va)
+        let { left: lb, found, right: rb } = split_member(b, va)
         if found {
           found_count += 1
           concat(aux(la, lb), aux(ra, rb))
@@ -475,7 +465,7 @@ pub fn[V : Compare] SortedSet::intersection(
         match (a, b) {
           (None, _) | (_, None) => None
           (Some({ value: va, left: la, right: ra, .. }), _) => {
-            let (lb, found, rb) = split_member(b, va)
+            let { left: lb, found, right: rb } = split_member(b, va)
             if found {
               found_count += 1
               Some(join(aux(la, lb), va, aux(ra, rb)))
@@ -1038,41 +1028,39 @@ test "union" {
 
 ///|
 #warnings("-deprecated")
-test "split" {
-  let (l, r) = split(from_array([7, 2, 9, 4, 5, 6, 3, 8, 1]).root, 5)
+test "split_member" {
+  let { left: l, found, right: r } = split_member(
+    from_array([7, 2, 9, 4, 5, 6, 3, 8, 1]).root,
+    5,
+  )
+  inspect(found, content="true")
   inspect(l, content="Some([1, 2, 3, 4])")
   inspect(r, content="Some([6, 7, 8, 9])")
-  let (l, r) = split(from_array([7, 2, 9, 4, 5, 6, 3, 8, 1]).root, 0)
+  let { left: l, found, right: r } = split_member(
+    from_array([7, 2, 9, 4, 5, 6, 3, 8, 1]).root,
+    0,
+  )
+  inspect(found, content="false")
   inspect(l, content="None")
   inspect(r, content="Some([1, 2, 3, 4, 5, 6, 7, 8, 9])")
-  let (l, r) = split(from_array([7, 2, 9, 4, 5, 6, 3, 8, 1]).root, 10)
+  let { left: l, found, right: r } = split_member(
+    from_array([7, 2, 9, 4, 5, 6, 3, 8, 1]).root,
+    10,
+  )
+  inspect(found, content="false")
   inspect(l, content="Some([1, 2, 3, 4, 5, 6, 7, 8, 9])")
   inspect(r, content="None")
-  let (l, r) = split(from_array([7, 2, 9, 4, 5, 6, 3, 8, 1]).root, 4)
+  let { left: l, found, right: r } = split_member(
+    from_array([7, 2, 9, 4, 5, 6, 3, 8, 1]).root,
+    4,
+  )
+  inspect(found, content="true")
   inspect(l, content="Some([1, 2, 3])")
   inspect(r, content="Some([5, 6, 7, 8, 9])")
-  let (l, r) = split(from_array([]).root, 7)
+  let { left: l, found, right: r } = split_member(from_array([]).root, 7)
+  inspect(found, content="false")
   inspect(l, content="None")
   inspect(r, content="None")
-  let (l, r) = split(
-    from_array([
-      0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21,
-      22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
-      41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
-      60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78,
-      79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97,
-      98, 99, 100,
-    ]).root,
-    50,
-  )
-  inspect(
-    l,
-    content="Some([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49])",
-  )
-  inspect(
-    r,
-    content="Some([51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100])",
-  )
 }
 
 ///|

sorted_set/set_test.mbt

@@ -564,3 +564,18 @@ test "intersection and difference preserve self's stored representatives" {
   let out3 = lone.difference(far).to_array()
   assert_true(out3.length() == 1 && out3[0].payload == 1)
 }
+
+///|
+test "union size is exact for overlapping, disjoint and identical sets" {
+  let a = @sorted_set.from_array([1, 2, 3, 4, 5])
+  let b = @sorted_set.from_array([4, 5, 6, 7])
+  inspect(a.union(b).length(), content="7")
+  inspect(b.union(a).length(), content="7")
+  inspect(a.union(@sorted_set.from_array([10, 11])).length(), content="7")
+  inspect(a.union(a).length(), content="5")
+  // size must agree with an actual element count
+  let u = a.union(b)
+  let mut n = 0
+  u.each(_x => n = n + 1)
+  inspect(u.length() == n, content="true")
+}