near_vm_types/entity/
secondary_map.rs

1// This file contains code from external sources.
2// Attributions: https://github.com/wasmerio/wasmer/blob/2.3.0/ATTRIBUTIONS.md
3
4//! Densely numbered entity references as mapping keys.
5
6use crate::entity::EntityRef;
7use crate::entity::iter::{Iter, IterMut};
8use crate::entity::keys::Keys;
9use crate::lib::std::cmp::min;
10use crate::lib::std::marker::PhantomData;
11use crate::lib::std::ops::{Index, IndexMut};
12use crate::lib::std::slice;
13use crate::lib::std::vec::Vec;
14use rkyv::Archive;
15
16/// A mapping `K -> V` for densely indexed entity references.
17///
18/// The `SecondaryMap` data structure uses the dense index space to implement a map with a vector.
19/// Unlike `PrimaryMap`, an `SecondaryMap` can't be used to allocate entity references. It is used
20/// to associate secondary information with entities.
21///
22/// The map does not track if an entry for a key has been inserted or not. Instead it behaves as if
23/// all keys have a default entry from the beginning.
24#[derive(Debug, Clone, rkyv::Serialize, rkyv::Deserialize, rkyv::Archive)]
25pub struct SecondaryMap<K, V>
26where
27    K: EntityRef,
28    V: Clone,
29{
30    pub(crate) elems: Vec<V>,
31    pub(crate) default: V,
32    pub(crate) unused: PhantomData<K>,
33}
34
35/// Shared `SecondaryMap` implementation for all value types.
36impl<K, V> SecondaryMap<K, V>
37where
38    K: EntityRef,
39    V: Clone,
40{
41    /// Create a new empty map.
42    pub fn new() -> Self
43    where
44        V: Default,
45    {
46        Self { elems: Vec::new(), default: Default::default(), unused: PhantomData }
47    }
48
49    /// Create a new, empty map with the specified capacity.
50    ///
51    /// The map will be able to hold exactly `capacity` elements without reallocating.
52    pub fn with_capacity(capacity: usize) -> Self
53    where
54        V: Default,
55    {
56        Self {
57            elems: Vec::with_capacity(capacity),
58            default: Default::default(),
59            unused: PhantomData,
60        }
61    }
62
63    /// Create a new empty map with a specified default value.
64    ///
65    /// This constructor does not require V to implement Default.
66    pub fn with_default(default: V) -> Self {
67        Self { elems: Vec::new(), default, unused: PhantomData }
68    }
69
70    /// Returns the number of elements the map can hold without reallocating.
71    pub fn capacity(&self) -> usize {
72        self.elems.capacity()
73    }
74
75    /// Get the element at `k` if it exists.
76    #[inline(always)]
77    pub fn get(&self, k: K) -> Option<&V> {
78        self.elems.get(k.index())
79    }
80
81    /// Is this map completely empty?
82    #[inline(always)]
83    pub fn is_empty(&self) -> bool {
84        self.elems.is_empty()
85    }
86
87    /// Remove all entries from this map.
88    #[inline(always)]
89    pub fn clear(&mut self) {
90        self.elems.clear()
91    }
92
93    /// Iterate over all the keys and values in this map.
94    pub fn iter(&self) -> Iter<K, V> {
95        Iter::new(self.elems.iter())
96    }
97
98    /// Iterate over all the keys and values in this map, mutable edition.
99    pub fn iter_mut(&mut self) -> IterMut<K, V> {
100        IterMut::new(self.elems.iter_mut())
101    }
102
103    /// Iterate over all the keys in this map.
104    pub fn keys(&self) -> Keys<K> {
105        Keys::with_len(self.elems.len())
106    }
107
108    /// Iterate over all the values in this map.
109    pub fn values(&self) -> slice::Iter<V> {
110        self.elems.iter()
111    }
112
113    /// Iterate over all the values in this map, mutable edition.
114    pub fn values_mut(&mut self) -> slice::IterMut<V> {
115        self.elems.iter_mut()
116    }
117
118    /// Resize the map to have `n` entries by adding default entries as needed.
119    pub fn resize(&mut self, n: usize) {
120        self.elems.resize(n, self.default.clone());
121    }
122}
123
124impl<K, V> Default for SecondaryMap<K, V>
125where
126    K: EntityRef,
127    V: Clone + Default,
128{
129    fn default() -> Self {
130        Self::new()
131    }
132}
133
134/// Immutable indexing into an `SecondaryMap`.
135///
136/// All keys are permitted. Untouched entries have the default value.
137impl<K, V> Index<K> for SecondaryMap<K, V>
138where
139    K: EntityRef,
140    V: Clone,
141{
142    type Output = V;
143
144    #[inline(always)]
145    fn index(&self, k: K) -> &V {
146        self.elems.get(k.index()).unwrap_or(&self.default)
147    }
148}
149
150/// Immutable indexing into an `SecondaryMap`.
151///
152/// All keys are permitted. Untouched entries have the default value.
153impl<K, V> Index<&K::Archived> for ArchivedSecondaryMap<K, V>
154where
155    K: EntityRef + Archive,
156    K::Archived: EntityRef,
157    V: Archive + Clone,
158{
159    type Output = <V as rkyv::Archive>::Archived;
160
161    fn index(&self, k: &K::Archived) -> &Self::Output {
162        &self.elems.get(k.index()).unwrap_or(&self.default)
163    }
164}
165
166/// Mutable indexing into an `SecondaryMap`.
167///
168/// The map grows as needed to accommodate new keys.
169impl<K, V> IndexMut<K> for SecondaryMap<K, V>
170where
171    K: EntityRef,
172    V: Clone,
173{
174    #[inline(always)]
175    fn index_mut(&mut self, k: K) -> &mut V {
176        let i = k.index();
177        if i >= self.elems.len() {
178            self.elems.resize(i + 1, self.default.clone());
179        }
180        &mut self.elems[i]
181    }
182}
183
184impl<K, V> PartialEq for SecondaryMap<K, V>
185where
186    K: EntityRef,
187    V: Clone + PartialEq,
188{
189    fn eq(&self, other: &Self) -> bool {
190        let min_size = min(self.elems.len(), other.elems.len());
191        self.default == other.default
192            && self.elems[..min_size] == other.elems[..min_size]
193            && self.elems[min_size..].iter().all(|e| *e == self.default)
194            && other.elems[min_size..].iter().all(|e| *e == other.default)
195    }
196}
197
198impl<K, V> Eq for SecondaryMap<K, V>
199where
200    K: EntityRef,
201    V: Clone + PartialEq + Eq,
202{
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    // `EntityRef` impl for testing.
210    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
211    struct E(u32);
212
213    impl EntityRef for E {
214        fn new(i: usize) -> Self {
215            Self(i as u32)
216        }
217        fn index(self) -> usize {
218            self.0 as usize
219        }
220    }
221
222    #[test]
223    fn basic() {
224        let r0 = E(0);
225        let r1 = E(1);
226        let r2 = E(2);
227        let mut m = SecondaryMap::new();
228
229        let v: Vec<E> = m.keys().collect();
230        assert_eq!(v, []);
231
232        m[r2] = 3;
233        m[r1] = 5;
234
235        assert_eq!(m[r1], 5);
236        assert_eq!(m[r2], 3);
237
238        let v: Vec<E> = m.keys().collect();
239        assert_eq!(v, [r0, r1, r2]);
240
241        let shared = &m;
242        assert_eq!(shared[r0], 0);
243        assert_eq!(shared[r1], 5);
244        assert_eq!(shared[r2], 3);
245    }
246}
near_vm_types/entity/secondary_map.rs

near_vm_types/entity/
secondary_map.rs
