Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bobzhang merged 1 commit into from
Aug 20, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
118 changes: 53 additions & 65 deletions sorted_set/set.mbt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 }
Expand All @@ -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 }
}
}
}
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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)))
Expand Down Expand Up @@ -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])",
)
}

///|
Expand Down
15 changes: 15 additions & 0 deletions sorted_set/set_test.mbt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
}
Loading