indexical 0.10.3

//! Map-like collections for indexed keys.

use std::{
    collections::hash_map,
    ops::{Index, IndexMut},
};

use rustc_hash::FxHashMap;

use crate::{
    FromIndexicalIterator, IndexedDomain, IndexedValue, ToIndex,
    pointer::{ArcFamily, PointerFamily, RcFamily, RefFamily},
    vec::IndexVec,
};

/// A mapping from indexed keys to values, implemented sparsely with a hash map.
///
/// This is more memory-efficient than the [`DenseIndexMap`] with a small
/// number of keys.
pub struct SparseIndexMap<'a, K: IndexedValue + 'a, V, P: PointerFamily<'a>> {
    map: FxHashMap<K::Index, V>,
    domain: P::Pointer<IndexedDomain<K>>,
}

/// [`SparseIndexMap`] specialized to the [`RcFamily`].
pub type SparseRcIndexMap<K, V> = SparseIndexMap<'static, K, V, RcFamily>;

/// [`SparseIndexMap`] specialized to the [`ArcFamily`].
pub type SparseArcIndexMap<K, V> = SparseIndexMap<'static, K, V, ArcFamily>;

/// [`SparseIndexMap`] specialized to the [`RefFamily`].
pub type SparseRefIndexMap<'a, K, V> = SparseIndexMap<'a, K, V, RefFamily<'a>>;

impl<'a, K, V, P> SparseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a,
    P: PointerFamily<'a>,
{
    /// Constructs an empty map within the given domain.
    pub fn new(domain: &P::Pointer<IndexedDomain<K>>) -> Self {
        SparseIndexMap {
            map: FxHashMap::default(),
            domain: domain.clone(),
        }
    }

    /// Returns an immutable reference to a value for a given key if it exists.
    pub fn get<M>(&self, key: impl ToIndex<K, M>) -> Option<&V> {
        let idx = key.to_index(&self.domain);
        self.map.get(&idx)
    }

    /// Returns a mutable reference to a value for a given key if it exists.
    pub fn get_mut<M>(&mut self, key: impl ToIndex<K, M>) -> Option<&mut V> {
        let idx = key.to_index(&self.domain);
        self.map.get_mut(&idx)
    }

    /// Returns a reference to a value for a given key.
    ///
    /// # Safety
    /// This function has undefined behavior if `key` is not in `self`.
    pub unsafe fn get_unchecked<M>(&self, key: impl ToIndex<K, M>) -> &V {
        let idx = key.to_index(&self.domain);
        unsafe { self.map.get(&idx).unwrap_unchecked() }
    }

    /// Returns a mutable reference to a value for a given key.
    ///
    /// # Safety
    /// This function has undefined behavior if `key` is not in `self`.
    pub unsafe fn get_unchecked_mut<M>(&mut self, key: impl ToIndex<K, M>) -> &mut V {
        let idx = key.to_index(&self.domain);
        unsafe { self.map.get_mut(&idx).unwrap_unchecked() }
    }

    /// Inserts the key/value pair into `self`.
    pub fn insert<M>(&mut self, key: impl ToIndex<K, M>, value: V) {
        let idx = key.to_index(&self.domain);
        self.map.insert(idx, value);
    }

    /// Returns an iterator over the values of the map.
    pub fn values(&self) -> impl ExactSizeIterator<Item = &V> {
        self.map.values()
    }

    /// Returns a mutable entry into the map for the given key.
    pub fn entry<M>(&mut self, key: impl ToIndex<K, M>) -> hash_map::Entry<'_, K::Index, V> {
        let idx = key.to_index(&self.domain);
        self.map.entry(idx)
    }

    /// Returns the number of entries in the map.
    pub fn len(&self) -> usize {
        self.map.len()
    }

    /// Returns true if the map has no elements.
    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    /// Returns an iterator over pairs of keys and values in the map.
    pub fn iter(&self) -> hash_map::Iter<'_, K::Index, V> {
        self.map.iter()
    }
}

impl<'a, K, V, P> PartialEq for SparseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a,
    P: PointerFamily<'a>,
    V: PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        self.map == other.map
    }
}

impl<'a, K, V, P> Eq for SparseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a,
    P: PointerFamily<'a>,
    V: Eq,
{
}

impl<'a, K, V, P> Index<K::Index> for SparseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a,
    P: PointerFamily<'a>,
{
    type Output = V;

    fn index(&self, index: K::Index) -> &Self::Output {
        self.get(index).unwrap()
    }
}

impl<'a, K, V, P> IndexMut<K::Index> for SparseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a,
    P: PointerFamily<'a>,
{
    fn index_mut(&mut self, index: K::Index) -> &mut Self::Output {
        self.get_mut(index).unwrap()
    }
}

impl<'a, K, V, P, M, U> FromIndexicalIterator<'a, K, P, M, (U, V)> for SparseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a,
    P: PointerFamily<'a>,
    U: ToIndex<K, M>,
{
    fn from_indexical_iter(
        iter: impl Iterator<Item = (U, V)>,
        domain: &P::Pointer<IndexedDomain<K>>,
    ) -> Self {
        let map = iter
            .map(|(u, v)| (u.to_index(domain), v))
            .collect::<FxHashMap<_, _>>();
        SparseIndexMap {
            map,
            domain: domain.clone(),
        }
    }
}

impl<'a, 'b, K, V, P> IntoIterator for &'b SparseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a + 'b,
    V: 'b,
    P: PointerFamily<'a>,
{
    type Item = (&'b K::Index, &'b V);
    type IntoIter = hash_map::Iter<'b, K::Index, V>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

/// A mapping from indexed keys to values, implemented densely with a vector.
///
/// This is more time-efficient than the [`SparseIndexMap`] for lookup,
/// but it consumes more memory for missing elements.
pub struct DenseIndexMap<'a, K: IndexedValue + 'a, V, P: PointerFamily<'a>>(
    IndexVec<'a, K, Option<V>, P>,
);

/// [`DenseIndexMap`] specialized to the [`RcFamily`].
pub type DenseRcIndexMap<K, V> = DenseIndexMap<'static, K, V, RcFamily>;

/// [`DenseIndexMap`] specialized to the [`ArcFamily`].
pub type DenseArcIndexMap<K, V> = DenseIndexMap<'static, K, V, ArcFamily>;

/// [`DenseIndexMap`] specialized to the [`RefFamily`].
pub type DenseRefIndexMap<'a, K, V> = DenseIndexMap<'a, K, V, RefFamily<'a>>;

impl<'a, K, V, P> DenseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a,
    P: PointerFamily<'a>,
{
    /// Constructs a new map with an initial element of `mk_elem(i)` for each `i` in `domain`.
    pub fn new(domain: &P::Pointer<IndexedDomain<K>>) -> Self {
        Self(IndexVec::from_fn(|_| None, domain))
    }

    /// Returns an immutable reference to a value for a given key if it exists.
    pub fn get<M>(&self, idx: impl ToIndex<K, M>) -> Option<&V> {
        self.0.get(idx).as_ref()
    }

    /// Returns a mutable reference to a value for a given key if it exists.
    pub fn get_mut<M>(&mut self, idx: impl ToIndex<K, M>) -> Option<&mut V> {
        self.0.get_mut(idx).as_mut()
    }

    /// Returns a reference to a value for a given key.
    ///
    /// # Safety
    /// This function has undefined behavior if `key` is not in `self`.
    pub unsafe fn get_unchecked<M>(&self, idx: impl ToIndex<K, M>) -> &V {
        unsafe { self.get(idx).unwrap_unchecked() }
    }

    /// Returns a mutable reference to a value for a given key.
    ///
    /// # Safety
    /// This function has undefined behavior if `key` is not in `self`.
    pub unsafe fn get_unchecked_mut<M>(&mut self, idx: impl ToIndex<K, M>) -> &mut V {
        unsafe { self.get_mut(idx).unwrap_unchecked() }
    }

    /// Inserts the key/value pair into `self`.
    pub fn insert<M>(&mut self, idx: impl ToIndex<K, M>, value: V) {
        let idx = idx.to_index(&self.0.domain);
        self.0[idx] = Some(value);
    }

    /// Returns an iterator over the values of the map.
    pub fn values(&self) -> impl Iterator<Item = &V> {
        self.0.iter().filter_map(Option::as_ref)
    }
}

impl<'a, K, V, P> Index<K::Index> for DenseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a,
    P: PointerFamily<'a>,
{
    type Output = V;

    fn index(&self, index: K::Index) -> &Self::Output {
        self.get(index).unwrap()
    }
}

impl<'a, K, V, P> IndexMut<K::Index> for DenseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a,
    P: PointerFamily<'a>,
{
    fn index_mut(&mut self, index: K::Index) -> &mut Self::Output {
        self.get_mut(index).unwrap()
    }
}

impl<'a, K, V, P, M, U> FromIndexicalIterator<'a, K, P, M, (U, V)> for DenseIndexMap<'a, K, V, P>
where
    K: IndexedValue + 'a,
    P: PointerFamily<'a>,
    U: ToIndex<K, M>,
{
    fn from_indexical_iter(
        iter: impl Iterator<Item = (U, V)>,
        domain: &P::Pointer<IndexedDomain<K>>,
    ) -> Self {
        let mut map = DenseIndexMap::new(domain);
        for (u, v) in iter {
            map.insert(u, v);
        }
        map
    }
}