Skip to main content

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