1
Fork 0

Improve get_by_key_enumerated more

This commit is contained in:
Yuki Okushi 2021-06-18 18:22:41 +09:00
parent b2b7c859c1
commit cb3b3cf6ab
No known key found for this signature in database
GPG key ID: DABA5B072961C18A
3 changed files with 11 additions and 39 deletions

View file

@ -1,6 +1,5 @@
//! A variant of `SortedMap` that preserves insertion order. //! A variant of `SortedMap` that preserves insertion order.
use std::borrow::Borrow;
use std::hash::{Hash, Hasher}; use std::hash::{Hash, Hasher};
use std::iter::FromIterator; use std::iter::FromIterator;
@ -76,11 +75,7 @@ impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
/// ///
/// If there are multiple items that are equivalent to `key`, they will be yielded in /// If there are multiple items that are equivalent to `key`, they will be yielded in
/// insertion order. /// insertion order.
pub fn get_by_key<Q: 'a>(&'a self, key: &Q) -> impl 'a + Iterator<Item = &'a V> pub fn get_by_key(&'a self, key: K) -> impl 'a + Iterator<Item = &'a V> {
where
Q: Ord + ?Sized,
K: Borrow<Q>,
{
self.get_by_key_enumerated(key).map(|(_, v)| v) self.get_by_key_enumerated(key).map(|(_, v)| v)
} }
@ -89,35 +84,12 @@ impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
/// ///
/// If there are multiple items that are equivalent to `key`, they will be yielded in /// If there are multiple items that are equivalent to `key`, they will be yielded in
/// insertion order. /// insertion order.
pub fn get_by_key_enumerated<Q>(&self, key: &Q) -> impl '_ + Iterator<Item = (I, &V)> pub fn get_by_key_enumerated(&'a self, key: K) -> impl '_ + Iterator<Item = (I, &V)> {
where let lower_bound = self.idx_sorted_by_item_key.partition_point(|&i| self.items[i].0 < key);
Q: Ord + ?Sized, self.idx_sorted_by_item_key[lower_bound..]
K: Borrow<Q>, .iter()
{ .take_while(move |&&i| self.items[i].0.eq(&key))
match self.binary_search_idx(key) { .map(move |&idx| (idx, &self.items[idx].1))
Err(_) => self.idxs_to_items_enumerated(&[]),
Ok(idx) => {
let start = self.idx_sorted_by_item_key[..idx]
.partition_point(|&i| self.items[i].0.borrow() != key);
let end = idx
+ self.idx_sorted_by_item_key[idx..]
.partition_point(|&i| self.items[i].0.borrow() == key);
self.idxs_to_items_enumerated(&self.idx_sorted_by_item_key[start..end])
}
}
}
fn binary_search_idx<Q>(&self, key: &Q) -> Result<usize, usize>
where
Q: Ord + ?Sized,
K: Borrow<Q>,
{
self.idx_sorted_by_item_key.binary_search_by(|&idx| self.items[idx].0.borrow().cmp(key))
}
fn idxs_to_items_enumerated(&'a self, idxs: &'a [I]) -> impl 'a + Iterator<Item = (I, &'a V)> {
idxs.iter().map(move |&idx| (idx, &self.items[idx].1))
} }
} }

View file

@ -14,11 +14,11 @@ fn test_sorted_index_multi_map() {
} }
// `get_by_key` works. // `get_by_key` works.
assert_eq!(set.get_by_key(&3).copied().collect::<Vec<_>>(), vec![0]); assert_eq!(set.get_by_key(3).copied().collect::<Vec<_>>(), vec![0]);
assert!(set.get_by_key(&4).next().is_none()); assert!(set.get_by_key(4).next().is_none());
// `get_by_key` returns items in insertion order. // `get_by_key` returns items in insertion order.
let twos: Vec<_> = set.get_by_key_enumerated(&2).collect(); let twos: Vec<_> = set.get_by_key_enumerated(2).collect();
let idxs: Vec<usize> = twos.iter().map(|(i, _)| *i).collect(); let idxs: Vec<usize> = twos.iter().map(|(i, _)| *i).collect();
let values: Vec<usize> = twos.iter().map(|(_, &v)| v).collect(); let values: Vec<usize> = twos.iter().map(|(_, &v)| v).collect();

View file

@ -124,7 +124,7 @@ impl<'tcx> AssocItems<'tcx> {
&self, &self,
name: Symbol, name: Symbol,
) -> impl '_ + Iterator<Item = &ty::AssocItem> { ) -> impl '_ + Iterator<Item = &ty::AssocItem> {
self.items.get_by_key(&name).copied() self.items.get_by_key(name).copied()
} }
/// Returns an iterator over all associated items with the given name. /// Returns an iterator over all associated items with the given name.