perf(sorted_set): O(n+m) difference, intersection, symmetric_difference

sorted_set/set.mbt

@@ -192,6 +192,54 @@ fn[V : Compare] split(root : Node[V]?, value : V) -> (Node[V]?, 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]?) {
+  match root {
+    None => (None, false, None)
+    Some(node) => {
+      let comp = value.compare(node.value)
+      if comp == 0 {
+        (node.left, true, node.right)
+      } else if comp < 0 {
+        let (l, found, r) = split_member(node.left, value)
+        (l, found, Some(join(r, node.value, node.right)))
+      } else {
+        let (l, found, r) = split_member(node.right, value)
+        (Some(join(node.left, node.value, l)), found, r)
+      }
+    }
+  }
+}
+
+///|
+/// Concatenates two trees where all elements in left < all elements in right.
+fn[V] concat(left : Node[V]?, right : Node[V]?) -> Node[V]? {
+  match (left, right) {
+    (None, _) => right
+    (_, None) => left
+    (Some(_), Some(_)) => {
+      let (min_val, right2) = remove_min(right.unwrap())
+      Some(join(left, min_val, right2))
+    }
+  }
+}
+
+///|
+/// Removes and returns the minimum value from a tree.
+fn[V] remove_min(node : Node[V]) -> (V, Node[V]?) {
+  match node.left {
+    None => (node.value, node.right)
+    Some(left) => {
+      let (min_val, new_left) = remove_min(left)
+      (min_val, Some(join(new_left, node.value, node.right)))
+    }
+  }
+}
+
 ///|
 fn[V] join(left : Node[V]?, value : V, right : Node[V]?) -> Node[V] {
   let (hl, hr) = (height(left), height(right))
@@ -261,9 +309,33 @@ pub fn[V : Compare] SortedSet::difference(
   self : SortedSet[V],
   src : SortedSet[V],
 ) -> SortedSet[V] {
-  let ret = new()
-  self.each(x => if !src.contains(x) { ret.add(x) })
-  ret
+  fn aux(a : Node[V]?, b : Node[V]?) -> Node[V]? {
+    match (a, b) {
+      (None, _) => None
+      (_, None) => a
+      (Some({ value: va, left: la, right: ra, .. }), _) => {
+        let (lb, found, rb) = split_member(b, va)
+        if found {
+          concat(aux(la, lb), aux(ra, rb))
+        } else {
+          Some(join(aux(la, lb), va, aux(ra, rb)))
+        }
+      }
+    }
+  }
+
+  let t = aux(copy_tree(self.root), copy_tree(src.root))
+  let mut ct = 0
+  fn count(node : Node[V]?) -> Unit {
+    if node is Some(n) {
+      ct += 1
+      count(n.left)
+      count(n.right)
+    }
+  }
+
+  count(t)
+  { root: t, size: ct }
 }
 
 ///|
@@ -293,10 +365,33 @@ pub fn[V : Compare] SortedSet::symmetric_difference(
   self : SortedSet[V],
   other : SortedSet[V],
 ) -> SortedSet[V] {
-  // TODO: Optimize this function to avoid creating two intermediate sets.
-  let set1 = self.difference(other)
-  let set2 = other.difference(self)
-  set1.union(set2)
+  fn aux(a : Node[V]?, b : Node[V]?) -> Node[V]? {
+    match (a, b) {
+      (None, _) => b
+      (_, None) => a
+      (Some({ value: va, left: la, right: ra, .. }), _) => {
+        let (lb, found, rb) = split_member(b, va)
+        if found {
+          concat(aux(la, lb), aux(ra, rb))
+        } else {
+          Some(join(aux(la, lb), va, aux(ra, rb)))
+        }
+      }
+    }
+  }
+
+  let t = aux(copy_tree(self.root), copy_tree(other.root))
+  let mut ct = 0
+  fn count(node : Node[V]?) -> Unit {
+    if node is Some(n) {
+      ct += 1
+      count(n.left)
+      count(n.right)
+    }
+  }
+
+  count(t)
+  { root: t, size: ct }
 }
 
 ///|
@@ -306,9 +401,32 @@ pub fn[V : Compare] SortedSet::intersection(
   self : SortedSet[V],
   src : SortedSet[V],
 ) -> SortedSet[V] {
-  let ret = new()
-  self.each(x => if src.contains(x) { ret.add(x) })
-  ret
+  fn aux(a : Node[V]?, b : Node[V]?) -> Node[V]? {
+    match (a, b) {
+      (None, _) | (_, None) => None
+      (Some({ value: va, left: la, right: ra, .. }), _) => {
+        let (lb, found, rb) = split_member(b, va)
+        if found {
+          Some(join(aux(la, lb), va, aux(ra, rb)))
+        } else {
+          concat(aux(la, lb), aux(ra, rb))
+        }
+      }
+    }
+  }
+
+  let t = aux(copy_tree(self.root), copy_tree(src.root))
+  let mut ct = 0
+  fn count(node : Node[V]?) -> Unit {
+    if node is Some(n) {
+      ct += 1
+      count(n.left)
+      count(n.right)
+    }
+  }
+
+  count(t)
+  { root: t, size: ct }
 }
 
 ///|