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
105impl<'a, K, V, P> PartialEq for SparseIndexMap<'a, K, V, P>
106where
107    K: IndexedValue + 'a,
108    P: PointerFamily<'a>,
109    V: PartialEq,
110{
111    fn eq(&self, other: &Self) -> bool {
112        self.map == other.map
113    }
114}
115
116impl<'a, K, V, P> Eq for SparseIndexMap<'a, K, V, P>
117where
118    K: IndexedValue + 'a,
119    P: PointerFamily<'a>,
120    V: Eq,
121{
122}
123
124impl<'a, K, V, P> Index<K::Index> for SparseIndexMap<'a, K, V, P>
125where
126    K: IndexedValue + 'a,
127    P: PointerFamily<'a>,
128{
129    type Output = V;
130
131    fn index(&self, index: K::Index) -> &Self::Output {
132        self.get(index).unwrap()
133    }
134}
135
136impl<'a, K, V, P> IndexMut<K::Index> for SparseIndexMap<'a, K, V, P>
137where
138    K: IndexedValue + 'a,
139    P: PointerFamily<'a>,
140{
141    fn index_mut(&mut self, index: K::Index) -> &mut Self::Output {
142        self.get_mut(index).unwrap()
143    }
144}
145
146impl<'a, K, V, P, M, U> FromIndexicalIterator<'a, K, P, M, (U, V)> for SparseIndexMap<'a, K, V, P>
147where
148    K: IndexedValue + 'a,
149    P: PointerFamily<'a>,
150    U: ToIndex<K, M>,
151{
152    fn from_indexical_iter(
153        iter: impl Iterator<Item = (U, V)>,
154        domain: &P::Pointer<IndexedDomain<K>>,
155    ) -> Self {
156        let map = iter
157            .map(|(u, v)| (u.to_index(domain), v))
158            .collect::<FxHashMap<_, _>>();
159        SparseIndexMap {
160            map,
161            domain: domain.clone(),
162        }
163    }
164}
165
166impl<'a, 'b, K, V, P> IntoIterator for &'b SparseIndexMap<'a, K, V, P>
167where
168    K: IndexedValue + 'a + 'b,
169    V: 'b,
170    P: PointerFamily<'a>,
171{
172    type Item = (&'b K::Index, &'b V);
173    type IntoIter = hash_map::Iter<'b, K::Index, V>;
174
175    fn into_iter(self) -> Self::IntoIter {
176        self.map.iter()
177    }
178}
179
180/// A mapping from indexed keys to values, implemented densely with a vector.
181///
182/// This is more time-efficient than the [`SparseIndexMap`] for lookup,
183/// but it consumes more memory for missing elements.
184pub struct DenseIndexMap<'a, K: IndexedValue + 'a, V, P: PointerFamily<'a>>(
185    IndexVec<'a, K, Option<V>, P>,
186);
187
188/// [`DenseIndexMap`] specialized to the [`RcFamily`].
189pub type DenseRcIndexMap<K, V> = DenseIndexMap<'static, K, V, RcFamily>;
190
191/// [`DenseIndexMap`] specialized to the [`ArcFamily`].
192pub type DenseArcIndexMap<K, V> = DenseIndexMap<'static, K, V, ArcFamily>;
193
194/// [`DenseIndexMap`] specialized to the [`RefFamily`].
195pub type DenseRefIndexMap<'a, K, V> = DenseIndexMap<'a, K, V, RefFamily<'a>>;
196
197impl<'a, K, V, P> DenseIndexMap<'a, K, V, P>
198where
199    K: IndexedValue + 'a,
200    P: PointerFamily<'a>,
201{
202    /// Constructs a new map with an initial element of `mk_elem(i)` for each `i` in `domain`.
203    pub fn new(domain: &P::Pointer<IndexedDomain<K>>) -> Self {
204        Self(IndexVec::from_fn(|_| None, domain))
205    }
206
207    /// Returns an immutable reference to a value for a given key if it exists.
208    pub fn get<M>(&self, idx: impl ToIndex<K, M>) -> Option<&V> {
209        self.0.get(idx).as_ref()
210    }
211
212    /// Returns a mutable reference to a value for a given key if it exists.
213    pub fn get_mut<M>(&mut self, idx: impl ToIndex<K, M>) -> Option<&mut V> {
214        self.0.get_mut(idx).as_mut()
215    }
216
217    /// Returns a reference to a value for a given key.
218    ///
219    /// # Safety
220    /// This function has undefined behavior if `key` is not in `self`.
221    pub unsafe fn get_unchecked<M>(&self, idx: impl ToIndex<K, M>) -> &V {
222        unsafe { self.get(idx).unwrap_unchecked() }
223    }
224
225    /// Returns a mutable reference to a value for a given key.
226    ///
227    /// # Safety
228    /// This function has undefined behavior if `key` is not in `self`.
229    pub unsafe fn get_unchecked_mut<M>(&mut self, idx: impl ToIndex<K, M>) -> &mut V {
230        unsafe { self.get_mut(idx).unwrap_unchecked() }
231    }
232
233    /// Inserts the key/value pair into `self`.
234    pub fn insert<M>(&mut self, idx: impl ToIndex<K, M>, value: V) {
235        let idx = idx.to_index(&self.0.domain);
236        self.0[idx] = Some(value);
237    }
238
239    /// Returns an iterator over the values of the map.
240    pub fn values(&self) -> impl Iterator<Item = &V> {
241        self.0.iter().filter_map(Option::as_ref)
242    }
243}
244
245impl<'a, K, V, P> Index<K::Index> for DenseIndexMap<'a, K, V, P>
246where
247    K: IndexedValue + 'a,
248    P: PointerFamily<'a>,
249{
250    type Output = V;
251
252    fn index(&self, index: K::Index) -> &Self::Output {
253        self.get(index).unwrap()
254    }
255}
256
257impl<'a, K, V, P> IndexMut<K::Index> for DenseIndexMap<'a, K, V, P>
258where
259    K: IndexedValue + 'a,
260    P: PointerFamily<'a>,
261{
262    fn index_mut(&mut self, index: K::Index) -> &mut Self::Output {
263        self.get_mut(index).unwrap()
264    }
265}
266
267impl<'a, K, V, P, M, U> FromIndexicalIterator<'a, K, P, M, (U, V)> for DenseIndexMap<'a, K, V, P>
268where
269    K: IndexedValue + 'a,
270    P: PointerFamily<'a>,
271    U: ToIndex<K, M>,
272{
273    fn from_indexical_iter(
274        iter: impl Iterator<Item = (U, V)>,
275        domain: &P::Pointer<IndexedDomain<K>>,
276    ) -> Self {
277        let mut map = DenseIndexMap::new(domain);
278        for (u, v) in iter {
279            map.insert(u, v);
280        }
281        map
282    }
283}