cranelift_entity/
map.rs

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

cranelift_entity/
map.rs
