use smallvec::SmallVec;
use std::collections::HashMap;
use std::hash::Hash;
#[derive(Clone)]
pub struct SMap<K, V>
where
K: Eq + Hash + Ord,
{
pub data: SmallVec<[(K, V); 64]>,
pub ext_data: Option<HashMap<K, V>>,
}
impl<K, V> Default for SMap<K, V>
where
K: Eq + Hash + Ord,
{
fn default() -> Self {
Self::new()
}
}
impl<K, V> SMap<K, V>
where
K: Eq + Hash + Ord,
{
pub fn new() -> Self {
Self {
data: SmallVec::new(),
ext_data: None,
}
}
pub fn insert(&mut self, key: K, mut value: V) -> Option<V> {
if let Some(ext_data) = &mut self.ext_data {
ext_data.insert(key, value)
} else if self.data.len() == 64 {
let mut ext_data = HashMap::with_capacity(128);
for (k, v) in self.data.drain(..) {
ext_data.insert(k, v);
}
ext_data.insert(key, value)
} else {
match self.data.binary_search_by(|probe| probe.0.cmp(&key)) {
Ok(index) => {
std::mem::swap(&mut value, &mut self.data[index].1);
Some(value)
}
Err(index) => {
self.data.insert(index, (key, value));
None
}
}
}
}
pub fn get(&self, key: &K) -> Option<&V> {
if let Some(ext_data) = &self.ext_data {
ext_data.get(key)
} else {
match self.data.binary_search_by(|probe| probe.0.cmp(&key)) {
Ok(index) => Some(&self.data[index].1),
Err(_) => None,
}
}
}
pub fn remove(&mut self, key: &K) -> Option<V> {
if let Some(ext_data) = &mut self.ext_data {
ext_data.remove(key)
} else {
match self.data.binary_search_by(|probe| probe.0.cmp(&key)) {
Ok(index) => Some(self.data.remove(index).1),
Err(_) => None,
}
}
}
pub fn len(&self) -> usize {
self.data.len() + self.ext_data.as_ref().map_or(0, HashMap::len)
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn iter(&self) -> SMapIter<'_, K, V> {
let state = match &self.ext_data {
Some(ext) => IterState::HashMap(ext.iter()),
None => IterState::SmallVec(self.data.iter()),
};
SMapIter { inner: state }
}
pub fn keys(&self) -> SMapKeys<'_, K, V> {
let state = match &self.ext_data {
Some(ext) => IterState::HashMap(ext.iter()),
None => IterState::SmallVec(self.data.iter()),
};
SMapKeys { inner: state }
}
pub fn values(&self) -> SMapValues<'_, K, V> {
let state = match &self.ext_data {
Some(ext) => IterState::HashMap(ext.iter()),
None => IterState::SmallVec(self.data.iter()),
};
SMapValues { inner: state }
}
}
enum IterState<'a, K, V> {
SmallVec(std::slice::Iter<'a, (K, V)>),
HashMap(std::collections::hash_map::Iter<'a, K, V>),
}
pub struct SMapIter<'a, K, V> {
inner: IterState<'a, K, V>,
}
impl<'a, K, V> Iterator for SMapIter<'a, K, V> {
type Item = (&'a K, &'a V);
fn next(&mut self) -> Option<Self::Item> {
match &mut self.inner {
IterState::SmallVec(iter) => {
let item = iter.next()?;
Some((&item.0, &item.1))
}
IterState::HashMap(iter) => iter.next(),
}
}
}
pub struct SMapKeys<'a, K, V> {
inner: IterState<'a, K, V>,
}
impl<'a, K, V> Iterator for SMapKeys<'a, K, V> {
type Item = &'a K;
fn next(&mut self) -> Option<Self::Item> {
match &mut self.inner {
IterState::SmallVec(iter) => Some(&iter.next()?.0),
IterState::HashMap(iter) => iter.next().map(|(k, _)| k),
}
}
}
pub struct SMapValues<'a, K, V> {
inner: IterState<'a, K, V>,
}
impl<'a, K, V> Iterator for SMapValues<'a, K, V> {
type Item = &'a V;
fn next(&mut self) -> Option<Self::Item> {
match &mut self.inner {
IterState::SmallVec(iter) => Some(&iter.next()?.1),
IterState::HashMap(iter) => iter.next().map(|(_, v)| v),
}
}
}