1use 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#[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
38impl<K, V> SecondaryMap<K, V>
40where
41 K: EntityRef,
42 V: Clone,
43{
44 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 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 pub fn with_default(default: V) -> Self {
74 Self {
75 elems: Vec::new(),
76 default,
77 unused: PhantomData,
78 }
79 }
80
81 pub fn capacity(&self) -> usize {
83 self.elems.capacity()
84 }
85
86 #[inline(always)]
88 pub fn get(&self, k: K) -> Option<&V> {
89 self.elems.get(k.index())
90 }
91
92 #[inline(always)]
94 pub fn is_empty(&self) -> bool {
95 self.elems.is_empty()
96 }
97
98 #[inline(always)]
100 pub fn clear(&mut self) {
101 self.elems.clear()
102 }
103
104 pub fn iter(&self) -> Iter<'_, K, V> {
106 Iter::new(self.elems.iter())
107 }
108
109 pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
111 IterMut::new(self.elems.iter_mut())
112 }
113
114 pub fn keys(&self) -> Keys<K> {
116 Keys::with_len(self.elems.len())
117 }
118
119 pub fn values(&self) -> slice::Iter<'_, V> {
121 self.elems.iter()
122 }
123
124 pub fn values_mut(&mut self) -> slice::IterMut<'_, V> {
126 self.elems.iter_mut()
127 }
128
129 pub fn resize(&mut self, n: usize) {
131 self.elems.resize(n, self.default.clone());
132 }
133
134 #[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
169impl<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
185impl<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 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; 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; 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 #[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}