diff options
| author | mark <markm@cs.wisc.edu> | 2020-08-27 22:58:48 -0500 |
|---|---|---|
| committer | Vadim Petrochenkov <vadim.petrochenkov@gmail.com> | 2020-08-30 18:45:07 +0300 |
| commit | 9e5f7d5631b8f4009ac1c693e585d4b7108d4275 (patch) | |
| tree | 158a05eb3f204a8e72939b58427d0c2787a4eade /compiler/rustc_data_structures/src/binary_search_util | |
| parent | db534b3ac286cf45688c3bbae6aa6e77439e52d2 (diff) | |
| download | rust-9e5f7d5631b8f4009ac1c693e585d4b7108d4275.tar.gz rust-9e5f7d5631b8f4009ac1c693e585d4b7108d4275.zip | |
mv compiler to compiler/
Diffstat (limited to 'compiler/rustc_data_structures/src/binary_search_util')
| -rw-r--r-- | compiler/rustc_data_structures/src/binary_search_util/mod.rs | 69 | ||||
| -rw-r--r-- | compiler/rustc_data_structures/src/binary_search_util/tests.rs | 23 |
2 files changed, 92 insertions, 0 deletions
diff --git a/compiler/rustc_data_structures/src/binary_search_util/mod.rs b/compiler/rustc_data_structures/src/binary_search_util/mod.rs new file mode 100644 index 00000000000..ede5757a479 --- /dev/null +++ b/compiler/rustc_data_structures/src/binary_search_util/mod.rs @@ -0,0 +1,69 @@ +#[cfg(test)] +mod tests; + +/// Uses a sorted slice `data: &[E]` as a kind of "multi-map". The +/// `key_fn` extracts a key of type `K` from the data, and this +/// function finds the range of elements that match the key. `data` +/// must have been sorted as if by a call to `sort_by_key` for this to +/// work. +pub fn binary_search_slice<E, K>(data: &'d [E], key_fn: impl Fn(&E) -> K, key: &K) -> &'d [E] +where + K: Ord, +{ + let mid = match data.binary_search_by_key(key, &key_fn) { + Ok(mid) => mid, + Err(_) => return &[], + }; + let size = data.len(); + + // We get back *some* element with the given key -- so do + // a galloping search backwards to find the *first* one. + let mut start = mid; + let mut previous = mid; + let mut step = 1; + loop { + start = start.saturating_sub(step); + if start == 0 || key_fn(&data[start]) != *key { + break; + } + previous = start; + step *= 2; + } + step = previous - start; + while step > 1 { + let half = step / 2; + let mid = start + half; + if key_fn(&data[mid]) != *key { + start = mid; + } + step -= half; + } + // adjust by one if we have overshot + if start < size && key_fn(&data[start]) != *key { + start += 1; + } + + // Now search forward to find the *last* one. + let mut end = mid; + let mut previous = mid; + let mut step = 1; + loop { + end = end.saturating_add(step).min(size); + if end == size || key_fn(&data[end]) != *key { + break; + } + previous = end; + step *= 2; + } + step = end - previous; + while step > 1 { + let half = step / 2; + let mid = end - half; + if key_fn(&data[mid]) != *key { + end = mid; + } + step -= half; + } + + &data[start..end] +} diff --git a/compiler/rustc_data_structures/src/binary_search_util/tests.rs b/compiler/rustc_data_structures/src/binary_search_util/tests.rs new file mode 100644 index 00000000000..d74febb5c0f --- /dev/null +++ b/compiler/rustc_data_structures/src/binary_search_util/tests.rs @@ -0,0 +1,23 @@ +use super::*; + +type Element = (usize, &'static str); + +fn test_map() -> Vec<Element> { + let mut data = vec![(3, "three-a"), (0, "zero"), (3, "three-b"), (22, "twenty-two")]; + data.sort_by_key(get_key); + data +} + +fn get_key(data: &Element) -> usize { + data.0 +} + +#[test] +fn binary_search_slice_test() { + let map = test_map(); + assert_eq!(binary_search_slice(&map, get_key, &0), &[(0, "zero")]); + assert_eq!(binary_search_slice(&map, get_key, &1), &[]); + assert_eq!(binary_search_slice(&map, get_key, &3), &[(3, "three-a"), (3, "three-b")]); + assert_eq!(binary_search_slice(&map, get_key, &22), &[(22, "twenty-two")]); + assert_eq!(binary_search_slice(&map, get_key, &23), &[]); +} |
