use std::borrow::Borrow;
use std::hash::{Hash, Hasher};
use std::iter::FromIterator;
use crate::stable_hasher::{HashStable, StableHasher};
use rustc_index::vec::{Idx, IndexVec};
#[derive(Clone, Debug)]
pub struct SortedIndexMultiMap<I: Idx, K, V> {
items: IndexVec<I, (K, V)>,
idx_sorted_by_item_key: Vec<I>,
}
impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
pub fn new() -> Self {
SortedIndexMultiMap { items: IndexVec::new(), idx_sorted_by_item_key: Vec::new() }
}
pub fn len(&self) -> usize {
self.items.len()
}
pub fn is_empty(&self) -> bool {
self.items.is_empty()
}
pub fn into_iter(self) -> impl DoubleEndedIterator<Item = (K, V)> {
self.items.into_iter()
}
pub fn into_iter_enumerated(self) -> impl DoubleEndedIterator<Item = (I, (K, V))> {
self.items.into_iter_enumerated()
}
pub fn iter(&self) -> impl '_ + DoubleEndedIterator<Item = (&K, &V)> {
self.items.iter().map(|(ref k, ref v)| (k, v))
}
pub fn iter_enumerated(&self) -> impl '_ + DoubleEndedIterator<Item = (I, (&K, &V))> {
self.items.iter_enumerated().map(|(i, (ref k, ref v))| (i, (k, v)))
}
pub fn get(&self, idx: I) -> Option<&(K, V)> {
self.items.get(idx)
}
pub fn get_by_key<Q: 'a>(&'a self, key: &Q) -> impl 'a + Iterator<Item = &'a V>
where
Q: Ord + ?Sized,
K: Borrow<Q>,
{
self.get_by_key_enumerated(key).map(|(_, v)| v)
}
pub fn get_by_key_enumerated<Q>(&self, key: &Q) -> impl '_ + Iterator<Item = (I, &V)>
where
Q: Ord + ?Sized,
K: Borrow<Q>,
{
match self.binary_search_idx(key) {
Err(_) => self.idxs_to_items_enumerated(&[]),
Ok(idx) => {
let start = self.find_lower_bound(key, idx);
let end = self.find_upper_bound(key, idx);
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 find_lower_bound<Q>(&self, key: &Q, initial: usize) -> usize
where
Q: Ord + ?Sized,
K: Borrow<Q>,
{
debug_assert!(self.items[self.idx_sorted_by_item_key[initial]].0.borrow() == key);
let mut start = initial;
while start != 0 && self.items[self.idx_sorted_by_item_key[start - 1]].0.borrow() == key {
start -= 1;
}
start
}
fn find_upper_bound<Q>(&self, key: &Q, initial: usize) -> usize
where
Q: Ord + ?Sized,
K: Borrow<Q>,
{
debug_assert!(self.items[self.idx_sorted_by_item_key[initial]].0.borrow() == key);
let mut end = initial + 1;
let len = self.items.len();
while end < len && self.items[self.idx_sorted_by_item_key[end]].0.borrow() == key {
end += 1;
}
end
}
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))
}
}
impl<I: Idx, K: Eq, V: Eq> Eq for SortedIndexMultiMap<I, K, V> {}
impl<I: Idx, K: PartialEq, V: PartialEq> PartialEq for SortedIndexMultiMap<I, K, V> {
fn eq(&self, other: &Self) -> bool {
self.items == other.items
}
}
impl<I: Idx, K, V> Hash for SortedIndexMultiMap<I, K, V>
where
K: Hash,
V: Hash,
{
fn hash<H: Hasher>(&self, hasher: &mut H) {
self.items.hash(hasher)
}
}
impl<I: Idx, K, V, C> HashStable<C> for SortedIndexMultiMap<I, K, V>
where
K: HashStable<C>,
V: HashStable<C>,
{
fn hash_stable(&self, ctx: &mut C, hasher: &mut StableHasher) {
self.items.hash_stable(ctx, hasher)
}
}
impl<I: Idx, K: Ord, V> FromIterator<(K, V)> for SortedIndexMultiMap<I, K, V> {
fn from_iter<J>(iter: J) -> Self
where
J: IntoIterator<Item = (K, V)>,
{
let items = IndexVec::from_iter(iter);
let mut idx_sorted_by_item_key: Vec<_> = items.indices().collect();
idx_sorted_by_item_key.sort_by_key(|&idx| &items[idx].0);
SortedIndexMultiMap { items, idx_sorted_by_item_key }
}
}
impl<I: Idx, K, V> std::ops::Index<I> for SortedIndexMultiMap<I, K, V> {
type Output = V;
fn index(&self, idx: I) -> &Self::Output {
&self.items[idx].1
}
}
#[cfg(tests)]
mod tests;