1use crate::EntityRef;
3use crate::boxed_slice::BoxedSlice;
4use crate::iter::{IntoIter, Iter, IterMut};
5use crate::keys::Keys;
6use alloc::boxed::Box;
7use alloc::vec::Vec;
8use core::fmt;
9use core::marker::PhantomData;
10use core::ops::{Index, IndexMut};
11use core::slice;
12#[cfg(feature = "enable-serde")]
13use serde_derive::{Deserialize, Serialize};
14
15#[derive(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_disjoint_mut<const N: usize>(
166 &mut self,
167 indices: [K; N],
168 ) -> Result<[&mut V; N], slice::GetDisjointMutError> {
169 self.elems.get_disjoint_mut(indices.map(|k| k.index()))
170 }
171
172 pub fn binary_search_values_by_key<'a, B, F>(&'a self, b: &B, f: F) -> Result<K, K>
184 where
185 F: FnMut(&'a V) -> B,
186 B: Ord,
187 {
188 self.elems
189 .binary_search_by_key(b, f)
190 .map(|i| K::new(i))
191 .map_err(|i| K::new(i))
192 }
193
194 pub fn get_raw_mut(&mut self, k: K) -> Option<*mut V> {
203 if k.index() < self.elems.len() {
204 unsafe { Some(self.elems.as_mut_ptr().add(k.index())) }
208 } else {
209 None
210 }
211 }
212}
213
214impl<K, V> Default for PrimaryMap<K, V>
215where
216 K: EntityRef,
217{
218 fn default() -> PrimaryMap<K, V> {
219 PrimaryMap::new()
220 }
221}
222
223impl<K, V> Index<K> for PrimaryMap<K, V>
226where
227 K: EntityRef,
228{
229 type Output = V;
230
231 fn index(&self, k: K) -> &V {
232 &self.elems[k.index()]
233 }
234}
235
236impl<K, V> IndexMut<K> for PrimaryMap<K, V>
238where
239 K: EntityRef,
240{
241 fn index_mut(&mut self, k: K) -> &mut V {
242 &mut self.elems[k.index()]
243 }
244}
245
246impl<K, V> IntoIterator for PrimaryMap<K, V>
247where
248 K: EntityRef,
249{
250 type Item = (K, V);
251 type IntoIter = IntoIter<K, V>;
252
253 fn into_iter(self) -> Self::IntoIter {
254 IntoIter::new(self.elems.into_iter())
255 }
256}
257
258impl<'a, K, V> IntoIterator for &'a PrimaryMap<K, V>
259where
260 K: EntityRef,
261{
262 type Item = (K, &'a V);
263 type IntoIter = Iter<'a, K, V>;
264
265 fn into_iter(self) -> Self::IntoIter {
266 Iter::new(self.elems.iter())
267 }
268}
269
270impl<'a, K, V> IntoIterator for &'a mut PrimaryMap<K, V>
271where
272 K: EntityRef,
273{
274 type Item = (K, &'a mut V);
275 type IntoIter = IterMut<'a, K, V>;
276
277 fn into_iter(self) -> Self::IntoIter {
278 IterMut::new(self.elems.iter_mut())
279 }
280}
281
282impl<K, V> FromIterator<V> for PrimaryMap<K, V>
283where
284 K: EntityRef,
285{
286 fn from_iter<T>(iter: T) -> Self
287 where
288 T: IntoIterator<Item = V>,
289 {
290 Self {
291 elems: Vec::from_iter(iter),
292 unused: PhantomData,
293 }
294 }
295}
296
297impl<K, V> From<Vec<V>> for PrimaryMap<K, V>
298where
299 K: EntityRef,
300{
301 fn from(elems: Vec<V>) -> Self {
302 Self {
303 elems,
304 unused: PhantomData,
305 }
306 }
307}
308
309impl<K, V> From<PrimaryMap<K, V>> for Vec<V>
310where
311 K: EntityRef,
312{
313 fn from(map: PrimaryMap<K, V>) -> Self {
314 map.elems
315 }
316}
317
318impl<K: EntityRef + fmt::Debug, V: fmt::Debug> fmt::Debug for PrimaryMap<K, V> {
319 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
320 let mut struct_ = f.debug_struct("PrimaryMap");
321 for (k, v) in self {
322 struct_.field(&alloc::format!("{k:?}"), v);
323 }
324 struct_.finish()
325 }
326}
327
328#[cfg(test)]
329mod tests {
330 use super::*;
331
332 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
334 struct E(u32);
335
336 impl EntityRef for E {
337 fn new(i: usize) -> Self {
338 E(i as u32)
339 }
340 fn index(self) -> usize {
341 self.0 as usize
342 }
343 }
344
345 #[test]
346 fn basic() {
347 let r0 = E(0);
348 let r1 = E(1);
349 let m = PrimaryMap::<E, isize>::new();
350
351 let v: Vec<E> = m.keys().collect();
352 assert_eq!(v, []);
353
354 assert!(!m.is_valid(r0));
355 assert!(!m.is_valid(r1));
356 }
357
358 #[test]
359 fn push() {
360 let mut m = PrimaryMap::new();
361 let k0: E = m.push(12);
362 let k1 = m.push(33);
363
364 assert_eq!(m[k0], 12);
365 assert_eq!(m[k1], 33);
366
367 let v: Vec<E> = m.keys().collect();
368 assert_eq!(v, [k0, k1]);
369 }
370
371 #[test]
372 fn iter() {
373 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
374 m.push(12);
375 m.push(33);
376
377 let mut i = 0;
378 for (key, value) in &m {
379 assert_eq!(key.index(), i);
380 match i {
381 0 => assert_eq!(*value, 12),
382 1 => assert_eq!(*value, 33),
383 _ => panic!(),
384 }
385 i += 1;
386 }
387 i = 0;
388 for (key_mut, value_mut) in m.iter_mut() {
389 assert_eq!(key_mut.index(), i);
390 match i {
391 0 => assert_eq!(*value_mut, 12),
392 1 => assert_eq!(*value_mut, 33),
393 _ => panic!(),
394 }
395 i += 1;
396 }
397 }
398
399 #[test]
400 fn iter_rev() {
401 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
402 m.push(12);
403 m.push(33);
404
405 let mut i = 2;
406 for (key, value) in m.iter().rev() {
407 i -= 1;
408 assert_eq!(key.index(), i);
409 match i {
410 0 => assert_eq!(*value, 12),
411 1 => assert_eq!(*value, 33),
412 _ => panic!(),
413 }
414 }
415
416 i = 2;
417 for (key, value) in m.iter_mut().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 #[test]
428 fn keys() {
429 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
430 m.push(12);
431 m.push(33);
432
433 let mut i = 0;
434 for key in m.keys() {
435 assert_eq!(key.index(), i);
436 i += 1;
437 }
438 }
439
440 #[test]
441 fn keys_rev() {
442 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
443 m.push(12);
444 m.push(33);
445
446 let mut i = 2;
447 for key in m.keys().rev() {
448 i -= 1;
449 assert_eq!(key.index(), i);
450 }
451 }
452
453 #[test]
454 fn values() {
455 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
456 m.push(12);
457 m.push(33);
458
459 let mut i = 0;
460 for value in m.values() {
461 match i {
462 0 => assert_eq!(*value, 12),
463 1 => assert_eq!(*value, 33),
464 _ => panic!(),
465 }
466 i += 1;
467 }
468 i = 0;
469 for value_mut in m.values_mut() {
470 match i {
471 0 => assert_eq!(*value_mut, 12),
472 1 => assert_eq!(*value_mut, 33),
473 _ => panic!(),
474 }
475 i += 1;
476 }
477 }
478
479 #[test]
480 fn values_rev() {
481 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
482 m.push(12);
483 m.push(33);
484
485 let mut i = 2;
486 for value in m.values().rev() {
487 i -= 1;
488 match i {
489 0 => assert_eq!(*value, 12),
490 1 => assert_eq!(*value, 33),
491 _ => panic!(),
492 }
493 }
494 i = 2;
495 for value_mut in m.values_mut().rev() {
496 i -= 1;
497 match i {
498 0 => assert_eq!(*value_mut, 12),
499 1 => assert_eq!(*value_mut, 33),
500 _ => panic!(),
501 }
502 }
503 }
504
505 #[test]
506 fn from_iter() {
507 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
508 m.push(12);
509 m.push(33);
510
511 let n = m.values().collect::<PrimaryMap<E, _>>();
512 assert!(m.len() == n.len());
513 for (me, ne) in m.values().zip(n.values()) {
514 assert!(*me == **ne);
515 }
516 }
517
518 #[test]
519 fn from_vec() {
520 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
521 m.push(12);
522 m.push(33);
523
524 let n = PrimaryMap::<E, &usize>::from(m.values().collect::<Vec<_>>());
525 assert!(m.len() == n.len());
526 for (me, ne) in m.values().zip(n.values()) {
527 assert!(*me == **ne);
528 }
529 }
530
531 #[test]
532 fn get_many_mut() {
533 let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
534 let _0 = m.push(0);
535 let _1 = m.push(1);
536 let _2 = m.push(2);
537
538 assert_eq!([&mut 0, &mut 2], m.get_disjoint_mut([_0, _2]).unwrap());
539 assert_eq!(
540 m.get_disjoint_mut([_0, _0]),
541 Err(slice::GetDisjointMutError::OverlappingIndices)
542 );
543 assert_eq!(
544 m.get_disjoint_mut([E(4)]),
545 Err(slice::GetDisjointMutError::IndexOutOfBounds)
546 );
547 }
548}