Skip to content

Commit

Permalink
Improve slice::binary_search_by
Browse files Browse the repository at this point in the history
  • Loading branch information
krtab committed Jun 26, 2024
1 parent 4bc39f0 commit e4a10ae
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 13 deletions.
26 changes: 14 additions & 12 deletions library/core/src/slice/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2788,13 +2788,12 @@ impl<T> [T] {
F: FnMut(&'a T) -> Ordering,
{
// INVARIANTS:
// - 0 <= left <= left + size = right <= self.len()
// - 0 <= left <= left + size <= self.len()
// - f returns Less for everything in self[..left]
// - f returns Greater for everything in self[right..]
// - f returns Greater for everything in self[left + size..]
let mut size = self.len();
let mut left = 0;
let mut right = size;
while left < right {
while size > 1 {
let mid = left + size / 2;

// SAFETY: the while condition means `size` is strictly positive, so
Expand All @@ -2807,21 +2806,24 @@ impl<T> [T] {
// fewer branches and instructions than if/else or matching on
// cmp::Ordering.
// This is x86 asm for u8: https://rust.godbolt.org/z/698eYffTx.
left = if cmp == Less { mid + 1 } else { left };
right = if cmp == Greater { mid } else { right };

left = if cmp == Less { mid } else { left };
size = if cmp == Greater { size / 2 } else { size - size / 2 };
if cmp == Equal {
// SAFETY: same as the `get_unchecked` above
unsafe { hint::assert_unchecked(mid < self.len()) };
return Ok(mid);
}

size = right - left;
}

// SAFETY: directly true from the overall invariant.
// Note that this is `<=`, unlike the assume in the `Ok` path.
unsafe { hint::assert_unchecked(left <= self.len()) };
Err(left)
if size == 0 {
Err(left)
} else {
// SAFETY: allowed per the invariants
let cmp = f(unsafe { self.get_unchecked(left) });
let res_idx = if cmp == Less { left + 1 } else { left };
if cmp == Equal { Ok(res_idx) } else { Err(res_idx) }
}
}

/// Binary searches this slice with a key extraction function.
Expand Down
2 changes: 1 addition & 1 deletion library/core/tests/slice.rs
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ fn test_binary_search_implementation_details() {
assert_eq!(b.binary_search(&3), Ok(5));
let b = [1, 1, 1, 1, 1, 3, 3, 3, 3];
assert_eq!(b.binary_search(&1), Ok(4));
assert_eq!(b.binary_search(&3), Ok(7));
assert_eq!(b.binary_search(&3), Ok(6));
let b = [1, 1, 1, 1, 3, 3, 3, 3, 3];
assert_eq!(b.binary_search(&1), Ok(2));
assert_eq!(b.binary_search(&3), Ok(4));
Expand Down

0 comments on commit e4a10ae

Please sign in to comment.