Skip to content

Commit

Permalink
Rewrite binary search implementation
Browse files Browse the repository at this point in the history
This restores the original binary search implementation from #45333
which has the nice property of having a loop count that only depends on
the size of the slice. This, along with explicit conditional moves
from #128250, means that the entire binary search loop can be perfectly
predicted by the branch predictor.

Additionally, LLVM is able to unroll the loop when the slice length is
known at compile-time. This results in a very compact code sequence of
3-4 instructions per binary search step and zero branches.

Fixes #53823
  • Loading branch information
Amanieu committed Jul 26, 2024
1 parent 472058a commit d000d5f
Showing 1 changed file with 31 additions and 26 deletions.
57 changes: 31 additions & 26 deletions library/core/src/slice/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2787,40 +2787,45 @@ impl<T> [T] {
where
F: FnMut(&'a T) -> Ordering,
{
// INVARIANTS:
// - 0 <= left <= left + size = right <= self.len()
// - f returns Less for everything in self[..left]
// - f returns Greater for everything in self[right..]
let mut size = self.len();
let mut left = 0;
let mut right = size;
while left < right {
let mid = left + size / 2;

// SAFETY: the while condition means `size` is strictly positive, so
// `size/2 < size`. Thus `left + size/2 < left + size`, which
// coupled with the `left + size <= self.len()` invariant means
// we have `left + size/2 < self.len()`, and this is in-bounds.
if size == 0 {
return Err(0);
}
let mut base = 0usize;

// This loop intentionally doesn't have an early exit if the comparison
// returns Equal. We want the number of loop iterations to depend *only*
// on the size of the input slice so that the CPU can reliably predict
// the loop count.
while size > 1 {
let half = size / 2;
let mid = base + half;

// SAFETY: the call is made safe by the following inconstants:
// - `mid >= 0`: by definition
// - `mid < size`: `mid = size / 2 + size / 4 + size / 8 ...`
let cmp = f(unsafe { self.get_unchecked(mid) });

// Binary search interacts poorly with branch prediction, so force
// the compiler to use conditional moves if supported by the target
// architecture.
left = select_unpredictable(cmp == Less, mid + 1, left);
right = select_unpredictable(cmp == Greater, mid, right);
if cmp == Equal {
// SAFETY: same as the `get_unchecked` above
unsafe { hint::assert_unchecked(mid < self.len()) };
return Ok(mid);
}

size = right - left;
base = select_unpredictable(cmp == Greater, base, mid);
size -= half;
}

// 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)
// SAFETY: base is always in [0, size) because base <= mid.
let cmp = f(unsafe { self.get_unchecked(base) });
if cmp == Equal {
// SAFETY: same as the `get_unchecked` above.
unsafe { hint::assert_unchecked(base < self.len()) };
Ok(base)
} else {
let result = base + (cmp == Less) as usize;
// SAFETY: same as the `get_unchecked` above.
// Note that this is `<=`, unlike the assume in the `Ok` path.
unsafe { hint::assert_unchecked(result <= self.len()) };
Err(result)
}
}

/// Binary searches this slice with a key extraction function.
Expand Down

0 comments on commit d000d5f

Please sign in to comment.