musli_zerocopy/slice/
binary_search.rs

1use core::cmp::Ordering;
2
3use crate::buf::{Buf, Visit};
4use crate::error::Error;
5use crate::slice::Slice;
6use crate::traits::ZeroCopy;
7
8/// The result of a [`binary_search()`].
9#[derive(Debug, PartialEq, Eq)]
10pub enum BinarySearch {
11    /// Found the element we were looking for during the search.
12    Found(usize),
13    /// Exact match could not be found, but this is the closest index where the
14    /// searched for element could be inserted while still maintaining search
15    /// order.
16    Missing(usize),
17}
18
19/// Binary searches this slice for a given element. If the slice is not sorted,
20/// the returned result is unspecified and meaningless.
21///
22/// If the value is found then [`BinarySearch::Found`] is returned, containing
23/// the index of the matching element. If there are multiple matches, then any
24/// one of the matches could be returned. The index is chosen deterministically,
25/// but is subject to change in future versions of Rust. If the value is not
26/// found then [`BinarySearch::Missing`] is returned, containing the index
27/// where a matching element could be inserted while maintaining sorted order.
28///
29/// # Examples
30///
31/// Looks up a series of four elements. The first is found, with a uniquely
32/// determined position; the second and third are not found; the fourth could
33/// match any position in `[1, 4]`.
34///
35/// ```
36/// use musli_zerocopy::OwnedBuf;
37/// use musli_zerocopy::slice::{binary_search, BinarySearch};
38///
39/// let mut buf = OwnedBuf::new();
40/// let slice = buf.store_slice(&[0, 1, 1, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55]);
41///
42/// assert_eq!(binary_search(&buf, slice, &13)?, BinarySearch::Found(9));
43/// assert_eq!(binary_search(&buf, slice, &4)?, BinarySearch::Missing(7));
44/// assert_eq!(binary_search(&buf, slice, &100)?, BinarySearch::Missing(13));
45///
46/// let r = binary_search(&buf, slice, &1)?;
47/// assert!(match r { BinarySearch::Found(1..=4) => true, _ => false, });
48/// # Ok::<_, musli_zerocopy::Error>(())
49/// ```
50pub fn binary_search<S, Q>(buf: &Buf, slice: S, x: &Q) -> Result<BinarySearch, Error>
51where
52    S: Slice<Item: ZeroCopy + Ord>,
53    Q: Visit<Target = S::Item>,
54{
55    binary_search_by(buf, slice, |value| x.visit(buf, |x| value.cmp(x)))
56}
57
58/// Binary searches this slice with a comparator function.
59///
60/// The comparator function should return an order code that indicates whether
61/// its argument is `Less`, `Equal` or `Greater` the desired target. If the
62/// slice is not sorted or if the comparator function does not implement an
63/// order consistent with the sort order of the underlying slice, the returned
64/// result is unspecified and meaningless.
65///
66/// If the value is found then [`BinarySearch::Found`] is returned, containing
67/// the index of the matching element. If there are multiple matches, then any
68/// one of the matches could be returned. The index is chosen deterministically,
69/// but is subject to change in future versions of Rust. If the value is not
70/// found then [`BinarySearch::Missing`] is returned, containing the index where
71/// a matching element could be inserted while maintaining sorted order.
72///
73/// # Examples
74///
75/// Looks up a series of four elements. The first is found, with a
76/// uniquely determined position; the second and third are not
77/// found; the fourth could match any position in `[1, 4]`.
78///
79/// ```
80/// use musli_zerocopy::OwnedBuf;
81/// use musli_zerocopy::slice::{binary_search_by, BinarySearch};
82///
83/// let mut buf = OwnedBuf::new();
84/// let slice = buf.store_slice(&[0, 1, 1, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55]);
85///
86/// let s = 13;
87/// assert_eq!(binary_search_by(&buf, slice, |v| Ok(v.cmp(&s)))?, BinarySearch::Found(9));
88/// let s = 4;
89/// assert_eq!(binary_search_by(&buf, slice, |v| Ok(v.cmp(&s)))?, BinarySearch::Missing(7));
90/// let s = 100;
91/// assert_eq!(binary_search_by(&buf, slice, |v| Ok(v.cmp(&s)))?, BinarySearch::Missing(13));
92/// let s = 1;
93/// let r = binary_search_by(&buf, slice, |v| Ok(v.cmp(&s)))?;
94/// assert!(match r { BinarySearch::Found(1..=4) => true, _ => false, });
95/// # Ok::<_, musli_zerocopy::Error>(())
96/// ```
97pub fn binary_search_by<S, F>(buf: &Buf, slice: S, mut f: F) -> Result<BinarySearch, Error>
98where
99    S: Slice<Item: ZeroCopy>,
100    F: FnMut(&S::Item) -> Result<Ordering, Error>,
101{
102    // INVARIANTS:
103    // - 0 <= left <= left + size = right <= slice.len()
104    // - f returns Less for everything in slice[..left]
105    // - f returns Greater for everything in slice[right..]
106    let mut size = slice.len();
107    let mut left = 0;
108    let mut right = size;
109
110    while left < right {
111        let mid = left + size / 2;
112
113        // The while condition means `size` is strictly positive, so `size/2
114        // < size`. Thus `left + size/2 < left + size`, which coupled with
115        // the `left + size <= slice.len()` invariant means we have `left +
116        // size/2 < slice.len()`, and this is in-bounds.
117        let value = buf.load(slice.get_unchecked(mid))?;
118        let cmp = f(value)?;
119
120        // The reason why we use if/else control flow rather than match
121        // is because match reorders comparison operations, which is perf sensitive.
122        // This is x86 asm for u8: https://rust.godbolt.org/z/8Y8Pra.
123        if cmp == Ordering::Less {
124            left = mid + 1;
125        } else if cmp == Ordering::Greater {
126            right = mid;
127        } else {
128            return Ok(BinarySearch::Found(mid));
129        }
130
131        size = right - left;
132    }
133
134    Ok(BinarySearch::Missing(left))
135}