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> 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
168impl<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 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; 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; 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 #[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}