1 #[cfg(test)]
2 mod tests;
3
4 /// Uses a sorted slice `data: &[E]` as a kind of "multi-map". The
5 /// `key_fn` extracts a key of type `K` from the data, and this
6 /// function finds the range of elements that match the key. `data`
7 /// must have been sorted as if by a call to `sort_by_key` for this to
8 /// work.
binary_search_slice<'d, E, K>(data: &'d [E], key_fn: impl Fn(&E) -> K, key: &K) -> &'d [E] where K: Ord,9 pub fn binary_search_slice<'d, E, K>(data: &'d [E], key_fn: impl Fn(&E) -> K, key: &K) -> &'d [E]
10 where
11 K: Ord,
12 {
13 let Ok(mid) = data.binary_search_by_key(key, &key_fn) else {
14 return &[];
15 };
16 let size = data.len();
17
18 // We get back *some* element with the given key -- so do
19 // a galloping search backwards to find the *first* one.
20 let mut start = mid;
21 let mut previous = mid;
22 let mut step = 1;
23 loop {
24 start = start.saturating_sub(step);
25 if start == 0 || key_fn(&data[start]) != *key {
26 break;
27 }
28 previous = start;
29 step *= 2;
30 }
31 step = previous - start;
32 while step > 1 {
33 let half = step / 2;
34 let mid = start + half;
35 if key_fn(&data[mid]) != *key {
36 start = mid;
37 }
38 step -= half;
39 }
40 // adjust by one if we have overshot
41 if start < size && key_fn(&data[start]) != *key {
42 start += 1;
43 }
44
45 // Now search forward to find the *last* one.
46 let mut end = mid;
47 let mut previous = mid;
48 let mut step = 1;
49 loop {
50 end = end.saturating_add(step).min(size);
51 if end == size || key_fn(&data[end]) != *key {
52 break;
53 }
54 previous = end;
55 step *= 2;
56 }
57 step = end - previous;
58 while step > 1 {
59 let half = step / 2;
60 let mid = end - half;
61 if key_fn(&data[mid]) != *key {
62 end = mid;
63 }
64 step -= half;
65 }
66
67 &data[start..end]
68 }
69