summaryrefslogtreecommitdiffstats
path: root/compiler/rustc_data_structures/src/binary_search_util
diff options
context:
space:
mode:
authorDaniel Baumann <daniel.baumann@progress-linux.org>2024-04-17 12:02:58 +0000
committerDaniel Baumann <daniel.baumann@progress-linux.org>2024-04-17 12:02:58 +0000
commit698f8c2f01ea549d77d7dc3338a12e04c11057b9 (patch)
tree173a775858bd501c378080a10dca74132f05bc50 /compiler/rustc_data_structures/src/binary_search_util
parentInitial commit. (diff)
downloadrustc-698f8c2f01ea549d77d7dc3338a12e04c11057b9.tar.xz
rustc-698f8c2f01ea549d77d7dc3338a12e04c11057b9.zip
Adding upstream version 1.64.0+dfsg1.upstream/1.64.0+dfsg1
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'compiler/rustc_data_structures/src/binary_search_util')
-rw-r--r--compiler/rustc_data_structures/src/binary_search_util/mod.rs68
-rw-r--r--compiler/rustc_data_structures/src/binary_search_util/tests.rs23
2 files changed, 91 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 000000000..d40172a2e
--- /dev/null
+++ b/compiler/rustc_data_structures/src/binary_search_util/mod.rs
@@ -0,0 +1,68 @@
+#[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<'d, E, K>(data: &'d [E], key_fn: impl Fn(&E) -> K, key: &K) -> &'d [E]
+where
+ K: Ord,
+{
+ let Ok(mid) = data.binary_search_by_key(key, &key_fn) else {
+ 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 000000000..d74febb5c
--- /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), &[]);
+}