1use crate::boxed_slice::BoxedSlice;
3use crate::iter::{IntoIter, Iter, IterMut};
4use crate::keys::Keys;
5use crate::EntityRef;
6use alloc::boxed::Box;
7use alloc::vec::Vec;
8use core::marker::PhantomData;
9use core::mem;
10use core::ops::{Index, IndexMut};
11use core::slice;
12#[cfg(feature = "enable-serde")]
13use serde_derive::{Deserialize, Serialize};
14
15#[derive(Debug, Clone, Hash, PartialEq, Eq)]
31#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
32pub struct PrimaryMap<K, V>
33where
34 K: EntityRef,
35{
36 elems: Vec<V>,
37 unused: PhantomData<K>,
38}
39
40impl<K, V> PrimaryMap<K, V>
41where
42 K: EntityRef,
43{
44 pub fn new() -> Self {
46 Self {
47 elems: Vec::new(),
48 unused: PhantomData,
49 }
50 }
51
52 pub fn with_capacity(capacity: usize) -> Self {
54 Self {
55 elems: Vec::with_capacity(capacity),
56 unused: PhantomData,
57 }
58 }
59
60 pub fn is_valid(&self, k: K) -> bool {
62 k.index() < self.elems.len()
63 }
64
65 pub fn get(&self, k: K) -> Option<&V> {
67 self.elems.get(k.index())
68 }
69
70 pub fn get_mut(&mut self, k: K) -> Option<&mut V> {
72 self.elems.get_mut(k.index())
73 }
74
75 pub fn is_empty(&self) -> bool {
77 self.elems.is_empty()
78 }
79
80 pub fn len(&self) -> usize {
82 self.elems.len()
83 }
84
85 pub fn keys(&self) -> Keys<K> {
87 Keys::with_len(self.elems.len())
88 }
89
90 pub fn values(&self) -> slice::Iter<V> {
92 self.elems.iter()
93 }
94
95 pub fn values_mut(&mut self) -> slice::IterMut<V> {
97 self.elems.iter_mut()
98 }
99
100 pub fn iter(&self) -> Iter<K, V> {
102 Iter::new(self.elems.iter())
103 }
104
105 pub fn iter_mut(&mut self) -> IterMut<K, V> {
107 IterMut::new(self.elems.iter_mut())
108 }
109
110 pub fn clear(&mut self) {
112 self.elems.clear()
113 }
114
115 pub fn next_key(&self) -> K {
117 K::new(self.elems.len())
118 }
119
120 pub fn push(&mut self, v: V) -> K {
122 let k = self.next_key();
123 self.elems.push(v);
124 k
125 }
126
127 pub fn last(&self) -> Option<(K, &V)> {
129 let len = self.elems.len();
130 let last = self.elems.last()?;
131 Some((K::new(len - 1), last))
132 }
133
134 pub fn last_mut(&mut self) -> Option<(K, &mut V)> {
136 let len = self.elems.len();
137 let last = self.elems.last_mut()?;
138 Some((K::new(len - 1), last))
139 }
140
141 pub fn reserve(&mut self, additional: usize) {
143 self.elems.reserve(additional)
144 }
145
146 pub fn reserve_exact(&mut self, additional: usize) {
148 self.elems.reserve_exact(additional)
149 }
150
151 pub fn shrink_to_fit(&mut self) {
153 self.elems.shrink_to_fit()
154 }
155
156 pub fn into_boxed_slice(self) -> BoxedSlice<K, V> {
158 unsafe { BoxedSlice::<K, V>::from_raw(Box::<[V]>::into_raw(self.elems.into_boxed_slice())) }
159 }
160
161 pub fn get_many_mut<const N: usize>(
169 &mut self,
170 indices: [K; N],
171 ) -> Result<[&mut V; N], GetManyMutError<K>> {
172 for (i, &idx) in indices.iter().enumerate() {
173 if idx.index() >= self.len() {
174 return Err(GetManyMutError::DoesNotExist(idx));
175 }
176 for &idx2 in &indices[..i] {
177 if idx == idx2 {
178 return Err(GetManyMutError::MultipleOf(idx));
179 }
180 }
181 }
182
183 let slice: *mut V = self.elems.as_mut_ptr();
184 let mut arr: mem::MaybeUninit<[&mut V; N]> = mem::MaybeUninit::uninit();
185 let arr_ptr = arr.as_mut_ptr();
186
187 unsafe {
188 for i in 0..N {
189 let idx = *indices.get_unchecked(i);
190 *(*arr_ptr).get_unchecked_mut(i) = &mut *slice.add(idx.index());
191 }
192 Ok(arr.assume_init())
193 }
194 }
195
196 pub fn binary_search_values_by_key<'a, B, F>(&'a self, b: &B, f: F) -> Result<K, K>
208 where
209 F: FnMut(&'a V) -> B,
210 B: Ord,
211 {
212 self.elems
213 .binary_search_by_key(b, f)
214 .map(|i| K::new(i))
215 .map_err(|i| K::new(i))
216 }
217
218 pub fn get_raw_mut(&mut self, k: K) -> Option<*mut V> {
227 if k.index() < self.elems.len() {
228 unsafe { Some(self.elems.as_mut_ptr().add(k.index())) }
232 } else {
233 None
234 }
235 }
236}
237
238#[derive(Debug, PartialEq, Eq, Clone, Copy)]
239pub enum GetManyMutError<K> {
240 DoesNotExist(K),
241 MultipleOf(K),
242}
243
244impl<K, V> Default for PrimaryMap<K, V>
245where
246 K: EntityRef,
247{
248 fn default() -> PrimaryMap<K, V> {
249 PrimaryMap::new()
250 }
251}
252
253impl<K, V> Index<K> for PrimaryMap<K, V>
256where
257 K: EntityRef,
258{
259 type Output = V;
260
261 fn index(&self, k: K) -> &V {
262 &self.elems[k.index()]
263 }
264}
265
266impl<K, V> IndexMut<K> for PrimaryMap<K, V>
268where
269 K: EntityRef,
270{
271 fn index_mut(&mut self, k: K) -> &mut V {
272 &mut self.elems[k.index()]
273 }
274}
275
276impl<K, V> IntoIterator for PrimaryMap<K, V>
277where
278 K: EntityRef,
279{
280 type Item = (K, V);
281 type IntoIter = IntoIter<K, V>;
282
283 fn into_iter(self) -> Self::IntoIter {
284 IntoIter::new(self.elems.into_iter())
285 }
286}
287
288impl<'a, K, V> IntoIterator for &'a PrimaryMap<K, V>
289where
290 K: EntityRef,
291{
292 type Item = (K, &'a V);
293 type IntoIter = Iter<'a, K, V>;
294
295 fn into_iter(self) -> Self::IntoIter {
296 Iter::new(self.elems.iter())
297 }
298}
299
300impl<'a, K, V> IntoIterator for &'a mut PrimaryMap<K, V>
301where
302 K: EntityRef,
303{
304 type Item = (K, &'a mut V);
305 type IntoIter = IterMut<'a, K, V>;
306
307 fn into_iter(self) -> Self::IntoIter {
308 IterMut::new(self.elems.iter_mut())
309 }
310}
311
312impl<K, V> FromIterator<V> for PrimaryMap<K, V>
313where
314 K: EntityRef,
315{
316 fn from_iter<T>(iter: T) -> Self
317 where
318 T: IntoIterator<Item = V>,
319 {
320 Self {
321 elems: Vec::from_iter(iter),
322 unused: PhantomData,
323 }
324 }
325}
326
327impl<K, V> From<Vec<V>> for PrimaryMap<K, V>
328where
329 K: EntityRef,
330{
331 fn from(elems: Vec<V>) -> Self {
332 Self {
333 elems,
334 unused: PhantomData,
335 }
336 }
337}
338
339#[cfg(test)]
340mod tests {
341 use super::*;
342
343 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
345 struct E(u32);
346
347 impl EntityRef for E {
348 fn new(i: usize) -> Self {
349 E(i as u32)
350 }
351 fn index(self) -> usize {
352 self.0 as usize
353 }
354 }
355
356 #[test]
357 fn basic() {
358 let r0 = E(0);
359 let r1 = E(1);
360 let m = PrimaryMap::<E, isize>::new();
361
362 let v: Vec<E> = m.keys().collect();
363 assert_eq!(v, []);
364
365 assert!(!m.is_valid(r0));
366 assert!(!m.is_valid(r1));
367 }
368
369 #[test]
370 fn push() {
371 let mut m = PrimaryMap::new();
372 let k0: E = m.push(12);
373 let k1 = m.push(33);
374
375 assert_eq!(m[k0], 12);
376 assert_eq!(m[k1], 33);
377
378 let v: Vec<E> = m.keys().collect();
379 assert_eq!(v, [k0, k1]);
380 }
381
382 #[test]
383 fn iter() {
384 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
385 m.push(12);
386 m.push(33);
387
388 let mut i = 0;
389 for (key, value) in &m {
390 assert_eq!(key.index(), i);
391 match i {
392 0 => assert_eq!(*value, 12),
393 1 => assert_eq!(*value, 33),
394 _ => panic!(),
395 }
396 i += 1;
397 }
398 i = 0;
399 for (key_mut, value_mut) in m.iter_mut() {
400 assert_eq!(key_mut.index(), i);
401 match i {
402 0 => assert_eq!(*value_mut, 12),
403 1 => assert_eq!(*value_mut, 33),
404 _ => panic!(),
405 }
406 i += 1;
407 }
408 }
409
410 #[test]
411 fn iter_rev() {
412 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
413 m.push(12);
414 m.push(33);
415
416 let mut i = 2;
417 for (key, value) in m.iter().rev() {
418 i -= 1;
419 assert_eq!(key.index(), i);
420 match i {
421 0 => assert_eq!(*value, 12),
422 1 => assert_eq!(*value, 33),
423 _ => panic!(),
424 }
425 }
426
427 i = 2;
428 for (key, value) in m.iter_mut().rev() {
429 i -= 1;
430 assert_eq!(key.index(), i);
431 match i {
432 0 => assert_eq!(*value, 12),
433 1 => assert_eq!(*value, 33),
434 _ => panic!(),
435 }
436 }
437 }
438 #[test]
439 fn keys() {
440 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
441 m.push(12);
442 m.push(33);
443
444 let mut i = 0;
445 for key in m.keys() {
446 assert_eq!(key.index(), i);
447 i += 1;
448 }
449 }
450
451 #[test]
452 fn keys_rev() {
453 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
454 m.push(12);
455 m.push(33);
456
457 let mut i = 2;
458 for key in m.keys().rev() {
459 i -= 1;
460 assert_eq!(key.index(), i);
461 }
462 }
463
464 #[test]
465 fn values() {
466 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
467 m.push(12);
468 m.push(33);
469
470 let mut i = 0;
471 for value in m.values() {
472 match i {
473 0 => assert_eq!(*value, 12),
474 1 => assert_eq!(*value, 33),
475 _ => panic!(),
476 }
477 i += 1;
478 }
479 i = 0;
480 for value_mut in m.values_mut() {
481 match i {
482 0 => assert_eq!(*value_mut, 12),
483 1 => assert_eq!(*value_mut, 33),
484 _ => panic!(),
485 }
486 i += 1;
487 }
488 }
489
490 #[test]
491 fn values_rev() {
492 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
493 m.push(12);
494 m.push(33);
495
496 let mut i = 2;
497 for value in m.values().rev() {
498 i -= 1;
499 match i {
500 0 => assert_eq!(*value, 12),
501 1 => assert_eq!(*value, 33),
502 _ => panic!(),
503 }
504 }
505 i = 2;
506 for value_mut in m.values_mut().rev() {
507 i -= 1;
508 match i {
509 0 => assert_eq!(*value_mut, 12),
510 1 => assert_eq!(*value_mut, 33),
511 _ => panic!(),
512 }
513 }
514 }
515
516 #[test]
517 fn from_iter() {
518 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
519 m.push(12);
520 m.push(33);
521
522 let n = m.values().collect::<PrimaryMap<E, _>>();
523 assert!(m.len() == n.len());
524 for (me, ne) in m.values().zip(n.values()) {
525 assert!(*me == **ne);
526 }
527 }
528
529 #[test]
530 fn from_vec() {
531 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
532 m.push(12);
533 m.push(33);
534
535 let n = PrimaryMap::<E, &usize>::from(m.values().collect::<Vec<_>>());
536 assert!(m.len() == n.len());
537 for (me, ne) in m.values().zip(n.values()) {
538 assert!(*me == **ne);
539 }
540 }
541
542 #[test]
543 fn get_many_mut() {
544 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
545 let _0 = m.push(0);
546 let _1 = m.push(1);
547 let _2 = m.push(2);
548
549 assert_eq!([&mut 0, &mut 2], m.get_many_mut([_0, _2]).unwrap());
550 assert_eq!(
551 m.get_many_mut([_0, _0]),
552 Err(GetManyMutError::MultipleOf(_0))
553 );
554 assert_eq!(
555 m.get_many_mut([E(4)]),
556 Err(GetManyMutError::DoesNotExist(E(4)))
557 );
558 }
559}