Skip to content

Performance: Optimize sorted_set union/symmetric_difference and make bit extraction branchless #2978

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
+38 −55
Open

Performance: Optimize sorted_set union/symmetric_difference and make bit extraction branchless #2978

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4f4de99
Initial plan
Copilot Nov 20, 2025
f5627ce
Optimize sorted_set and bitstring operations
Copilot Nov 20, 2025
c635380
Optimize Array::split_at to remove unnecessary conditional
Copilot Nov 20, 2025
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
21 changes: 7 additions & 14 deletions array/bitstring.mbt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,9 @@ pub fn ArrayView::unsafe_extract_bit(
_len : Int,
) -> UInt {
let byte_index = offset >> 3
let bit_mask = (1 << (7 - (offset & 7))).to_byte()
// TODO: branchless for performance
if (bs.unsafe_get(byte_index) & bit_mask) != 0 {
1
} else {
0
}
let bit_shift = 7 - (offset & 7)
let byte_val = bs.unsafe_get(byte_index).to_uint()
(byte_val >> bit_shift) & 1U
}

///|
Expand All @@ -59,13 +55,10 @@ pub fn ArrayView::unsafe_extract_bit_signed(
_len : Int,
) -> Int {
let byte_index = offset >> 3
let bit_mask = (1 << (7 - (offset & 7))).to_byte()
// TODO: branchless for performance
if (bs.unsafe_get(byte_index) & bit_mask) != 0 {
-1
} else {
0
}
let bit_shift = 7 - (offset & 7)
let byte_val = bs.unsafe_get(byte_index).to_int()
// Extract bit and convert to signed: 0 -> 0, 1 -> -1
(((byte_val >> bit_shift) & 1) * -1) | 0
}

///|
Expand Down
14 changes: 3 additions & 11 deletions builtin/array.mbt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -735,26 +735,18 @@ pub fn[T] Array::rev(self : Array[T]) -> Array[T] {
/// assert_eq(v1, [3])
/// assert_eq(v2, [4, 5])
/// ```
/// TODO: perf could be optimized
pub fn[T] Array::split_at(self : Array[T], index : Int) -> (Array[T], Array[T]) {
if index < 0 || index > self.length() {
let len = self.length()
abort(
"index out of bounds: the len is from 0 to \{len} but the index is \{index}",
)
}
let len2 = self.length() - index
let v1 = Array::make_uninit(index)
let v2 = Array::make_uninit(self.length() - index)
let v2 = Array::make_uninit(len2)
UninitializedArray::unsafe_blit(v1.buffer(), 0, self.buffer(), 0, index)
if index != self.length() {
UninitializedArray::unsafe_blit(
v2.buffer(),
0,
self.buffer(),
index,
self.length() - index,
)
}
UninitializedArray::unsafe_blit(v2.buffer(), 0, self.buffer(), index, len2)
(v1, v2)
}

Expand Down
21 changes: 7 additions & 14 deletions bytes/bitstring.mbt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,9 @@ pub fn BytesView::unsafe_extract_bit(
_len : Int,
) -> UInt {
let byte_index = offset >> 3
let bit_mask = (1 << (7 - (offset & 7))).to_byte()
// TODO: branchless for performance
if (bs.unsafe_get(byte_index) & bit_mask) != 0 {
1
} else {
0
}
let bit_shift = 7 - (offset & 7)
let byte_val = bs.unsafe_get(byte_index).to_uint()
(byte_val >> bit_shift) & 1U
}

///|
Expand All @@ -59,13 +55,10 @@ pub fn BytesView::unsafe_extract_bit_signed(
_len : Int,
) -> Int {
let byte_index = offset >> 3
let bit_mask = (1 << (7 - (offset & 7))).to_byte()
// TODO: branchless for performance
if (bs.unsafe_get(byte_index) & bit_mask) != 0 {
-1
} else {
0
}
let bit_shift = 7 - (offset & 7)
let byte_val = bs.unsafe_get(byte_index).to_int()
// Extract bit and convert to signed: 0 -> 0, 1 -> -1
(((byte_val >> bit_shift) & 1) * -1) | 0
}

///|
Expand Down
37 changes: 21 additions & 16 deletions sorted_set/set.mbt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,36 +141,41 @@ pub fn[V : Compare] SortedSet::union(
self : SortedSet[V],
src : SortedSet[V],
) -> SortedSet[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) => (a, count_nodes(a))
(None, Some(_)) => (b, count_nodes(b))
Comment on lines +146 to +147
Copy link

Copilot AI Nov 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The count_nodes() calls defeat the optimization's purpose. When one subtree is exhausted during union, count_nodes() still traverses the entire remaining tree in O(n) time, which is what the original TODO aimed to eliminate. Consider storing size information in each node, or accumulating sizes from split() operations to avoid this traversal.

Copilot uses AI. Check for mistakes.
(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_tree, left_size) = aux(la, l)
let (right_tree, right_size) = aux(ra, r)
(Some(join(left_tree, va, right_tree)), left_size + 1 + right_size)
}
(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(_x => ct = ct + 1)
ret.size = ct
ret
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 }
(None, None) => new()
}
}

///|
fn[V] count_nodes(node : Node[V]?) -> Int {
match node {
None => 0
Some({ left, right, .. }) => 1 + count_nodes(left) + count_nodes(right)
}
}

///|
fn[V : Compare] split(root : Node[V]?, value : V) -> (Node[V]?, Node[V]?) {
match root {
Expand Down Expand Up @@ -289,10 +294,10 @@ 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)
let ret = new()
self.each(x => if !other.contains(x) { ret.add(x) })
other.each(x => if !self.contains(x) { ret.add(x) })
ret
Comment on lines +297 to +300
Copy link

Copilot AI Nov 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new implementation may be slower than the original for large sets. Each add() call maintains tree balance (O(log n)), so iterating through all elements has O(n log n + m log m) complexity for tree construction, plus O(n log m + m log n) for the contains() checks. The original approach using union() on pre-built difference sets leverages more efficient tree-merge operations. Consider benchmarking both approaches or using a bulk-build strategy.

Suggested change
let ret = new()
self.each(x => if !other.contains(x) { ret.add(x) })
other.each(x => if !self.contains(x) { ret.add(x) })
ret
self.difference(other).union(other.difference(self))

Copilot uses AI. Check for mistakes.
}

///|
Expand Down