From 447d63a6022c759542f68c245d90a3f9967ccd1f Mon Sep 17 00:00:00 2001
From: Hongbo Zhang <bobzhang1988@gmail.com>
Date: Fri, 13 Jun 2025 16:46:00 +0800
Subject: [PATCH] Optimize sorted_set union size calculation

---
 sorted_set/set.mbt | 40 ++++++++++++++++++++++++++--------------
 1 file changed, 26 insertions(+), 14 deletions(-)

diff --git a/sorted_set/set.mbt b/sorted_set/set.mbt
index 42b252866..8f71cdf6d 100644
--- a/sorted_set/set.mbt
+++ b/sorted_set/set.mbt
@@ -81,6 +81,20 @@ fn[V] copy_tree(node : Node[V]?) -> Node[V]? {
   }
 }
 
+///|
+/// Copies a tree and returns both the copied tree and its size.
+fn[V] copy_tree_with_size(node : Node[V]?) -> (Node[V]?, Int) {
+  match node {
+    None => (None, 0)
+    Some(node) => {
+      let (left, lsize) = copy_tree_with_size(node.left)
+      let (right, rsize) = copy_tree_with_size(node.right)
+      let new_node = new_node(node.value, left~, right~, height=node.height)
+      (Some(new_node), lsize + rsize + 1)
+    }
+  }
+}
+
 ///|
 fn[V] new_node(
   value : V,
@@ -152,29 +166,27 @@ pub fn[V : Compare] contains(self : T[V], value : V) -> Bool {
 ///|
 /// Returns the union of two sets.
 pub fn[V : Compare] union(self : T[V], src : T[V]) -> T[V] {
-  fn aux(a : Node[V]?, b : Node[V]?) -> Node[V]? {
+  fn aux(a : Node[V]?, b : Node[V]?) -> (Node[V]?, Int) {
     match (a, b) {
-      (Some(_), None) => a
-      (None, Some(_)) => b
+      (Some(_), None) => copy_tree_with_size(a)
+      (None, Some(_)) => copy_tree_with_size(b)
       (Some({ value: va, left: la, right: ra, .. }), Some(_)) => {
         let (l, r) = split(b, va)
-        Some(join(aux(la, l), va, aux(ra, r)))
+        let (left, lsize) = aux(la, l)
+        let (right, rsize) = aux(ra, r)
+        let node = join(left, va, right)
+        (Some(node), lsize + rsize + 1)
       }
-      (None, None) => None
+      (None, None) => (None, 0)
     }
   }
 
   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(fn(_x) { ct = ct + 1 })
-      ret.size = ct
-      ret
+      let (t1, _) = copy_tree_with_size(self.root)
+      let (t2, _) = copy_tree_with_size(src.root)
+      let (t, size) = aux(t1, t2)
+      { root: t, size }
     }
     (Some(_), None) => { root: copy_tree(self.root), size: self.size }
     (None, Some(_)) => { root: copy_tree(src.root), size: src.size }