just create a binary search slice helper fn

This commit is contained in:
Niko Matsakis 2019-06-17 07:55:54 -04:00
parent 89a205bf44
commit 3e01c7416a
5 changed files with 72 additions and 114 deletions

View 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]
}

View 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), &[]);
}

View File

@ -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 @@ pub mod tiny_list;
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;

View File

@ -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()
}
}

View File

@ -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), &[]);
}