Skip to main content

indexical/
map.rs

1//! Map-like collections for indexed keys.
2
3use std::{
4    collections::hash_map,
5    ops::{Index, IndexMut},
6};
7
8use rustc_hash::FxHashMap;
9
10use crate::{
11    FromIndexicalIterator, IndexedDomain, IndexedValue, ToIndex,
12    pointer::{ArcFamily, PointerFamily, RcFamily, RefFamily},
13    vec::IndexVec,
14};
15
16/// A mapping from indexed keys to values, implemented sparsely with a hash map.
17///
18/// This is more memory-efficient than the [`DenseIndexMap`] with a small
19/// number of keys.
20pub struct SparseIndexMap<'a, K: IndexedValue + 'a, V, P: PointerFamily<'a>> {
21    map: FxHashMap<K::Index, V>,
22    domain: P::Pointer<IndexedDomain<K>>,
23}
24
25/// [`SparseIndexMap`] specialized to the [`RcFamily`].
26pub type SparseRcIndexMap<K, V> = SparseIndexMap<'static, K, V, RcFamily>;
27
28/// [`SparseIndexMap`] specialized to the [`ArcFamily`].
29pub type SparseArcIndexMap<K, V> = SparseIndexMap<'static, K, V, ArcFamily>;
30
31/// [`SparseIndexMap`] specialized to the [`RefFamily`].
32pub type SparseRefIndexMap<'a, K, V> = SparseIndexMap<'a, K, V, RefFamily<'a>>;
33
34impl<'a, K, V, P> SparseIndexMap<'a, K, V, P>
35where
36    K: IndexedValue + 'a,
37    P: PointerFamily<'a>,
38{
39    /// Constructs an empty map within the given domain.
40    pub fn new(domain: &P::Pointer<IndexedDomain<K>>) -> Self {
41        SparseIndexMap {
42            map: FxHashMap::default(),
43            domain: domain.clone(),
44        }
45    }
46
47    /// Returns an immutable reference to a value for a given key if it exists.
48    pub fn get<M>(&self, key: impl ToIndex<K, M>) -> Option<&V> {
49        let idx = key.to_index(&self.domain);
50        self.map.get(&idx)
51    }
52
53    /// Returns a mutable reference to a value for a given key if it exists.
54    pub fn get_mut<M>(&mut self, key: impl ToIndex<K, M>) -> Option<&mut V> {
55        let idx = key.to_index(&self.domain);
56        self.map.get_mut(&idx)
57    }
58
59    /// Returns a reference to a value for a given key.
60    ///
61    /// # Safety
62    /// This function has undefined behavior if `key` is not in `self`.
63    pub unsafe fn get_unchecked<M>(&self, key: impl ToIndex<K, M>) -> &V {
64        let idx = key.to_index(&self.domain);
65        unsafe { self.map.get(&idx).unwrap_unchecked() }
66    }
67
68    /// Returns a mutable reference to a value for a given key.
69    ///
70    /// # Safety
71    /// This function has undefined behavior if `key` is not in `self`.
72    pub unsafe fn get_unchecked_mut<M>(&mut self, key: impl ToIndex<K, M>) -> &mut V {
73        let idx = key.to_index(&self.domain);
74        unsafe { self.map.get_mut(&idx).unwrap_unchecked() }
75    }
76
77    /// Inserts the key/value pair into `self`.
78    pub fn insert<M>(&mut self, key: impl ToIndex<K, M>, value: V) {
79        let idx = key.to_index(&self.domain);
80        self.map.insert(idx, value);
81    }
82
83    /// Returns an iterator over the values of the map.
84    pub fn values(&self) -> impl ExactSizeIterator<Item = &V> {
85        self.map.values()
86    }
87
88    /// Returns a mutable entry into the map for the given key.
89    pub fn entry<M>(&mut self, key: impl ToIndex<K, M>) -> hash_map::Entry<'_, K::Index, V> {
90        let idx = key.to_index(&self.domain);
91        self.map.entry(idx)
92    }
93
94    /// Returns the number of entries in the map.
95    pub fn len(&self) -> usize {
96        self.map.len()
97    }
98
99    /// Returns true if the map has no elements.
100    pub fn is_empty(&self) -> bool {
101        self.map.is_empty()
102    }
103
104    /// Returns an iterator over pairs of keys and values in the map.
105    pub fn iter(&self) -> hash_map::Iter<'_, K::Index, V> {
106        self.map.iter()
107    }
108}
109
110impl<'a, K, V, P> PartialEq for SparseIndexMap<'a, K, V, P>
111where
112    K: IndexedValue + 'a,
113    P: PointerFamily<'a>,
114    V: PartialEq,
115{
116    fn eq(&self, other: &Self) -> bool {
117        self.map == other.map
118    }
119}
120
121impl<'a, K, V, P> Eq for SparseIndexMap<'a, K, V, P>
122where
123    K: IndexedValue + 'a,
124    P: PointerFamily<'a>,
125    V: Eq,
126{
127}
128
129impl<'a, K, V, P> Index<K::Index> for SparseIndexMap<'a, K, V, P>
130where
131    K: IndexedValue + 'a,
132    P: PointerFamily<'a>,
133{
134    type Output = V;
135
136    fn index(&self, index: K::Index) -> &Self::Output {
137        self.get(index).unwrap()
138    }
139}
140
141impl<'a, K, V, P> IndexMut<K::Index> for SparseIndexMap<'a, K, V, P>
142where
143    K: IndexedValue + 'a,
144    P: PointerFamily<'a>,
145{
146    fn index_mut(&mut self, index: K::Index) -> &mut Self::Output {
147        self.get_mut(index).unwrap()
148    }
149}
150
151impl<'a, K, V, P, M, U> FromIndexicalIterator<'a, K, P, M, (U, V)> for SparseIndexMap<'a, K, V, P>
152where
153    K: IndexedValue + 'a,
154    P: PointerFamily<'a>,
155    U: ToIndex<K, M>,
156{
157    fn from_indexical_iter(
158        iter: impl Iterator<Item = (U, V)>,
159        domain: &P::Pointer<IndexedDomain<K>>,
160    ) -> Self {
161        let map = iter
162            .map(|(u, v)| (u.to_index(domain), v))
163            .collect::<FxHashMap<_, _>>();
164        SparseIndexMap {
165            map,
166            domain: domain.clone(),
167        }
168    }
169}
170
171impl<'a, 'b, K, V, P> IntoIterator for &'b SparseIndexMap<'a, K, V, P>
172where
173    K: IndexedValue + 'a + 'b,
174    V: 'b,
175    P: PointerFamily<'a>,
176{
177    type Item = (&'b K::Index, &'b V);
178    type IntoIter = hash_map::Iter<'b, K::Index, V>;
179
180    fn into_iter(self) -> Self::IntoIter {
181        self.iter()
182    }
183}
184
185/// A mapping from indexed keys to values, implemented densely with a vector.
186///
187/// This is more time-efficient than the [`SparseIndexMap`] for lookup,
188/// but it consumes more memory for missing elements.
189pub struct DenseIndexMap<'a, K: IndexedValue + 'a, V, P: PointerFamily<'a>>(
190    IndexVec<'a, K, Option<V>, P>,
191);
192
193/// [`DenseIndexMap`] specialized to the [`RcFamily`].
194pub type DenseRcIndexMap<K, V> = DenseIndexMap<'static, K, V, RcFamily>;
195
196/// [`DenseIndexMap`] specialized to the [`ArcFamily`].
197pub type DenseArcIndexMap<K, V> = DenseIndexMap<'static, K, V, ArcFamily>;
198
199/// [`DenseIndexMap`] specialized to the [`RefFamily`].
200pub type DenseRefIndexMap<'a, K, V> = DenseIndexMap<'a, K, V, RefFamily<'a>>;
201
202impl<'a, K, V, P> DenseIndexMap<'a, K, V, P>
203where
204    K: IndexedValue + 'a,
205    P: PointerFamily<'a>,
206{
207    /// Constructs a new map with an initial element of `mk_elem(i)` for each `i` in `domain`.
208    pub fn new(domain: &P::Pointer<IndexedDomain<K>>) -> Self {
209        Self(IndexVec::from_fn(|_| None, domain))
210    }
211
212    /// Returns an immutable reference to a value for a given key if it exists.
213    pub fn get<M>(&self, idx: impl ToIndex<K, M>) -> Option<&V> {
214        self.0.get(idx).as_ref()
215    }
216
217    /// Returns a mutable reference to a value for a given key if it exists.
218    pub fn get_mut<M>(&mut self, idx: impl ToIndex<K, M>) -> Option<&mut V> {
219        self.0.get_mut(idx).as_mut()
220    }
221
222    /// Returns a reference to a value for a given key.
223    ///
224    /// # Safety
225    /// This function has undefined behavior if `key` is not in `self`.
226    pub unsafe fn get_unchecked<M>(&self, idx: impl ToIndex<K, M>) -> &V {
227        unsafe { self.get(idx).unwrap_unchecked() }
228    }
229
230    /// Returns a mutable reference to a value for a given key.
231    ///
232    /// # Safety
233    /// This function has undefined behavior if `key` is not in `self`.
234    pub unsafe fn get_unchecked_mut<M>(&mut self, idx: impl ToIndex<K, M>) -> &mut V {
235        unsafe { self.get_mut(idx).unwrap_unchecked() }
236    }
237
238    /// Inserts the key/value pair into `self`.
239    pub fn insert<M>(&mut self, idx: impl ToIndex<K, M>, value: V) {
240        let idx = idx.to_index(&self.0.domain);
241        self.0[idx] = Some(value);
242    }
243
244    /// Returns an iterator over the values of the map.
245    pub fn values(&self) -> impl Iterator<Item = &V> {
246        self.0.iter().filter_map(Option::as_ref)
247    }
248}
249
250impl<'a, K, V, P> Index<K::Index> for DenseIndexMap<'a, K, V, P>
251where
252    K: IndexedValue + 'a,
253    P: PointerFamily<'a>,
254{
255    type Output = V;
256
257    fn index(&self, index: K::Index) -> &Self::Output {
258        self.get(index).unwrap()
259    }
260}
261
262impl<'a, K, V, P> IndexMut<K::Index> for DenseIndexMap<'a, K, V, P>
263where
264    K: IndexedValue + 'a,
265    P: PointerFamily<'a>,
266{
267    fn index_mut(&mut self, index: K::Index) -> &mut Self::Output {
268        self.get_mut(index).unwrap()
269    }
270}
271
272impl<'a, K, V, P, M, U> FromIndexicalIterator<'a, K, P, M, (U, V)> for DenseIndexMap<'a, K, V, P>
273where
274    K: IndexedValue + 'a,
275    P: PointerFamily<'a>,
276    U: ToIndex<K, M>,
277{
278    fn from_indexical_iter(
279        iter: impl Iterator<Item = (U, V)>,
280        domain: &P::Pointer<IndexedDomain<K>>,
281    ) -> Self {
282        let mut map = DenseIndexMap::new(domain);
283        for (u, v) in iter {
284            map.insert(u, v);
285        }
286        map
287    }
288}