use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use std::{
cmp::Ordering,
collections::HashMap,
fmt::Debug,
hash::{Hash, Hasher},
ops::{Index, IndexMut, Range},
};
#[derive(Clone, Debug, Serialize, Deserialize, Default)]
pub struct SortedVecMap<K, V, const N: usize = 8> {
inner: SmallVec<[(K, V); N]>,
}
impl<K: Ord, V> SortedVecMap<K, V> {
#[inline]
pub const fn new() -> Self {
Self {
inner: SmallVec::new_const(),
}
}
pub fn with_capacity(capacity: usize) -> Self {
Self {
inner: SmallVec::with_capacity(capacity),
}
}
#[inline]
pub fn len(&self) -> usize {
self.inner.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
#[inline]
pub fn capacity(&self) -> usize {
self.inner.capacity()
}
#[inline]
pub fn clear(&mut self) {
self.inner.clear();
}
#[inline]
pub fn contains_key(&self, key: &K) -> bool {
self.inner.iter().any(|(k, _)| k == key)
}
#[inline]
pub fn get(&self, key: &K) -> Option<&V> {
self.inner.iter().find(|(k, _)| k == key).map(|(_, v)| v)
}
#[inline]
pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
self.inner
.iter_mut()
.find(|(k, _)| k == key)
.map(|(_, v)| v)
}
#[inline]
pub fn get_key_value(&self, key: &K) -> Option<(&K, &V)> {
self.inner
.iter()
.find(|(k, _)| k == key)
.map(|(k, v)| (k, v))
}
pub fn insert(&mut self, key: K, value: V) -> Option<V> {
for (i, (k, v)) in self.inner.iter_mut().enumerate() {
match key.cmp(k) {
Ordering::Less => {
self.inner.insert(i, (key, value));
return None;
}
Ordering::Equal => {
return Some(std::mem::replace(v, value));
}
Ordering::Greater => continue,
}
}
self.inner.push((key, value));
None
}
pub fn remove(&mut self, key: &K) -> Option<V> {
self.inner
.iter()
.position(|(k, _)| k == key)
.map(|pos| self.inner.remove(pos).1)
}
pub fn remove_entry(&mut self, key: &K) -> Option<(K, V)> {
self.inner
.iter()
.position(|(k, _)| k == key)
.map(|pos| self.inner.remove(pos))
}
pub fn retain<F>(&mut self, mut f: F)
where
F: FnMut(&K, &mut V) -> bool,
{
self.inner.retain_mut(|(k, v)| f(k, v));
}
pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
if let Some(pos) = self.inner.iter().position(|(k, _)| *k == key) {
Entry::Occupied(OccupiedEntry {
map: self,
position: pos,
})
} else {
Entry::Vacant(VacantEntry { key, map: self })
}
}
pub fn merge(mut self, mut other: Self) -> Self
where
K: Eq,
{
self.inner.append(&mut other.inner);
self.sort_and_dedup();
self
}
pub fn append(&mut self, other: &mut Self)
where
K: Eq,
{
self.inner.append(&mut other.inner);
self.sort_and_dedup();
}
#[inline]
pub fn keys(&self) -> impl Iterator<Item = &K> {
self.inner.iter().map(|(k, _)| k)
}
#[inline]
pub fn values(&self) -> impl Iterator<Item = &V> {
self.inner.iter().map(|(_, v)| v)
}
#[inline]
pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
self.inner.iter_mut().map(|(_, v)| v)
}
#[inline]
pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
self.inner.iter().map(|(k, v)| (k, v))
}
#[inline]
pub fn iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
self.inner.iter_mut().map(|(k, v)| (&*k, v))
}
fn sort_and_dedup(&mut self)
where
K: Eq,
{
self.inner.sort_by(|a, b| a.0.cmp(&b.0));
self.inner.dedup_by(|a, b| a.0 == b.0);
}
}
impl<K: Ord + Sync, V: Sync> SortedVecMap<K, V> {
#[inline]
pub fn par_iter(&self) -> impl ParallelIterator<Item = (&K, &V)> {
self.inner.par_iter().map(|(k, v)| (k, v))
}
}
pub enum Entry<'a, K, V> {
Occupied(OccupiedEntry<'a, K, V>),
Vacant(VacantEntry<'a, K, V>),
}
pub struct OccupiedEntry<'a, K, V> {
map: &'a mut SortedVecMap<K, V>,
position: usize,
}
pub struct VacantEntry<'a, K, V> {
key: K,
map: &'a mut SortedVecMap<K, V>,
}
impl<'a, K: Ord, V> Entry<'a, K, V> {
pub fn or_insert(self, default: V) -> &'a mut V {
match self {
Entry::Occupied(e) => e.into_mut(),
Entry::Vacant(e) => e.insert(default),
}
}
pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> &'a mut V {
match self {
Entry::Occupied(e) => e.into_mut(),
Entry::Vacant(e) => e.insert(default()),
}
}
pub fn or_insert_with_key<F: FnOnce(&K) -> V>(self, default: F) -> &'a mut V {
match self {
Entry::Occupied(e) => e.into_mut(),
Entry::Vacant(e) => {
let value = default(&e.key);
e.insert(value)
}
}
}
pub fn and_modify<F: FnOnce(&mut V)>(mut self, f: F) -> Self {
match &mut self {
Entry::Occupied(e) => f(e.get_mut()),
Entry::Vacant(_) => {}
}
self
}
pub fn key(&self) -> &K {
match self {
Entry::Occupied(e) => e.key(),
Entry::Vacant(e) => &e.key,
}
}
}
impl<'a, K: Ord, V: Default> Entry<'a, K, V> {
pub fn or_default(self) -> &'a mut V {
self.or_insert_with(Default::default)
}
}
impl<'a, K: Ord, V> VacantEntry<'a, K, V> {
pub fn key(&self) -> &K {
&self.key
}
pub fn into_key(self) -> K {
self.key
}
pub fn insert(self, value: V) -> &'a mut V {
let pos = self
.map
.inner
.iter()
.position(|(k, _)| k > &self.key)
.unwrap_or(self.map.inner.len());
self.map.inner.insert(pos, (self.key, value));
&mut self.map.inner[pos].1
}
}
impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
pub fn key(&self) -> &K {
&self.map.inner[self.position].0
}
pub fn get(&self) -> &V {
&self.map.inner[self.position].1
}
pub fn get_mut(&mut self) -> &mut V {
&mut self.map.inner[self.position].1
}
pub fn into_mut(self) -> &'a mut V {
&mut self.map.inner[self.position].1
}
pub fn insert(&mut self, value: V) -> V {
std::mem::replace(&mut self.map.inner[self.position].1, value)
}
pub fn remove(self) -> V {
self.map.inner.remove(self.position).1
}
pub fn remove_entry(self) -> (K, V) {
self.map.inner.remove(self.position)
}
}
impl<K: Ord, V> Extend<(K, V)> for SortedVecMap<K, V>
where
K: Eq,
{
fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
self.inner.extend(iter);
self.sort_and_dedup();
}
}
impl<K, V, const N: usize> IntoIterator for SortedVecMap<K, V, N> {
type Item = (K, V);
type IntoIter = smallvec::IntoIter<[(K, V); N]>;
fn into_iter(self) -> Self::IntoIter {
self.inner.into_iter()
}
}
impl<'a, K, V> IntoIterator for &'a SortedVecMap<K, V> {
type Item = (&'a K, &'a V);
type IntoIter = std::iter::Map<std::slice::Iter<'a, (K, V)>, fn(&'a (K, V)) -> (&'a K, &'a V)>;
fn into_iter(self) -> Self::IntoIter {
self.inner.iter().map(|(k, v)| (k, v))
}
}
impl<'a, K, V> IntoIterator for &'a mut SortedVecMap<K, V> {
type Item = (&'a K, &'a mut V);
type IntoIter =
std::iter::Map<std::slice::IterMut<'a, (K, V)>, fn(&'a mut (K, V)) -> (&'a K, &'a mut V)>;
fn into_iter(self) -> Self::IntoIter {
self.inner.iter_mut().map(|(k, v)| (&*k, v))
}
}
impl<K: Ord, V> FromIterator<(K, V)> for SortedVecMap<K, V>
where
K: Eq,
{
fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
let mut map = Self {
inner: iter.into_iter().collect(),
};
map.sort_and_dedup();
map
}
}
impl<K: Ord, V> From<HashMap<K, V>> for SortedVecMap<K, V>
where
K: Eq,
{
fn from(hash_map: HashMap<K, V>) -> Self {
let mut map = Self {
inner: hash_map.into_iter().collect(),
};
map.sort_and_dedup();
map
}
}
impl<K: PartialEq, V: PartialEq> PartialEq for SortedVecMap<K, V> {
fn eq(&self, other: &Self) -> bool {
self.inner == other.inner
}
}
impl<K: Eq, V: Eq> Eq for SortedVecMap<K, V> {}
impl<K: Ord, V> From<(K, V)> for SortedVecMap<K, V> {
fn from(tuple: (K, V)) -> Self {
let mut inner = SmallVec::new();
inner.push(tuple);
Self { inner }
}
}
impl<K: Ord, V: PartialOrd> PartialOrd for SortedVecMap<K, V> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
self.inner.partial_cmp(&other.inner)
}
}
impl<K: Ord, V: Ord> Ord for SortedVecMap<K, V> {
fn cmp(&self, other: &Self) -> Ordering {
self.inner.cmp(&other.inner)
}
}
impl<K: Hash, V: Hash> Hash for SortedVecMap<K, V> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.inner.hash(state);
}
}
impl<K: Ord, V, const N: usize> From<[(K, V); N]> for SortedVecMap<K, V>
where
K: Eq,
{
fn from(arr: [(K, V); N]) -> Self {
let mut map = Self {
inner: SmallVec::from_iter(arr),
};
map.sort_and_dedup();
map
}
}
impl<K: Ord, V> From<Vec<(K, V)>> for SortedVecMap<K, V>
where
K: Eq,
{
fn from(vec: Vec<(K, V)>) -> Self {
let mut map = Self {
inner: SmallVec::from_vec(vec),
};
map.sort_and_dedup();
map
}
}
impl<K: Ord, V> Index<&K> for SortedVecMap<K, V> {
type Output = V;
fn index(&self, key: &K) -> &Self::Output {
self.get(key).expect("key not found")
}
}
impl<K: Ord, V> IndexMut<&K> for SortedVecMap<K, V> {
fn index_mut(&mut self, key: &K) -> &mut Self::Output {
self.get_mut(key).expect("key not found")
}
}
impl<K, V> Index<Range<usize>> for SortedVecMap<K, V> {
type Output = [(K, V)];
fn index(&self, index: Range<usize>) -> &Self::Output {
&self.inner[index]
}
}
impl<K, V> IndexMut<Range<usize>> for SortedVecMap<K, V> {
fn index_mut(&mut self, index: Range<usize>) -> &mut Self::Output {
&mut self.inner[index]
}
}