just create a binary search slice helper fn
This commit is contained in:
parent
89a205bf44
commit
3e01c7416a
49
src/librustc_data_structures/binary_search_util/mod.rs
Normal file
49
src/librustc_data_structures/binary_search_util/mod.rs
Normal file
@ -0,0 +1,49 @@
|
||||
#[cfg(test)]
|
||||
mod test;
|
||||
|
||||
/// 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 &[],
|
||||
};
|
||||
|
||||
// We get back *some* element with the given key -- so
|
||||
// search backwards to find the *first* one.
|
||||
//
|
||||
// (It'd be more efficient to use a "galloping" search
|
||||
// here, but it's not really worth it for small-ish
|
||||
// amounts of data.)
|
||||
let mut start = mid;
|
||||
while start > 0 {
|
||||
if key_fn(&data[start - 1]) == *key {
|
||||
start -= 1;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Now search forward to find the *last* one.
|
||||
//
|
||||
// (It'd be more efficient to use a "galloping" search
|
||||
// here, but it's not really worth it for small-ish
|
||||
// amounts of data.)
|
||||
let mut end = mid + 1;
|
||||
let max = data.len();
|
||||
while end < max {
|
||||
if key_fn(&data[end]) == *key {
|
||||
end += 1;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
&data[start..end]
|
||||
}
|
22
src/librustc_data_structures/binary_search_util/test.rs
Normal file
22
src/librustc_data_structures/binary_search_util/test.rs
Normal file
@ -0,0 +1,22 @@
|
||||
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)
|
||||
}
|
||||
|
||||
fn get_key(data: &Element) -> usize {
|
||||
data.0
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn binary_search_slice() {
|
||||
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), &[]);
|
||||
}
|
@ -72,6 +72,7 @@ macro_rules! unlikely {
|
||||
pub mod macros;
|
||||
pub mod svh;
|
||||
pub mod base_n;
|
||||
pub mod binary_search_util;
|
||||
pub mod bit_set;
|
||||
pub mod box_region;
|
||||
pub mod const_cstr;
|
||||
@ -95,7 +96,6 @@ macro_rules! unlikely {
|
||||
pub mod thin_vec;
|
||||
pub mod transitive_relation;
|
||||
pub use ena::unify;
|
||||
pub mod vec_map;
|
||||
pub mod vec_linked_list;
|
||||
pub mod work_queue;
|
||||
pub mod fingerprint;
|
||||
|
@ -1,85 +0,0 @@
|
||||
#[cfg(test)]
|
||||
mod test;
|
||||
|
||||
/// A (multi-)map based on a sorted vector. This uses binary search to
|
||||
/// find the starting index for a given element and can be a fairly
|
||||
/// efficient data structure, particularly for small-ish sets of data.
|
||||
///
|
||||
/// To use, you supply the starting vector along with a "key fn" that
|
||||
/// extracts the key from each element.
|
||||
pub struct VecMap<E, KeyFn> {
|
||||
data: Vec<E>,
|
||||
key_fn: KeyFn,
|
||||
}
|
||||
|
||||
impl<E, K, KeyFn> VecMap<E, KeyFn>
|
||||
where
|
||||
KeyFn: Fn(&E) -> K,
|
||||
K: Ord + std::fmt::Debug,
|
||||
{
|
||||
pub fn new(
|
||||
mut data: Vec<E>,
|
||||
key_fn: KeyFn,
|
||||
) -> Self {
|
||||
data.sort_by_key(&key_fn);
|
||||
Self { data, key_fn }
|
||||
}
|
||||
|
||||
/// Extract the first index for the given key using binary search.
|
||||
/// Returns `None` if there is no such index.
|
||||
fn get_range(&self, key: &K) -> Option<(usize, usize)> {
|
||||
match self.data.binary_search_by_key(key, &self.key_fn) {
|
||||
Ok(mid) => {
|
||||
// We get back *some* element with the given key -- so
|
||||
// search backwards to find the *first* one.
|
||||
//
|
||||
// (It'd be more efficient to use a "galloping" search
|
||||
// here, but it's not really worth it for small-ish
|
||||
// amounts of data.)
|
||||
let mut start = mid;
|
||||
while start > 0 {
|
||||
if (self.key_fn)(&self.data[start - 1]) == *key {
|
||||
start -= 1;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Now search forward to find the *last* one.
|
||||
//
|
||||
// (It'd be more efficient to use a "galloping" search
|
||||
// here, but it's not really worth it for small-ish
|
||||
// amounts of data.)
|
||||
let mut end = mid + 1;
|
||||
let max = self.data.len();
|
||||
while end < max {
|
||||
if (self.key_fn)(&self.data[end]) == *key {
|
||||
end += 1;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Some((start, end))
|
||||
}
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Gets the (first) value associated with a given key.
|
||||
pub fn get_first(&self, key: &K) -> Option<&E> {
|
||||
let (start, _) = self.get_range(key)?;
|
||||
Some(&self.data[start])
|
||||
}
|
||||
|
||||
/// Gets a slice of values associated with the given key.
|
||||
pub fn get_all(&self, key: &K) -> &[E] {
|
||||
let (start, end) = self.get_range(key).unwrap_or((0, 0));
|
||||
&self.data[start..end]
|
||||
}
|
||||
|
||||
/// Gets a slice of values associated with the given key.
|
||||
pub fn get_iter<'k>(&'k self, key: &'k K) -> impl Iterator<Item = &'k E> {
|
||||
self.get_all(key).iter()
|
||||
}
|
||||
}
|
@ -1,28 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
type Element = (usize, &'static str);
|
||||
|
||||
fn test_map() -> VecMap<Element, impl Fn(&Element) -> usize> {
|
||||
let data = vec![(3, "three-a"), (0, "zero"), (3, "three-b"), (22, "twenty-two")];
|
||||
VecMap::new(data, |&(key, _)| key)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_first() {
|
||||
let map = test_map();
|
||||
assert_eq!(map.get_first(&0), Some(&(0, "zero")));
|
||||
assert_eq!(map.get_first(&1), None);
|
||||
assert_eq!(map.get_first(&3), Some(&(3, "three-a")));
|
||||
assert_eq!(map.get_first(&22), Some(&(22, "twenty-two")));
|
||||
assert_eq!(map.get_first(&23), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_all() {
|
||||
let map = test_map();
|
||||
assert_eq!(map.get_all(&0), &[(0, "zero")]);
|
||||
assert_eq!(map.get_all(&1), &[]);
|
||||
assert_eq!(map.get_all(&3), &[(3, "three-a"), (3, "three-b")]);
|
||||
assert_eq!(map.get_all(&22), &[(22, "twenty-two")]);
|
||||
assert_eq!(map.get_all(&23), &[]);
|
||||
}
|
Loading…
Reference in New Issue
Block a user