use crate::avl::{Iter, IterMut, Tree, WeakTree};
pub use crate::chunk::DEFAULT_SIZE;
use core::{
borrow::Borrow,
cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd},
default::Default,
fmt::{self, Debug, Formatter},
hash::{Hash, Hasher},
iter::FromIterator,
ops::{Index, IndexMut, RangeBounds, RangeFull},
};
#[cfg(feature = "serde")]
use serde::{
de::{MapAccess, Visitor},
ser::SerializeMap,
Deserialize, Deserializer, Serialize, Serializer,
};
#[cfg(feature = "serde")]
use core::marker::PhantomData;
#[cfg(feature = "rayon")]
use rayon::{
iter::{FromParallelIterator, IntoParallelIterator},
prelude::*,
};
#[derive(Clone)]
pub struct Map<K: Ord + Clone, V: Clone, const SIZE: usize>(Tree<K, V, SIZE>);
pub type MapS<K, V> = Map<K, V, { DEFAULT_SIZE / 2 }>;
pub type MapM<K, V> = Map<K, V, DEFAULT_SIZE>;
pub type MapL<K, V> = Map<K, V, { DEFAULT_SIZE * 2 }>;
#[derive(Clone)]
pub struct WeakMapRef<K: Ord + Clone, V: Clone, const SIZE: usize>(WeakTree<K, V, SIZE>);
pub type WeakMapRefS<K, V> = WeakMapRef<K, V, { DEFAULT_SIZE / 2 }>;
pub type WeakMapRefM<K, V> = WeakMapRef<K, V, DEFAULT_SIZE>;
pub type WeakMapRefL<K, V> = WeakMapRef<K, V, { DEFAULT_SIZE * 2 }>;
impl<K, V, const SIZE: usize> WeakMapRef<K, V, SIZE>
where
K: Ord + Clone,
V: Clone,
{
pub fn upgrade(&self) -> Option<Map<K, V, SIZE>> {
self.0.upgrade().map(Map)
}
}
impl<K, V, const SIZE: usize> Hash for Map<K, V, SIZE>
where
K: Hash + Ord + Clone,
V: Hash + Clone,
{
fn hash<H: Hasher>(&self, state: &mut H) {
self.0.hash(state)
}
}
impl<K, V, const SIZE: usize> Default for Map<K, V, SIZE>
where
K: Ord + Clone,
V: Clone,
{
fn default() -> Map<K, V, SIZE> {
Map::new()
}
}
impl<K, V, const SIZE: usize> PartialEq for Map<K, V, SIZE>
where
K: PartialEq + Ord + Clone,
V: PartialEq + Clone,
{
fn eq(&self, other: &Map<K, V, SIZE>) -> bool {
self.0 == other.0
}
}
impl<K, V, const SIZE: usize> Eq for Map<K, V, SIZE>
where
K: Eq + Ord + Clone,
V: Eq + Clone,
{
}
impl<K, V, const SIZE: usize> PartialOrd for Map<K, V, SIZE>
where
K: Ord + Clone,
V: PartialOrd + Clone,
{
fn partial_cmp(&self, other: &Map<K, V, SIZE>) -> Option<Ordering> {
self.0.partial_cmp(&other.0)
}
}
impl<K, V, const SIZE: usize> Ord for Map<K, V, SIZE>
where
K: Ord + Clone,
V: Ord + Clone,
{
fn cmp(&self, other: &Map<K, V, SIZE>) -> Ordering {
self.0.cmp(&other.0)
}
}
impl<K, V, const SIZE: usize> Debug for Map<K, V, SIZE>
where
K: Debug + Ord + Clone,
V: Debug + Clone,
{
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
self.0.fmt(f)
}
}
impl<'a, Q, K, V, const SIZE: usize> Index<&'a Q> for Map<K, V, SIZE>
where
Q: Ord,
K: Ord + Clone + Borrow<Q>,
V: Clone,
{
type Output = V;
fn index(&self, k: &Q) -> &V {
self.get(k).expect("element not found for key")
}
}
impl<'a, Q, K, V, const SIZE: usize> IndexMut<&'a Q> for Map<K, V, SIZE>
where
Q: Ord,
K: Ord + Clone + Borrow<Q>,
V: Clone,
{
fn index_mut(&mut self, k: &'a Q) -> &mut Self::Output {
self.get_mut_cow(k).expect("element not found for key")
}
}
impl<K, V, const SIZE: usize> FromIterator<(K, V)> for Map<K, V, SIZE>
where
K: Ord + Clone,
V: Clone,
{
fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
Map::new().insert_many(iter)
}
}
impl<'a, K, V, const SIZE: usize> IntoIterator for &'a Map<K, V, SIZE>
where
K: 'a + Ord + Clone,
V: 'a + Clone,
{
type Item = (&'a K, &'a V);
type IntoIter = Iter<'a, RangeFull, K, K, V, SIZE>;
fn into_iter(self) -> Self::IntoIter {
self.0.into_iter()
}
}
#[cfg(feature = "serde")]
impl<K, V, const SIZE: usize> Serialize for Map<K, V, SIZE>
where
K: Serialize + Clone + Ord,
V: Serialize + Clone,
{
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut map = serializer.serialize_map(Some(self.len()))?;
for (k, v) in self {
map.serialize_entry(k, v)?
}
map.end()
}
}
#[cfg(feature = "serde")]
struct MapVisitor<K: Clone + Ord, V: Clone, const SIZE: usize> {
marker: PhantomData<fn() -> Map<K, V, SIZE>>,
}
#[cfg(feature = "serde")]
impl<'a, K, V, const SIZE: usize> Visitor<'a> for MapVisitor<K, V, SIZE>
where
K: Deserialize<'a> + Clone + Ord,
V: Deserialize<'a> + Clone,
{
type Value = Map<K, V, SIZE>;
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
formatter.write_str("expected an immutable_chunkmap::Map")
}
fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
where
A: MapAccess<'a>,
{
let mut t = Map::<K, V, SIZE>::new();
while let Some((k, v)) = map.next_entry()? {
t.insert_cow(k, v);
}
Ok(t)
}
}
#[cfg(feature = "serde")]
impl<'a, K, V, const SIZE: usize> Deserialize<'a> for Map<K, V, SIZE>
where
K: Deserialize<'a> + Clone + Ord,
V: Deserialize<'a> + Clone,
{
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'a>,
{
deserializer.deserialize_map(MapVisitor {
marker: PhantomData,
})
}
}
#[cfg(feature = "rayon")]
impl<'a, K, V, const SIZE: usize> IntoParallelIterator for &'a Map<K, V, SIZE>
where
K: 'a + Ord + Clone + Send + Sync,
V: Clone + Send + Sync,
{
type Item = (&'a K, &'a V);
type Iter = rayon::vec::IntoIter<(&'a K, &'a V)>;
fn into_par_iter(self) -> Self::Iter {
self.into_iter().collect::<Vec<_>>().into_par_iter()
}
}
#[cfg(feature = "rayon")]
impl<K, V, const SIZE: usize> FromParallelIterator<(K, V)> for Map<K, V, SIZE>
where
K: Ord + Clone + Send + Sync,
V: Clone + Send + Sync,
{
fn from_par_iter<I>(i: I) -> Self
where
I: IntoParallelIterator<Item = (K, V)>,
{
i.into_par_iter()
.fold_with(Map::new(), |mut m, (k, v)| {
m.insert_cow(k, v);
m
})
.reduce_with(|m0, m1| m0.union(&m1, |_k, _v0, v1| Some(v1.clone())))
.unwrap_or_else(Map::new)
}
}
impl<K, V, const SIZE: usize> Map<K, V, SIZE>
where
K: Ord + Clone,
V: Clone,
{
pub fn new() -> Self {
Map(Tree::new())
}
pub fn downgrade(&self) -> WeakMapRef<K, V, SIZE> {
WeakMapRef(self.0.downgrade())
}
pub fn strong_count(&self) -> usize {
self.0.strong_count()
}
pub fn weak_count(&self) -> usize {
self.0.weak_count()
}
pub fn insert_many<E: IntoIterator<Item = (K, V)>>(&self, elts: E) -> Self {
Map(self.0.insert_many(elts))
}
pub fn remove_many<Q, E>(&self, elts: E) -> Self
where
E: IntoIterator<Item = Q>,
Q: Ord,
K: Borrow<Q>,
{
self.update_many(elts.into_iter().map(|q| (q, ())), |_, _, _| None)
}
pub fn update_many<Q, D, E, F>(&self, elts: E, mut f: F) -> Self
where
E: IntoIterator<Item = (Q, D)>,
Q: Ord,
K: Borrow<Q>,
F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
{
Map(self.0.update_many(elts, &mut f))
}
pub fn insert(&self, k: K, v: V) -> (Self, Option<V>) {
let (root, prev) = self.0.insert(k, v);
(Map(root), prev)
}
pub fn insert_cow(&mut self, k: K, v: V) -> Option<V> {
self.0.insert_cow(k, v)
}
pub fn update<Q, D, F>(&self, q: Q, d: D, mut f: F) -> (Self, Option<V>)
where
Q: Ord,
K: Borrow<Q>,
F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
{
let (root, prev) = self.0.update(q, d, &mut f);
(Map(root), prev)
}
pub fn update_cow<Q, D, F>(&mut self, q: Q, d: D, mut f: F) -> Option<V>
where
Q: Ord,
K: Borrow<Q>,
F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
{
self.0.update_cow(q, d, &mut f)
}
pub fn union<F>(&self, other: &Map<K, V, SIZE>, mut f: F) -> Self
where
F: FnMut(&K, &V, &V) -> Option<V>,
{
Map(Tree::union(&self.0, &other.0, &mut f))
}
pub fn intersect<F>(&self, other: &Map<K, V, SIZE>, mut f: F) -> Self
where
F: FnMut(&K, &V, &V) -> Option<V>,
{
Map(Tree::intersect(&self.0, &other.0, &mut f))
}
pub fn diff<F>(&self, other: &Map<K, V, SIZE>, mut f: F) -> Self
where
F: FnMut(&K, &V, &V) -> Option<V>,
K: Debug,
V: Debug,
{
Map(Tree::diff(&self.0, &other.0, &mut f))
}
pub fn get<'a, Q: ?Sized + Ord>(&'a self, k: &Q) -> Option<&'a V>
where
K: Borrow<Q>,
{
self.0.get(k)
}
pub fn get_key<'a, Q: ?Sized + Ord>(&'a self, k: &Q) -> Option<&'a K>
where
K: Borrow<Q>,
{
self.0.get_key(k)
}
pub fn get_full<'a, Q: ?Sized + Ord>(&'a self, k: &Q) -> Option<(&'a K, &'a V)>
where
K: Borrow<Q>,
{
self.0.get_full(k)
}
pub fn get_mut_cow<'a, Q: ?Sized + Ord>(&'a mut self, k: &Q) -> Option<&'a mut V>
where
K: Borrow<Q>,
{
self.0.get_mut_cow(k)
}
pub fn get_or_insert_cow<'a, F>(&'a mut self, k: K, f: F) -> &'a mut V
where
F: FnOnce() -> V,
{
self.0.get_or_insert_cow(k, f)
}
pub fn remove<Q: Sized + Ord>(&self, k: &Q) -> (Self, Option<V>)
where
K: Borrow<Q>,
{
let (t, prev) = self.0.remove(k);
(Map(t), prev)
}
pub fn remove_cow<Q: Sized + Ord>(&mut self, k: &Q) -> Option<V>
where
K: Borrow<Q>,
{
self.0.remove_cow(k)
}
pub fn len(&self) -> usize {
self.0.len()
}
pub fn range<'a, Q, R>(&'a self, r: R) -> Iter<'a, R, Q, K, V, SIZE>
where
Q: Ord + ?Sized + 'a,
K: Borrow<Q>,
R: RangeBounds<Q> + 'a,
{
self.0.range(r)
}
pub fn range_mut_cow<'a, Q, R>(&'a mut self, r: R) -> IterMut<'a, R, Q, K, V, SIZE>
where
Q: Ord + ?Sized + 'a,
K: Borrow<Q>,
R: RangeBounds<Q> + 'a,
{
self.0.range_mut_cow(r)
}
pub fn iter_mut_cow<'a>(&'a mut self) -> IterMut<'a, RangeFull, K, K, V, SIZE> {
self.0.iter_mut_cow()
}
}
impl<K, V, const SIZE: usize> Map<K, V, SIZE>
where
K: Ord + Clone,
V: Clone + Default,
{
pub fn get_or_default_cow<'a>(&'a mut self, k: K) -> &'a mut V {
self.get_or_insert_cow(k, V::default)
}
}
impl<K, V, const SIZE: usize> Map<K, V, SIZE>
where
K: Ord + Clone + Debug,
V: Clone + Debug,
{
#[allow(dead_code)]
pub fn invariant(&self) -> () {
self.0.invariant()
}
}