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
152/// Immutable indexing into an `SecondaryMap`.
153///
154/// All keys are permitted. Untouched entries have the default value.
155impl<K, V> Index<K> for SecondaryMap<K, V>
156where
157    K: EntityRef,
158    V: Clone,
159{
160    type Output = V;
161
162    #[inline(always)]
163    fn index(&self, k: K) -> &V {
164        self.elems.get(k.index()).unwrap_or(&self.default)
165    }
166}
167
168/// Mutable indexing into an `SecondaryMap`.
169///
170/// The map grows as needed to accommodate new keys.
171impl<K, V> IndexMut<K> for SecondaryMap<K, V>
172where
173    K: EntityRef,
174    V: Clone,
175{
176    #[inline(always)]
177    fn index_mut(&mut self, k: K) -> &mut V {
178        let i = k.index();
179        if i >= self.elems.len() {
180            return self.resize_for_index_mut(i);
181        }
182        &mut self.elems[i]
183    }
184}
185
186impl<K, V> PartialEq for SecondaryMap<K, V>
187where
188    K: EntityRef,
189    V: Clone + PartialEq,
190{
191    fn eq(&self, other: &Self) -> bool {
192        let min_size = min(self.elems.len(), other.elems.len());
193        self.default == other.default
194            && self.elems[..min_size] == other.elems[..min_size]
195            && self.elems[min_size..].iter().all(|e| *e == self.default)
196            && other.elems[min_size..].iter().all(|e| *e == other.default)
197    }
198}
199
200impl<K, V> Eq for SecondaryMap<K, V>
201where
202    K: EntityRef,
203    V: Clone + PartialEq + Eq,
204{
205}
206
207#[cfg(feature = "enable-serde")]
208impl<K, V> Serialize for SecondaryMap<K, V>
209where
210    K: EntityRef,
211    V: Clone + PartialEq + Serialize,
212{
213    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
214    where
215        S: Serializer,
216    {
217        // TODO: bincode encodes option as "byte for Some/None" and then optionally the content
218        // TODO: we can actually optimize it by encoding manually bitmask, then elements
219        let mut elems_cnt = self.elems.len();
220        while elems_cnt > 0 && self.elems[elems_cnt - 1] == self.default {
221            elems_cnt -= 1;
222        }
223        let mut seq = serializer.serialize_seq(Some(1 + elems_cnt))?;
224        seq.serialize_element(&Some(self.default.clone()))?;
225        for e in self.elems.iter().take(elems_cnt) {
226            let some_e = Some(e);
227            seq.serialize_element(if *e == self.default { &None } else { &some_e })?;
228        }
229        seq.end()
230    }
231}
232
233#[cfg(feature = "enable-serde")]
234impl<'de, K, V> Deserialize<'de> for SecondaryMap<K, V>
235where
236    K: EntityRef,
237    V: Clone + Deserialize<'de>,
238{
239    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
240    where
241        D: Deserializer<'de>,
242    {
243        use alloc::fmt;
244        struct SecondaryMapVisitor<K, V> {
245            unused: PhantomData<fn(K) -> V>,
246        }
247
248        impl<'de, K, V> Visitor<'de> for SecondaryMapVisitor<K, V>
249        where
250            K: EntityRef,
251            V: Clone + Deserialize<'de>,
252        {
253            type Value = SecondaryMap<K, V>;
254
255            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
256                formatter.write_str("struct SecondaryMap")
257            }
258
259            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
260            where
261                A: SeqAccess<'de>,
262            {
263                match seq.next_element()? {
264                    Some(Some(default_val)) => {
265                        let default_val: V = default_val; // compiler can't infer the type
266                        let mut m = SecondaryMap::with_default(default_val.clone());
267                        let mut idx = 0;
268                        while let Some(val) = seq.next_element()? {
269                            let val: Option<_> = val; // compiler can't infer the type
270                            m[K::new(idx)] = val.unwrap_or_else(|| default_val.clone());
271                            idx += 1;
272                        }
273                        Ok(m)
274                    }
275                    _ => Err(serde::de::Error::custom("Default value required")),
276                }
277            }
278        }
279
280        deserializer.deserialize_seq(SecondaryMapVisitor {
281            unused: PhantomData {},
282        })
283    }
284}
285
286impl<K: EntityRef + fmt::Debug, V: fmt::Debug + Clone> fmt::Debug for SecondaryMap<K, V> {
287    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
288        f.debug_struct("SecondaryMap")
289            .field("elems", &self.elems)
290            .field("default", &self.default)
291            .finish()
292    }
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298
299    // `EntityRef` impl for testing.
300    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
301    struct E(u32);
302
303    impl EntityRef for E {
304        fn new(i: usize) -> Self {
305            E(i as u32)
306        }
307        fn index(self) -> usize {
308            self.0 as usize
309        }
310    }
311
312    #[test]
313    fn basic() {
314        let r0 = E(0);
315        let r1 = E(1);
316        let r2 = E(2);
317        let mut m = SecondaryMap::new();
318
319        let v: Vec<E> = m.keys().collect();
320        assert_eq!(v, []);
321
322        m[r2] = 3;
323        m[r1] = 5;
324
325        assert_eq!(m[r1], 5);
326        assert_eq!(m[r2], 3);
327
328        let v: Vec<E> = m.keys().collect();
329        assert_eq!(v, [r0, r1, r2]);
330
331        let shared = &m;
332        assert_eq!(shared[r0], 0);
333        assert_eq!(shared[r1], 5);
334        assert_eq!(shared[r2], 3);
335    }
336}