cranelift_entity/
map.rs

1//! Densely numbered entity references as mapping keys.
2
3use crate::EntityRef;
4use crate::iter::{Iter, IterMut};
5use crate::keys::Keys;
6use alloc::vec::Vec;
7use core::cmp::min;
8use core::fmt;
9use core::marker::PhantomData;
10use core::ops::{Index, IndexMut};
11use core::slice;
12#[cfg(feature = "enable-serde")]
13use serde::{
14    Deserialize, Serialize,
15    de::{Deserializer, SeqAccess, Visitor},
16    ser::{SerializeSeq, Serializer},
17};
18
19/// A mapping `K -> V` for densely indexed entity references.
20///
21/// The `SecondaryMap` data structure uses the dense index space to implement a map with a vector.
22/// Unlike `PrimaryMap`, an `SecondaryMap` can't be used to allocate entity references. It is used
23/// to associate secondary information with entities.
24///
25/// The map does not track if an entry for a key has been inserted or not. Instead it behaves as if
26/// all keys have a default entry from the beginning.
27#[derive(Clone, Hash)]
28pub struct SecondaryMap<K, V>
29where
30    K: EntityRef,
31    V: Clone,
32{
33    elems: Vec<V>,
34    default: V,
35    unused: PhantomData<K>,
36}
37
38/// Shared `SecondaryMap` implementation for all value types.
39impl<K, V> SecondaryMap<K, V>
40where
41    K: EntityRef,
42    V: Clone,
43{
44    /// Create a new empty map.
45    pub fn new() -> Self
46    where
47        V: Default,
48    {
49        Self {
50            elems: Vec::new(),
51            default: Default::default(),
52            unused: PhantomData,
53        }
54    }
55
56    /// Create a new, empty map with the specified capacity.
57    ///
58    /// The map will be able to hold exactly `capacity` elements without reallocating.
59    pub fn with_capacity(capacity: usize) -> Self
60    where
61        V: Default,
62    {
63        Self {
64            elems: Vec::with_capacity(capacity),
65            default: Default::default(),
66            unused: PhantomData,
67        }
68    }
69
70    /// Create a new empty map with a specified default value.
71    ///
72    /// This constructor does not require V to implement Default.
73    pub fn with_default(default: V) -> Self {
74        Self {
75            elems: Vec::new(),
76            default,
77            unused: PhantomData,
78        }
79    }
80
81    /// Returns the number of elements the map can hold without reallocating.
82    pub fn capacity(&self) -> usize {
83        self.elems.capacity()
84    }
85
86    /// Get the element at `k` if it exists.
87    #[inline(always)]
88    pub fn get(&self, k: K) -> Option<&V> {
89        self.elems.get(k.index())
90    }
91
92    /// Is this map completely empty?
93    #[inline(always)]
94    pub fn is_empty(&self) -> bool {
95        self.elems.is_empty()
96    }
97
98    /// Remove all entries from this map.
99    #[inline(always)]
100    pub fn clear(&mut self) {
101        self.elems.clear()
102    }
103
104    /// Iterate over all the keys and values in this map.
105    pub fn iter(&self) -> Iter<'_, K, V> {
106        Iter::new(self.elems.iter())
107    }
108
109    /// Iterate over all the keys and values in this map, mutable edition.
110    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
111        IterMut::new(self.elems.iter_mut())
112    }
113
114    /// Iterate over all the keys in this map.
115    pub fn keys(&self) -> Keys<K> {
116        Keys::with_len(self.elems.len())
117    }
118
119    /// Iterate over all the values in this map.
120    pub fn values(&self) -> slice::Iter<'_, V> {
121        self.elems.iter()
122    }
123
124    /// Iterate over all the values in this map, mutable edition.
125    pub fn values_mut(&mut self) -> slice::IterMut<'_, V> {
126        self.elems.iter_mut()
127    }
128
129    /// Resize the map to have `n` entries by adding default entries as needed.
130    pub fn resize(&mut self, n: usize) {
131        self.elems.resize(n, self.default.clone());
132    }
133
134    /// Slow path for `index_mut` which resizes the vector.
135    #[cold]
136    fn resize_for_index_mut(&mut self, i: usize) -> &mut V {
137        self.elems.resize(i + 1, self.default.clone());
138        &mut self.elems[i]
139    }
140}
141
142impl<K, V> Default for SecondaryMap<K, V>
143where
144    K: EntityRef,
145    V: Clone + Default,
146{
147    fn default() -> SecondaryMap<K, V> {
148        SecondaryMap::new()
149    }
150}
151
152impl<K, V> FromIterator<(K, V)> for SecondaryMap<K, V>
153where
154    K: EntityRef,
155    V: Clone + Default,
156{
157    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
158        let iter = iter.into_iter();
159        let (min, max) = iter.size_hint();
160        let cap = max.unwrap_or_else(|| 2 * min);
161        let mut map = Self::with_capacity(cap);
162        for (k, v) in iter {
163            map[k] = v;
164        }
165        map
166    }
167}
168
169/// Immutable indexing into an `SecondaryMap`.
170///
171/// All keys are permitted. Untouched entries have the default value.
172impl<K, V> Index<K> for SecondaryMap<K, V>
173where
174    K: EntityRef,
175    V: Clone,
176{
177    type Output = V;
178
179    #[inline(always)]
180    fn index(&self, k: K) -> &V {
181        self.elems.get(k.index()).unwrap_or(&self.default)
182    }
183}
184
185/// Mutable indexing into an `SecondaryMap`.
186///
187/// The map grows as needed to accommodate new keys.
188impl<K, V> IndexMut<K> for SecondaryMap<K, V>
189where
190    K: EntityRef,
191    V: Clone,
192{
193    #[inline(always)]
194    fn index_mut(&mut self, k: K) -> &mut V {
195        let i = k.index();
196        if i >= self.elems.len() {
197            return self.resize_for_index_mut(i);
198        }
199        &mut self.elems[i]
200    }
201}
202
203impl<K, V> PartialEq for SecondaryMap<K, V>
204where
205    K: EntityRef,
206    V: Clone + PartialEq,
207{
208    fn eq(&self, other: &Self) -> bool {
209        let min_size = min(self.elems.len(), other.elems.len());
210        self.default == other.default
211            && self.elems[..min_size] == other.elems[..min_size]
212            && self.elems[min_size..].iter().all(|e| *e == self.default)
213            && other.elems[min_size..].iter().all(|e| *e == other.default)
214    }
215}
216
217impl<K, V> Eq for SecondaryMap<K, V>
218where
219    K: EntityRef,
220    V: Clone + PartialEq + Eq,
221{
222}
223
224#[cfg(feature = "enable-serde")]
225impl<K, V> Serialize for SecondaryMap<K, V>
226where
227    K: EntityRef,
228    V: Clone + PartialEq + Serialize,
229{
230    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
231    where
232        S: Serializer,
233    {
234        // TODO: bincode encodes option as "byte for Some/None" and then optionally the content
235        // TODO: we can actually optimize it by encoding manually bitmask, then elements
236        let mut elems_cnt = self.elems.len();
237        while elems_cnt > 0 && self.elems[elems_cnt - 1] == self.default {
238            elems_cnt -= 1;
239        }
240        let mut seq = serializer.serialize_seq(Some(1 + elems_cnt))?;
241        seq.serialize_element(&Some(self.default.clone()))?;
242        for e in self.elems.iter().take(elems_cnt) {
243            let some_e = Some(e);
244            seq.serialize_element(if *e == self.default { &None } else { &some_e })?;
245        }
246        seq.end()
247    }
248}
249
250#[cfg(feature = "enable-serde")]
251impl<'de, K, V> Deserialize<'de> for SecondaryMap<K, V>
252where
253    K: EntityRef,
254    V: Clone + Deserialize<'de>,
255{
256    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
257    where
258        D: Deserializer<'de>,
259    {
260        use alloc::fmt;
261        struct SecondaryMapVisitor<K, V> {
262            unused: PhantomData<fn(K) -> V>,
263        }
264
265        impl<'de, K, V> Visitor<'de> for SecondaryMapVisitor<K, V>
266        where
267            K: EntityRef,
268            V: Clone + Deserialize<'de>,
269        {
270            type Value = SecondaryMap<K, V>;
271
272            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
273                formatter.write_str("struct SecondaryMap")
274            }
275
276            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
277            where
278                A: SeqAccess<'de>,
279            {
280                match seq.next_element()? {
281                    Some(Some(default_val)) => {
282                        let default_val: V = default_val; // compiler can't infer the type
283                        let mut m = SecondaryMap::with_default(default_val.clone());
284                        let mut idx = 0;
285                        while let Some(val) = seq.next_element()? {
286                            let val: Option<_> = val; // compiler can't infer the type
287                            m[K::new(idx)] = val.unwrap_or_else(|| default_val.clone());
288                            idx += 1;
289                        }
290                        Ok(m)
291                    }
292                    _ => Err(serde::de::Error::custom("Default value required")),
293                }
294            }
295        }
296
297        deserializer.deserialize_seq(SecondaryMapVisitor {
298            unused: PhantomData {},
299        })
300    }
301}
302
303impl<K: EntityRef + fmt::Debug, V: fmt::Debug + Clone> fmt::Debug for SecondaryMap<K, V> {
304    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
305        f.debug_struct("SecondaryMap")
306            .field("elems", &self.elems)
307            .field("default", &self.default)
308            .finish()
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315
316    // `EntityRef` impl for testing.
317    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
318    struct E(u32);
319
320    impl EntityRef for E {
321        fn new(i: usize) -> Self {
322            E(i as u32)
323        }
324        fn index(self) -> usize {
325            self.0 as usize
326        }
327    }
328
329    #[test]
330    fn basic() {
331        let r0 = E(0);
332        let r1 = E(1);
333        let r2 = E(2);
334        let mut m = SecondaryMap::new();
335
336        let v: Vec<E> = m.keys().collect();
337        assert_eq!(v, []);
338
339        m[r2] = 3;
340        m[r1] = 5;
341
342        assert_eq!(m[r1], 5);
343        assert_eq!(m[r2], 3);
344
345        let v: Vec<E> = m.keys().collect();
346        assert_eq!(v, [r0, r1, r2]);
347
348        let shared = &m;
349        assert_eq!(shared[r0], 0);
350        assert_eq!(shared[r1], 5);
351        assert_eq!(shared[r2], 3);
352    }
353}