slotmap_fork_otter/basic.rs
1// Needed because assigning to non-Copy union is unsafe in stable but not in nightly.
2#![allow(unused_unsafe)]
3
4//! Contains the slot map implementation.
5
6#[cfg(all(nightly, any(doc, feature = "unstable")))]
7use alloc::collections::TryReserveError;
8use alloc::vec::Vec;
9use core::fmt;
10use core::iter::{Enumerate, FusedIterator};
11use core::marker::PhantomData;
12#[allow(unused_imports)] // MaybeUninit is only used on nightly at the moment.
13use core::mem::{ManuallyDrop, MaybeUninit};
14use core::ops::{Index, IndexMut};
15
16use crate::{DefaultKey, Key, KeyData};
17
18// Storage inside a slot or metadata for the freelist when vacant.
19union SlotUnion<T> {
20 value: ManuallyDrop<T>,
21 next_free: u32,
22}
23
24// A slot, which represents storage for a value and a current version.
25// Can be occupied or vacant.
26struct Slot<T> {
27 u: SlotUnion<T>,
28 version: u32, // Even = vacant, odd = occupied.
29}
30
31// Safe API to read a slot.
32enum SlotContent<'a, T: 'a> {
33 Occupied(&'a T),
34 Vacant(&'a u32),
35}
36
37enum SlotContentMut<'a, T: 'a> {
38 OccupiedMut(&'a mut T),
39 VacantMut(&'a mut u32),
40}
41
42use self::SlotContent::{Occupied, Vacant};
43use self::SlotContentMut::{OccupiedMut, VacantMut};
44
45impl<T> Slot<T> {
46 // Is this slot occupied?
47 #[inline(always)]
48 pub fn occupied(&self) -> bool {
49 self.version % 2 > 0
50 }
51
52 pub fn get(&self) -> SlotContent<T> {
53 unsafe {
54 if self.occupied() {
55 Occupied(&*self.u.value)
56 } else {
57 Vacant(&self.u.next_free)
58 }
59 }
60 }
61
62 pub fn get_mut(&mut self) -> SlotContentMut<T> {
63 unsafe {
64 if self.occupied() {
65 OccupiedMut(&mut *self.u.value)
66 } else {
67 VacantMut(&mut self.u.next_free)
68 }
69 }
70 }
71}
72
73impl<T> Drop for Slot<T> {
74 fn drop(&mut self) {
75 if core::mem::needs_drop::<T>() && self.occupied() {
76 // This is safe because we checked that we're occupied.
77 unsafe {
78 ManuallyDrop::drop(&mut self.u.value);
79 }
80 }
81 }
82}
83
84impl<T: Clone> Clone for Slot<T> {
85 fn clone(&self) -> Self {
86 Self {
87 u: match self.get() {
88 Occupied(value) => SlotUnion {
89 value: ManuallyDrop::new(value.clone()),
90 },
91 Vacant(&next_free) => SlotUnion { next_free },
92 },
93 version: self.version,
94 }
95 }
96}
97
98impl<T: fmt::Debug> fmt::Debug for Slot<T> {
99 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
100 let mut builder = fmt.debug_struct("Slot");
101 builder.field("version", &self.version);
102 match self.get() {
103 Occupied(value) => builder.field("value", value).finish(),
104 Vacant(next_free) => builder.field("next_free", next_free).finish(),
105 }
106 }
107}
108
109/// Slot map, storage with stable unique keys.
110///
111/// See [crate documentation](crate) for more details.
112#[derive(Debug, Clone)]
113pub struct SlotMap<K: Key, V> {
114 slots: Vec<Slot<V>>,
115 free_head: u32,
116 num_elems: u32,
117 _k: PhantomData<fn(K) -> K>,
118}
119
120impl<V> SlotMap<DefaultKey, V> {
121 /// Constructs a new, empty [`SlotMap`].
122 ///
123 /// # Examples
124 ///
125 /// ```
126 /// # use slotmap::*;
127 /// let mut sm: SlotMap<_, i32> = SlotMap::new();
128 /// ```
129 pub fn new() -> Self {
130 Self::with_capacity_and_key(0)
131 }
132
133 /// Creates an empty [`SlotMap`] with the given capacity.
134 ///
135 /// The slot map will not reallocate until it holds at least `capacity`
136 /// elements.
137 ///
138 /// # Examples
139 ///
140 /// ```
141 /// # use slotmap::*;
142 /// let mut sm: SlotMap<_, i32> = SlotMap::with_capacity(10);
143 /// ```
144 pub fn with_capacity(capacity: usize) -> Self {
145 Self::with_capacity_and_key(capacity)
146 }
147}
148
149impl<K: Key, V> SlotMap<K, V> {
150 /// Constructs a new, empty [`SlotMap`] with a custom key type.
151 ///
152 /// # Examples
153 ///
154 /// ```
155 /// # use slotmap::*;
156 /// new_key_type! {
157 /// struct PositionKey;
158 /// }
159 /// let mut positions: SlotMap<PositionKey, i32> = SlotMap::with_key();
160 /// ```
161 pub fn with_key() -> Self {
162 Self::with_capacity_and_key(0)
163 }
164
165 /// Creates an empty [`SlotMap`] with the given capacity and a custom key
166 /// type.
167 ///
168 /// The slot map will not reallocate until it holds at least `capacity`
169 /// elements.
170 ///
171 /// # Examples
172 ///
173 /// ```
174 /// # use slotmap::*;
175 /// new_key_type! {
176 /// struct MessageKey;
177 /// }
178 /// let mut messages = SlotMap::with_capacity_and_key(3);
179 /// let welcome: MessageKey = messages.insert("Welcome");
180 /// let good_day = messages.insert("Good day");
181 /// let hello = messages.insert("Hello");
182 /// ```
183 pub fn with_capacity_and_key(capacity: usize) -> Self {
184 // Create slots with a sentinel at index 0.
185 // We don't actually use the sentinel for anything currently, but
186 // HopSlotMap does, and if we want keys to remain valid through
187 // conversion we have to have one as well.
188 let mut slots = Vec::with_capacity(capacity + 1);
189 slots.push(Slot {
190 u: SlotUnion { next_free: 0 },
191 version: 0,
192 });
193
194 Self {
195 slots,
196 free_head: 1,
197 num_elems: 0,
198 _k: PhantomData,
199 }
200 }
201
202 /// Returns the number of elements in the slot map.
203 ///
204 /// # Examples
205 ///
206 /// ```
207 /// # use slotmap::*;
208 /// let mut sm = SlotMap::with_capacity(10);
209 /// sm.insert("len() counts actual elements, not capacity");
210 /// let key = sm.insert("removed elements don't count either");
211 /// sm.remove(key);
212 /// assert_eq!(sm.len(), 1);
213 /// ```
214 pub fn len(&self) -> usize {
215 self.num_elems as usize
216 }
217
218 /// Returns if the slot map is empty.
219 ///
220 /// # Examples
221 ///
222 /// ```
223 /// # use slotmap::*;
224 /// let mut sm = SlotMap::new();
225 /// let key = sm.insert("dummy");
226 /// assert_eq!(sm.is_empty(), false);
227 /// sm.remove(key);
228 /// assert_eq!(sm.is_empty(), true);
229 /// ```
230 pub fn is_empty(&self) -> bool {
231 self.num_elems == 0
232 }
233
234 /// Returns the number of elements the [`SlotMap`] can hold without
235 /// reallocating.
236 ///
237 /// # Examples
238 ///
239 /// ```
240 /// # use slotmap::*;
241 /// let sm: SlotMap<_, f64> = SlotMap::with_capacity(10);
242 /// assert_eq!(sm.capacity(), 10);
243 /// ```
244 pub fn capacity(&self) -> usize {
245 // One slot is reserved for the sentinel.
246 self.slots.capacity() - 1
247 }
248
249 /// Reserves capacity for at least `additional` more elements to be inserted
250 /// in the [`SlotMap`]. The collection may reserve more space to avoid
251 /// frequent reallocations.
252 ///
253 /// # Panics
254 ///
255 /// Panics if the new allocation size overflows [`usize`].
256 ///
257 /// # Examples
258 ///
259 /// ```
260 /// # use slotmap::*;
261 /// let mut sm = SlotMap::new();
262 /// sm.insert("foo");
263 /// sm.reserve(32);
264 /// assert!(sm.capacity() >= 33);
265 /// ```
266 pub fn reserve(&mut self, additional: usize) {
267 // One slot is reserved for the sentinel.
268 let needed = (self.len() + additional).saturating_sub(self.slots.len() - 1);
269 self.slots.reserve(needed);
270 }
271
272 /// Tries to reserve capacity for at least `additional` more elements to be
273 /// inserted in the [`SlotMap`]. The collection may reserve more space to
274 /// avoid frequent reallocations.
275 ///
276 /// # Examples
277 ///
278 /// ```
279 /// # use slotmap::*;
280 /// let mut sm = SlotMap::new();
281 /// sm.insert("foo");
282 /// sm.try_reserve(32).unwrap();
283 /// assert!(sm.capacity() >= 33);
284 /// ```
285 #[cfg(all(nightly, any(doc, feature = "unstable")))]
286 #[cfg_attr(all(nightly, doc), doc(cfg(feature = "unstable")))]
287 pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
288 // One slot is reserved for the sentinel.
289 let needed = (self.len() + additional).saturating_sub(self.slots.len() - 1);
290 self.slots.try_reserve(needed)
291 }
292
293 /// Returns [`true`] if the slot map contains `key`.
294 ///
295 /// # Examples
296 ///
297 /// ```
298 /// # use slotmap::*;
299 /// let mut sm = SlotMap::new();
300 /// let key = sm.insert(42);
301 /// assert_eq!(sm.contains_key(key), true);
302 /// sm.remove(key);
303 /// assert_eq!(sm.contains_key(key), false);
304 /// ```
305 pub fn contains_key(&self, key: K) -> bool {
306 let kd = key.data();
307 self.slots
308 .get(kd.idx as usize)
309 .map_or(false, |slot| slot.version == kd.version.get())
310 }
311
312 /// Inserts a value into the slot map. Returns a unique key that can be used
313 /// to access this value.
314 ///
315 /// # Panics
316 ///
317 /// Panics if the number of elements in the slot map equals
318 /// 2<sup>32</sup> - 2.
319 ///
320 /// # Examples
321 ///
322 /// ```
323 /// # use slotmap::*;
324 /// let mut sm = SlotMap::new();
325 /// let key = sm.insert(42);
326 /// assert_eq!(sm[key], 42);
327 /// ```
328 pub fn insert(&mut self, value: V) -> K {
329 self.insert_with_key(|_| value)
330 }
331
332 /// Inserts a value given by `f` into the slot map. The key where the
333 /// value will be stored is passed into `f`. This is useful to store values
334 /// that contain their own key.
335 ///
336 /// # Panics
337 ///
338 /// Panics if the number of elements in the slot map equals
339 /// 2<sup>32</sup> - 2.
340 ///
341 /// # Examples
342 ///
343 /// ```
344 /// # use slotmap::*;
345 /// let mut sm = SlotMap::new();
346 /// let key = sm.insert_with_key(|k| (k, 20));
347 /// assert_eq!(sm[key], (key, 20));
348 /// ```
349 pub fn insert_with_key<F>(&mut self, f: F) -> K
350 where
351 F: FnOnce(K) -> V,
352 {
353 // In case f panics, we don't make any changes until we have the value.
354 let new_num_elems = self.num_elems + 1;
355 if new_num_elems == core::u32::MAX {
356 panic!("SlotMap number of elements overflow");
357 }
358
359 if let Some(slot) = self.slots.get_mut(self.free_head as usize) {
360 let occupied_version = slot.version | 1;
361 let kd = KeyData::new(self.free_head, occupied_version);
362
363 // Get value first in case f panics.
364 let value = f(kd.into());
365
366 // Update.
367 unsafe {
368 self.free_head = slot.u.next_free;
369 slot.u.value = ManuallyDrop::new(value);
370 slot.version = occupied_version;
371 }
372 self.num_elems = new_num_elems;
373 return kd.into();
374 }
375
376 let version = 1;
377 let kd = KeyData::new(self.slots.len() as u32, version);
378
379 // Create new slot before adjusting freelist in case f or the allocation panics.
380 self.slots.push(Slot {
381 u: SlotUnion {
382 value: ManuallyDrop::new(f(kd.into())),
383 },
384 version,
385 });
386
387 self.free_head = kd.idx + 1;
388 self.num_elems = new_num_elems;
389 kd.into()
390 }
391
392 // Helper function to remove a value from a slot. Safe iff the slot is
393 // occupied. Returns the value removed.
394 #[inline(always)]
395 unsafe fn remove_from_slot(&mut self, idx: usize) -> V {
396 // Remove value from slot before overwriting union.
397 let slot = self.slots.get_unchecked_mut(idx);
398 let value = ManuallyDrop::take(&mut slot.u.value);
399
400 // Maintain freelist.
401 slot.u.next_free = self.free_head;
402 self.free_head = idx as u32;
403 self.num_elems -= 1;
404 slot.version = slot.version.wrapping_add(1);
405
406 value
407 }
408
409 /// Removes a key from the slot map, returning the value at the key if the
410 /// key was not previously removed.
411 ///
412 /// # Examples
413 ///
414 /// ```
415 /// # use slotmap::*;
416 /// let mut sm = SlotMap::new();
417 /// let key = sm.insert(42);
418 /// assert_eq!(sm.remove(key), Some(42));
419 /// assert_eq!(sm.remove(key), None);
420 /// ```
421 pub fn remove(&mut self, key: K) -> Option<V> {
422 let kd = key.data();
423 if self.contains_key(key) {
424 // This is safe because we know that the slot is occupied.
425 Some(unsafe { self.remove_from_slot(kd.idx as usize) })
426 } else {
427 None
428 }
429 }
430
431 /// Retains only the elements specified by the predicate.
432 ///
433 /// In other words, remove all key-value pairs `(k, v)` such that
434 /// `f(k, &mut v)` returns false. This method invalidates any removed keys.
435 ///
436 /// This function must iterate over all slots, empty or not. In the face of
437 /// many deleted elements it can be inefficient.
438 ///
439 /// # Examples
440 ///
441 /// ```
442 /// # use slotmap::*;
443 /// let mut sm = SlotMap::new();
444 ///
445 /// let k1 = sm.insert(0);
446 /// let k2 = sm.insert(1);
447 /// let k3 = sm.insert(2);
448 ///
449 /// sm.retain(|key, val| key == k1 || *val == 1);
450 ///
451 /// assert!(sm.contains_key(k1));
452 /// assert!(sm.contains_key(k2));
453 /// assert!(!sm.contains_key(k3));
454 ///
455 /// assert_eq!(2, sm.len());
456 /// ```
457 pub fn retain<F>(&mut self, mut f: F)
458 where
459 F: FnMut(K, &mut V) -> bool,
460 {
461 for i in 1..self.slots.len() {
462 // This is safe because removing elements does not shrink slots.
463 let slot = unsafe { self.slots.get_unchecked_mut(i) };
464 let version = slot.version;
465
466 let should_remove = if let OccupiedMut(value) = slot.get_mut() {
467 let key = KeyData::new(i as u32, version).into();
468 !f(key, value)
469 } else {
470 false
471 };
472
473 if should_remove {
474 // This is safe because we know that the slot was occupied.
475 unsafe { self.remove_from_slot(i) };
476 }
477 }
478 }
479
480 /// Clears the slot map. Keeps the allocated memory for reuse.
481 ///
482 /// This function must iterate over all slots, empty or not. In the face of
483 /// many deleted elements it can be inefficient.
484 ///
485 /// # Examples
486 ///
487 /// ```
488 /// # use slotmap::*;
489 /// let mut sm = SlotMap::new();
490 /// for i in 0..10 {
491 /// sm.insert(i);
492 /// }
493 /// assert_eq!(sm.len(), 10);
494 /// sm.clear();
495 /// assert_eq!(sm.len(), 0);
496 /// ```
497 pub fn clear(&mut self) {
498 self.drain();
499 }
500
501 /// Clears the slot map, returning all key-value pairs in arbitrary order as
502 /// an iterator. Keeps the allocated memory for reuse.
503 ///
504 /// When the iterator is dropped all elements in the slot map are removed,
505 /// even if the iterator was not fully consumed. If the iterator is not
506 /// dropped (using e.g. [`std::mem::forget`]), only the elements that were
507 /// iterated over are removed.
508 ///
509 /// This function must iterate over all slots, empty or not. In the face of
510 /// many deleted elements it can be inefficient.
511 ///
512 /// # Examples
513 ///
514 /// ```
515 /// # use slotmap::*;
516 /// let mut sm = SlotMap::new();
517 /// let k = sm.insert(0);
518 /// let v: Vec<_> = sm.drain().collect();
519 /// assert_eq!(sm.len(), 0);
520 /// assert_eq!(v, vec![(k, 0)]);
521 /// ```
522 pub fn drain(&mut self) -> Drain<K, V> {
523 Drain {
524 cur: 1,
525 num_left: self.len(),
526 sm: self,
527 }
528 }
529
530 /// Returns a reference to the value corresponding to the key.
531 ///
532 /// # Examples
533 ///
534 /// ```
535 /// # use slotmap::*;
536 /// let mut sm = SlotMap::new();
537 /// let key = sm.insert("bar");
538 /// assert_eq!(sm.get(key), Some(&"bar"));
539 /// sm.remove(key);
540 /// assert_eq!(sm.get(key), None);
541 /// ```
542 pub fn get(&self, key: K) -> Option<&V> {
543 let kd = key.data();
544 self.slots
545 .get(kd.idx as usize)
546 .filter(|slot| slot.version == kd.version.get())
547 .map(|slot| unsafe { &*slot.u.value })
548 }
549
550 /// Returns a reference to the value corresponding to the key without
551 /// version or bounds checking.
552 ///
553 /// # Safety
554 ///
555 /// This should only be used if `contains_key(key)` is true. Otherwise it is
556 /// potentially unsafe.
557 ///
558 /// # Examples
559 ///
560 /// ```
561 /// # use slotmap::*;
562 /// let mut sm = SlotMap::new();
563 /// let key = sm.insert("bar");
564 /// assert_eq!(unsafe { sm.get_unchecked(key) }, &"bar");
565 /// sm.remove(key);
566 /// // sm.get_unchecked(key) is now dangerous!
567 /// ```
568 pub unsafe fn get_unchecked(&self, key: K) -> &V {
569 debug_assert!(self.contains_key(key));
570 &self.slots.get_unchecked(key.data().idx as usize).u.value
571 }
572
573 /// Returns a mutable reference to the value corresponding to the key.
574 ///
575 /// # Examples
576 ///
577 /// ```
578 /// # use slotmap::*;
579 /// let mut sm = SlotMap::new();
580 /// let key = sm.insert(3.5);
581 /// if let Some(x) = sm.get_mut(key) {
582 /// *x += 3.0;
583 /// }
584 /// assert_eq!(sm[key], 6.5);
585 /// ```
586 pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
587 let kd = key.data();
588 self.slots
589 .get_mut(kd.idx as usize)
590 .filter(|slot| slot.version == kd.version.get())
591 .map(|slot| unsafe { &mut *slot.u.value })
592 }
593
594 /// Returns a mutable reference to the value corresponding to the key
595 /// without version or bounds checking.
596 ///
597 /// # Safety
598 ///
599 /// This should only be used if `contains_key(key)` is true. Otherwise it is
600 /// potentially unsafe.
601 ///
602 /// # Examples
603 ///
604 /// ```
605 /// # use slotmap::*;
606 /// let mut sm = SlotMap::new();
607 /// let key = sm.insert("foo");
608 /// unsafe { *sm.get_unchecked_mut(key) = "bar" };
609 /// assert_eq!(sm[key], "bar");
610 /// sm.remove(key);
611 /// // sm.get_unchecked_mut(key) is now dangerous!
612 /// ```
613 pub unsafe fn get_unchecked_mut(&mut self, key: K) -> &mut V {
614 debug_assert!(self.contains_key(key));
615 &mut self
616 .slots
617 .get_unchecked_mut(key.data().idx as usize)
618 .u
619 .value
620 }
621
622 /// Returns mutable references to the values corresponding to the given
623 /// keys. All keys must be valid and disjoint, otherwise None is returned.
624 ///
625 /// # Examples
626 ///
627 /// ```
628 /// # use slotmap::*;
629 /// let mut sm = SlotMap::new();
630 /// let ka = sm.insert("butter");
631 /// let kb = sm.insert("apples");
632 /// let kc = sm.insert("charlie");
633 /// sm.remove(kc); // Make key c invalid.
634 /// assert_eq!(sm.get_disjoint_mut([ka, kb, kc]), None); // Has invalid key.
635 /// assert_eq!(sm.get_disjoint_mut([ka, ka]), None); // Not disjoint.
636 /// let [a, b] = sm.get_disjoint_mut([ka, kb]).unwrap();
637 /// std::mem::swap(a, b);
638 /// assert_eq!(sm[ka], "apples");
639 /// assert_eq!(sm[kb], "butter");
640 /// ```
641 #[cfg(all(nightly, any(doc, feature = "unstable")))]
642 #[cfg_attr(all(nightly, doc), doc(cfg(feature = "unstable")))]
643 pub fn get_disjoint_mut<const N: usize>(&mut self, keys: [K; N]) -> Option<[&mut V; N]> {
644 // Create an uninitialized array of `MaybeUninit`. The `assume_init` is
645 // safe because the type we are claiming to have initialized here is a
646 // bunch of `MaybeUninit`s, which do not require initialization.
647 let mut ptrs: [MaybeUninit<*mut V>; N] = unsafe { MaybeUninit::uninit().assume_init() };
648
649 let mut i = 0;
650 while i < N {
651 let kd = keys[i].data();
652 if !self.contains_key(kd.into()) {
653 break;
654 }
655
656 // This key is valid, and thus the slot is occupied. Temporarily
657 // mark it as unoccupied so duplicate keys would show up as invalid.
658 // This gives us a linear time disjointness check.
659 unsafe {
660 let slot = self.slots.get_unchecked_mut(kd.idx as usize);
661 slot.version ^= 1;
662 ptrs[i] = MaybeUninit::new(&mut *slot.u.value);
663 }
664 i += 1;
665 }
666
667 // Undo temporary unoccupied markings.
668 for k in &keys[..i] {
669 let idx = k.data().idx as usize;
670 unsafe {
671 self.slots.get_unchecked_mut(idx).version ^= 1;
672 }
673 }
674
675 if i == N {
676 // All were valid and disjoint.
677 Some(unsafe { core::mem::transmute_copy::<_, [&mut V; N]>(&ptrs) })
678 } else {
679 None
680 }
681 }
682
683 /// Returns mutable references to the values corresponding to the given
684 /// keys. All keys must be valid and disjoint.
685 ///
686 /// # Safety
687 ///
688 /// This should only be used if `contains_key(key)` is true for every given
689 /// key and no two keys are equal. Otherwise it is potentially unsafe.
690 ///
691 /// # Examples
692 ///
693 /// ```
694 /// # use slotmap::*;
695 /// let mut sm = SlotMap::new();
696 /// let ka = sm.insert("butter");
697 /// let kb = sm.insert("apples");
698 /// let [a, b] = unsafe { sm.get_disjoint_unchecked_mut([ka, kb]) };
699 /// std::mem::swap(a, b);
700 /// assert_eq!(sm[ka], "apples");
701 /// assert_eq!(sm[kb], "butter");
702 /// ```
703 #[cfg(all(nightly, any(doc, feature = "unstable")))]
704 #[cfg_attr(all(nightly, doc), doc(cfg(feature = "unstable")))]
705 pub unsafe fn get_disjoint_unchecked_mut<const N: usize>(
706 &mut self,
707 keys: [K; N],
708 ) -> [&mut V; N] {
709 // Safe, see get_disjoint_mut.
710 let mut ptrs: [MaybeUninit<*mut V>; N] = MaybeUninit::uninit().assume_init();
711 for i in 0..N {
712 ptrs[i] = MaybeUninit::new(self.get_unchecked_mut(keys[i]));
713 }
714 core::mem::transmute_copy::<_, [&mut V; N]>(&ptrs)
715 }
716
717 /// An iterator visiting all key-value pairs in arbitrary order. The
718 /// iterator element type is `(K, &'a V)`.
719 ///
720 /// This function must iterate over all slots, empty or not. In the face of
721 /// many deleted elements it can be inefficient.
722 ///
723 /// # Examples
724 ///
725 /// ```
726 /// # use slotmap::*;
727 /// let mut sm = SlotMap::new();
728 /// let k0 = sm.insert(0);
729 /// let k1 = sm.insert(1);
730 /// let k2 = sm.insert(2);
731 ///
732 /// for (k, v) in sm.iter() {
733 /// println!("key: {:?}, val: {}", k, v);
734 /// }
735 /// ```
736 pub fn iter(&self) -> Iter<K, V> {
737 let mut it = self.slots.iter().enumerate();
738 it.next(); // Skip sentinel.
739 Iter {
740 slots: it,
741 num_left: self.len(),
742 _k: PhantomData,
743 }
744 }
745
746 /// An iterator visiting all key-value pairs in arbitrary order, with
747 /// mutable references to the values. The iterator element type is
748 /// `(K, &'a mut V)`.
749 ///
750 /// This function must iterate over all slots, empty or not. In the face of
751 /// many deleted elements it can be inefficient.
752 ///
753 /// # Examples
754 ///
755 /// ```
756 /// # use slotmap::*;
757 /// let mut sm = SlotMap::new();
758 /// let k0 = sm.insert(10);
759 /// let k1 = sm.insert(20);
760 /// let k2 = sm.insert(30);
761 ///
762 /// for (k, v) in sm.iter_mut() {
763 /// if k != k1 {
764 /// *v *= -1;
765 /// }
766 /// }
767 ///
768 /// assert_eq!(sm[k0], -10);
769 /// assert_eq!(sm[k1], 20);
770 /// assert_eq!(sm[k2], -30);
771 /// ```
772 pub fn iter_mut(&mut self) -> IterMut<K, V> {
773 let len = self.len();
774 let mut it = self.slots.iter_mut().enumerate();
775 it.next(); // Skip sentinel.
776 IterMut {
777 num_left: len,
778 slots: it,
779 _k: PhantomData,
780 }
781 }
782
783 /// An iterator visiting all keys in arbitrary order. The iterator element
784 /// type is `K`.
785 ///
786 /// This function must iterate over all slots, empty or not. In the face of
787 /// many deleted elements it can be inefficient.
788 ///
789 /// # Examples
790 ///
791 /// ```
792 /// # use slotmap::*;
793 /// # use std::collections::HashSet;
794 /// let mut sm = SlotMap::new();
795 /// let k0 = sm.insert(10);
796 /// let k1 = sm.insert(20);
797 /// let k2 = sm.insert(30);
798 /// let keys: HashSet<_> = sm.keys().collect();
799 /// let check: HashSet<_> = vec![k0, k1, k2].into_iter().collect();
800 /// assert_eq!(keys, check);
801 /// ```
802 pub fn keys(&self) -> Keys<K, V> {
803 Keys { inner: self.iter() }
804 }
805
806 /// An iterator visiting all values in arbitrary order. The iterator element
807 /// type is `&'a V`.
808 ///
809 /// This function must iterate over all slots, empty or not. In the face of
810 /// many deleted elements it can be inefficient.
811 ///
812 /// # Examples
813 ///
814 /// ```
815 /// # use slotmap::*;
816 /// # use std::collections::HashSet;
817 /// let mut sm = SlotMap::new();
818 /// let k0 = sm.insert(10);
819 /// let k1 = sm.insert(20);
820 /// let k2 = sm.insert(30);
821 /// let values: HashSet<_> = sm.values().collect();
822 /// let check: HashSet<_> = vec![&10, &20, &30].into_iter().collect();
823 /// assert_eq!(values, check);
824 /// ```
825 pub fn values(&self) -> Values<K, V> {
826 Values { inner: self.iter() }
827 }
828
829 /// An iterator visiting all values mutably in arbitrary order. The iterator
830 /// element type is `&'a mut V`.
831 ///
832 /// This function must iterate over all slots, empty or not. In the face of
833 /// many deleted elements it can be inefficient.
834 ///
835 /// # Examples
836 ///
837 /// ```
838 /// # use slotmap::*;
839 /// # use std::collections::HashSet;
840 /// let mut sm = SlotMap::new();
841 /// sm.insert(1);
842 /// sm.insert(2);
843 /// sm.insert(3);
844 /// sm.values_mut().for_each(|n| { *n *= 3 });
845 /// let values: HashSet<_> = sm.into_iter().map(|(_k, v)| v).collect();
846 /// let check: HashSet<_> = vec![3, 6, 9].into_iter().collect();
847 /// assert_eq!(values, check);
848 /// ```
849 pub fn values_mut(&mut self) -> ValuesMut<K, V> {
850 ValuesMut {
851 inner: self.iter_mut(),
852 }
853 }
854}
855
856impl<K: Key, V> Default for SlotMap<K, V> {
857 fn default() -> Self {
858 Self::with_key()
859 }
860}
861
862impl<K: Key, V> Index<K> for SlotMap<K, V> {
863 type Output = V;
864
865 fn index(&self, key: K) -> &V {
866 match self.get(key) {
867 Some(r) => r,
868 None => panic!("invalid SlotMap key used"),
869 }
870 }
871}
872
873impl<K: Key, V> IndexMut<K> for SlotMap<K, V> {
874 fn index_mut(&mut self, key: K) -> &mut V {
875 match self.get_mut(key) {
876 Some(r) => r,
877 None => panic!("invalid SlotMap key used"),
878 }
879 }
880}
881
882// Iterators.
883/// A draining iterator for [`SlotMap`].
884///
885/// This iterator is created by [`SlotMap::drain`].
886#[derive(Debug)]
887pub struct Drain<'a, K: 'a + Key, V: 'a> {
888 num_left: usize,
889 sm: &'a mut SlotMap<K, V>,
890 cur: usize,
891}
892
893/// An iterator that moves key-value pairs out of a [`SlotMap`].
894///
895/// This iterator is created by calling the `into_iter` method on [`SlotMap`],
896/// provided by the [`IntoIterator`] trait.
897#[derive(Debug, Clone)]
898pub struct IntoIter<K: Key, V> {
899 num_left: usize,
900 slots: Enumerate<alloc::vec::IntoIter<Slot<V>>>,
901 _k: PhantomData<fn(K) -> K>,
902}
903
904/// An iterator over the key-value pairs in a [`SlotMap`].
905///
906/// This iterator is created by [`SlotMap::iter`].
907#[derive(Debug, Clone)]
908pub struct Iter<'a, K: 'a + Key, V: 'a> {
909 num_left: usize,
910 slots: Enumerate<core::slice::Iter<'a, Slot<V>>>,
911 _k: PhantomData<fn(K) -> K>,
912}
913
914/// A mutable iterator over the key-value pairs in a [`SlotMap`].
915///
916/// This iterator is created by [`SlotMap::iter_mut`].
917#[derive(Debug)]
918pub struct IterMut<'a, K: 'a + Key, V: 'a> {
919 num_left: usize,
920 slots: Enumerate<core::slice::IterMut<'a, Slot<V>>>,
921 _k: PhantomData<fn(K) -> K>,
922}
923
924/// An iterator over the keys in a [`SlotMap`].
925///
926/// This iterator is created by [`SlotMap::keys`].
927#[derive(Debug, Clone)]
928pub struct Keys<'a, K: 'a + Key, V: 'a> {
929 inner: Iter<'a, K, V>,
930}
931
932/// An iterator over the values in a [`SlotMap`].
933///
934/// This iterator is created by [`SlotMap::values`].
935#[derive(Debug, Clone)]
936pub struct Values<'a, K: 'a + Key, V: 'a> {
937 inner: Iter<'a, K, V>,
938}
939
940/// A mutable iterator over the values in a [`SlotMap`].
941///
942/// This iterator is created by [`SlotMap::values_mut`].
943#[derive(Debug)]
944pub struct ValuesMut<'a, K: 'a + Key, V: 'a> {
945 inner: IterMut<'a, K, V>,
946}
947
948impl<'a, K: Key, V> Iterator for Drain<'a, K, V> {
949 type Item = (K, V);
950
951 fn next(&mut self) -> Option<(K, V)> {
952 let len = self.sm.slots.len();
953 while self.cur < len {
954 let idx = self.cur;
955 self.cur += 1;
956
957 // This is safe because removing doesn't shrink slots.
958 unsafe {
959 let slot = self.sm.slots.get_unchecked(idx);
960 if slot.occupied() {
961 let kd = KeyData::new(idx as u32, slot.version);
962 self.num_left -= 1;
963 return Some((kd.into(), self.sm.remove_from_slot(idx)));
964 }
965 }
966 }
967
968 None
969 }
970
971 fn size_hint(&self) -> (usize, Option<usize>) {
972 (self.num_left, Some(self.num_left))
973 }
974}
975
976impl<'a, K: Key, V> Drop for Drain<'a, K, V> {
977 fn drop(&mut self) {
978 self.for_each(|_drop| {});
979 }
980}
981
982impl<K: Key, V> Iterator for IntoIter<K, V> {
983 type Item = (K, V);
984
985 fn next(&mut self) -> Option<(K, V)> {
986 while let Some((idx, mut slot)) = self.slots.next() {
987 if slot.occupied() {
988 let kd = KeyData::new(idx as u32, slot.version);
989
990 // Prevent dropping after extracting the value.
991 slot.version = 0;
992
993 // This is safe because we know the slot was occupied.
994 let value = unsafe { ManuallyDrop::take(&mut slot.u.value) };
995
996 self.num_left -= 1;
997 return Some((kd.into(), value));
998 }
999 }
1000
1001 None
1002 }
1003
1004 fn size_hint(&self) -> (usize, Option<usize>) {
1005 (self.num_left, Some(self.num_left))
1006 }
1007}
1008
1009impl<'a, K: Key, V> Iterator for Iter<'a, K, V> {
1010 type Item = (K, &'a V);
1011
1012 fn next(&mut self) -> Option<(K, &'a V)> {
1013 while let Some((idx, slot)) = self.slots.next() {
1014 if let Occupied(value) = slot.get() {
1015 let kd = KeyData::new(idx as u32, slot.version);
1016 self.num_left -= 1;
1017 return Some((kd.into(), value));
1018 }
1019 }
1020
1021 None
1022 }
1023
1024 fn size_hint(&self) -> (usize, Option<usize>) {
1025 (self.num_left, Some(self.num_left))
1026 }
1027}
1028
1029impl<'a, K: Key, V> Iterator for IterMut<'a, K, V> {
1030 type Item = (K, &'a mut V);
1031
1032 fn next(&mut self) -> Option<(K, &'a mut V)> {
1033 while let Some((idx, slot)) = self.slots.next() {
1034 let version = slot.version;
1035 if let OccupiedMut(value) = slot.get_mut() {
1036 let kd = KeyData::new(idx as u32, version);
1037 self.num_left -= 1;
1038 return Some((kd.into(), value));
1039 }
1040 }
1041
1042 None
1043 }
1044
1045 fn size_hint(&self) -> (usize, Option<usize>) {
1046 (self.num_left, Some(self.num_left))
1047 }
1048}
1049
1050impl<'a, K: Key, V> Iterator for Keys<'a, K, V> {
1051 type Item = K;
1052
1053 fn next(&mut self) -> Option<K> {
1054 self.inner.next().map(|(key, _)| key)
1055 }
1056
1057 fn size_hint(&self) -> (usize, Option<usize>) {
1058 self.inner.size_hint()
1059 }
1060}
1061
1062impl<'a, K: Key, V> Iterator for Values<'a, K, V> {
1063 type Item = &'a V;
1064
1065 fn next(&mut self) -> Option<&'a V> {
1066 self.inner.next().map(|(_, value)| value)
1067 }
1068
1069 fn size_hint(&self) -> (usize, Option<usize>) {
1070 self.inner.size_hint()
1071 }
1072}
1073
1074impl<'a, K: Key, V> Iterator for ValuesMut<'a, K, V> {
1075 type Item = &'a mut V;
1076
1077 fn next(&mut self) -> Option<&'a mut V> {
1078 self.inner.next().map(|(_, value)| value)
1079 }
1080
1081 fn size_hint(&self) -> (usize, Option<usize>) {
1082 self.inner.size_hint()
1083 }
1084}
1085
1086impl<'a, K: Key, V> IntoIterator for &'a SlotMap<K, V> {
1087 type Item = (K, &'a V);
1088 type IntoIter = Iter<'a, K, V>;
1089
1090 fn into_iter(self) -> Self::IntoIter {
1091 self.iter()
1092 }
1093}
1094
1095impl<'a, K: Key, V> IntoIterator for &'a mut SlotMap<K, V> {
1096 type Item = (K, &'a mut V);
1097 type IntoIter = IterMut<'a, K, V>;
1098
1099 fn into_iter(self) -> Self::IntoIter {
1100 self.iter_mut()
1101 }
1102}
1103
1104impl<K: Key, V> IntoIterator for SlotMap<K, V> {
1105 type Item = (K, V);
1106 type IntoIter = IntoIter<K, V>;
1107
1108 fn into_iter(self) -> Self::IntoIter {
1109 let len = self.len();
1110 let mut it = self.slots.into_iter().enumerate();
1111 it.next(); // Skip sentinel.
1112 IntoIter {
1113 num_left: len,
1114 slots: it,
1115 _k: PhantomData,
1116 }
1117 }
1118}
1119
1120impl<'a, K: Key, V> FusedIterator for Iter<'a, K, V> {}
1121impl<'a, K: Key, V> FusedIterator for IterMut<'a, K, V> {}
1122impl<'a, K: Key, V> FusedIterator for Keys<'a, K, V> {}
1123impl<'a, K: Key, V> FusedIterator for Values<'a, K, V> {}
1124impl<'a, K: Key, V> FusedIterator for ValuesMut<'a, K, V> {}
1125impl<'a, K: Key, V> FusedIterator for Drain<'a, K, V> {}
1126impl<K: Key, V> FusedIterator for IntoIter<K, V> {}
1127
1128impl<'a, K: Key, V> ExactSizeIterator for Iter<'a, K, V> {}
1129impl<'a, K: Key, V> ExactSizeIterator for IterMut<'a, K, V> {}
1130impl<'a, K: Key, V> ExactSizeIterator for Keys<'a, K, V> {}
1131impl<'a, K: Key, V> ExactSizeIterator for Values<'a, K, V> {}
1132impl<'a, K: Key, V> ExactSizeIterator for ValuesMut<'a, K, V> {}
1133impl<'a, K: Key, V> ExactSizeIterator for Drain<'a, K, V> {}
1134impl<K: Key, V> ExactSizeIterator for IntoIter<K, V> {}
1135
1136// Serialization with serde.
1137#[cfg(feature = "serde")]
1138mod serialize {
1139 use super::*;
1140 use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
1141
1142 #[derive(Serialize, Deserialize)]
1143 struct SerdeSlot<T> {
1144 value: Option<T>,
1145 version: u32,
1146 }
1147
1148 impl<T: Serialize> Serialize for Slot<T> {
1149 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1150 where
1151 S: Serializer,
1152 {
1153 let serde_slot = SerdeSlot {
1154 version: self.version,
1155 value: match self.get() {
1156 Occupied(value) => Some(value),
1157 Vacant(_) => None,
1158 },
1159 };
1160 serde_slot.serialize(serializer)
1161 }
1162 }
1163
1164 impl<'de, T> Deserialize<'de> for Slot<T>
1165 where
1166 T: Deserialize<'de>,
1167 {
1168 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1169 where
1170 D: Deserializer<'de>,
1171 {
1172 let serde_slot: SerdeSlot<T> = Deserialize::deserialize(deserializer)?;
1173 let occupied = serde_slot.version % 2 == 1;
1174 if occupied ^ serde_slot.value.is_some() {
1175 return Err(de::Error::custom(&"inconsistent occupation in Slot"));
1176 }
1177
1178 Ok(Self {
1179 u: match serde_slot.value {
1180 Some(value) => SlotUnion {
1181 value: ManuallyDrop::new(value),
1182 },
1183 None => SlotUnion { next_free: 0 },
1184 },
1185 version: serde_slot.version,
1186 })
1187 }
1188 }
1189
1190 impl<K: Key, V: Serialize> Serialize for SlotMap<K, V> {
1191 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1192 where
1193 S: Serializer,
1194 {
1195 self.slots.serialize(serializer)
1196 }
1197 }
1198
1199 impl<'de, K, V> Deserialize<'de> for SlotMap<K, V>
1200 where
1201 K: Key,
1202 V: Deserialize<'de>,
1203 {
1204 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1205 where
1206 D: Deserializer<'de>,
1207 {
1208 let mut slots: Vec<Slot<V>> = Deserialize::deserialize(deserializer)?;
1209 if slots.len() >= u32::max_value() as usize {
1210 return Err(de::Error::custom(&"too many slots"));
1211 }
1212
1213 // Ensure the first slot exists and is empty for the sentinel.
1214 if slots.get(0).map_or(true, |slot| slot.version % 2 == 1) {
1215 return Err(de::Error::custom(&"first slot not empty"));
1216 }
1217
1218 slots[0].version = 0;
1219 slots[0].u.next_free = 0;
1220
1221 // We have our slots, rebuild freelist.
1222 let mut num_elems = 0;
1223 let mut next_free = slots.len();
1224 for (i, slot) in slots[1..].iter_mut().enumerate() {
1225 if slot.occupied() {
1226 num_elems += 1;
1227 } else {
1228 slot.u.next_free = next_free as u32;
1229 next_free = i + 1;
1230 }
1231 }
1232
1233 Ok(Self {
1234 num_elems,
1235 slots,
1236 free_head: next_free as u32,
1237 _k: PhantomData,
1238 })
1239 }
1240 }
1241}
1242
1243#[cfg(test)]
1244mod tests {
1245 use super::*;
1246 use quickcheck::quickcheck;
1247 use std::collections::HashMap;
1248
1249 #[derive(Clone)]
1250 struct CountDrop<'a>(&'a std::cell::RefCell<usize>);
1251
1252 impl<'a> Drop for CountDrop<'a> {
1253 fn drop(&mut self) {
1254 *self.0.borrow_mut() += 1;
1255 }
1256 }
1257
1258 #[cfg(all(nightly, feature = "unstable"))]
1259 #[test]
1260 fn check_drops() {
1261 let drops = std::cell::RefCell::new(0usize);
1262
1263 {
1264 let mut clone = {
1265 // Insert 1000 items.
1266 let mut sm = SlotMap::new();
1267 let mut sm_keys = Vec::new();
1268 for _ in 0..1000 {
1269 sm_keys.push(sm.insert(CountDrop(&drops)));
1270 }
1271
1272 // Remove even keys.
1273 for i in (0..1000).filter(|i| i % 2 == 0) {
1274 sm.remove(sm_keys[i]);
1275 }
1276
1277 // Should only have dropped 500 so far.
1278 assert_eq!(*drops.borrow(), 500);
1279
1280 // Let's clone ourselves and then die.
1281 sm.clone()
1282 };
1283
1284 // Now all original items should have been dropped exactly once.
1285 assert_eq!(*drops.borrow(), 1000);
1286
1287 // Reuse some empty slots.
1288 for _ in 0..250 {
1289 clone.insert(CountDrop(&drops));
1290 }
1291 }
1292
1293 // 1000 + 750 drops in total should have happened.
1294 assert_eq!(*drops.borrow(), 1750);
1295 }
1296
1297 #[cfg(all(nightly, feature = "unstable"))]
1298 #[test]
1299 fn disjoint() {
1300 // Intended to be run with miri to find any potential UB.
1301 let mut sm = SlotMap::new();
1302
1303 // Some churn.
1304 for i in 0..20usize {
1305 sm.insert(i);
1306 }
1307 sm.retain(|_, i| *i % 2 == 0);
1308
1309 let keys: Vec<_> = sm.keys().collect();
1310 for i in 0..keys.len() {
1311 for j in 0..keys.len() {
1312 if let Some([r0, r1]) = sm.get_disjoint_mut([keys[i], keys[j]]) {
1313 *r0 ^= *r1;
1314 *r1 = r1.wrapping_add(*r0);
1315 } else {
1316 assert!(i == j);
1317 }
1318 }
1319 }
1320
1321 for i in 0..keys.len() {
1322 for j in 0..keys.len() {
1323 for k in 0..keys.len() {
1324 if let Some([r0, r1, r2]) = sm.get_disjoint_mut([keys[i], keys[j], keys[k]]) {
1325 *r0 ^= *r1;
1326 *r0 = r0.wrapping_add(*r2);
1327 *r1 ^= *r0;
1328 *r1 = r1.wrapping_add(*r2);
1329 *r2 ^= *r0;
1330 *r2 = r2.wrapping_add(*r1);
1331 } else {
1332 assert!(i == j || j == k || i == k);
1333 }
1334 }
1335 }
1336 }
1337 }
1338
1339 quickcheck! {
1340 fn qc_slotmap_equiv_hashmap(operations: Vec<(u8, u32)>) -> bool {
1341 let mut hm = HashMap::new();
1342 let mut hm_keys = Vec::new();
1343 let mut unique_key = 0u32;
1344 let mut sm = SlotMap::new();
1345 let mut sm_keys = Vec::new();
1346
1347 #[cfg(not(feature = "serde"))]
1348 let num_ops = 3;
1349 #[cfg(feature = "serde")]
1350 let num_ops = 4;
1351
1352 for (op, val) in operations {
1353 match op % num_ops {
1354 // Insert.
1355 0 => {
1356 hm.insert(unique_key, val);
1357 hm_keys.push(unique_key);
1358 unique_key += 1;
1359
1360 sm_keys.push(sm.insert(val));
1361 }
1362
1363 // Delete.
1364 1 => {
1365 if hm_keys.is_empty() { continue; }
1366
1367 let idx = val as usize % hm_keys.len();
1368 if hm.remove(&hm_keys[idx]) != sm.remove(sm_keys[idx]) {
1369 return false;
1370 }
1371 }
1372
1373 // Access.
1374 2 => {
1375 if hm_keys.is_empty() { continue; }
1376 let idx = val as usize % hm_keys.len();
1377 let (hm_key, sm_key) = (&hm_keys[idx], sm_keys[idx]);
1378
1379 if hm.contains_key(hm_key) != sm.contains_key(sm_key) ||
1380 hm.get(hm_key) != sm.get(sm_key) {
1381 return false;
1382 }
1383 }
1384
1385 // Serde round-trip.
1386 #[cfg(feature = "serde")]
1387 3 => {
1388 let ser = serde_json::to_string(&sm).unwrap();
1389 sm = serde_json::from_str(&ser).unwrap();
1390 }
1391
1392 _ => unreachable!(),
1393 }
1394 }
1395
1396 let mut smv: Vec<_> = sm.values().collect();
1397 let mut hmv: Vec<_> = hm.values().collect();
1398 smv.sort();
1399 hmv.sort();
1400 smv == hmv
1401 }
1402 }
1403
1404 #[cfg(feature = "serde")]
1405 #[test]
1406 fn slotmap_serde() {
1407 let mut sm = SlotMap::new();
1408 // Self-referential structure.
1409 let first = sm.insert_with_key(|k| (k, 23i32));
1410 let second = sm.insert((first, 42));
1411
1412 // Make some empty slots.
1413 let empties = vec![sm.insert((first, 0)), sm.insert((first, 0))];
1414 empties.iter().for_each(|k| {
1415 sm.remove(*k);
1416 });
1417
1418 let third = sm.insert((second, 0));
1419 sm[first].0 = third;
1420
1421 let ser = serde_json::to_string(&sm).unwrap();
1422 let de: SlotMap<DefaultKey, (DefaultKey, i32)> = serde_json::from_str(&ser).unwrap();
1423 assert_eq!(de.len(), sm.len());
1424
1425 let mut smkv: Vec<_> = sm.iter().collect();
1426 let mut dekv: Vec<_> = de.iter().collect();
1427 smkv.sort();
1428 dekv.sort();
1429 assert_eq!(smkv, dekv);
1430 }
1431
1432 #[cfg(feature = "serde")]
1433 #[test]
1434 fn slotmap_serde_freelist() {
1435 let mut sm = SlotMap::new();
1436 let k = sm.insert(5i32);
1437 sm.remove(k);
1438
1439 let ser = serde_json::to_string(&sm).unwrap();
1440 let mut de: SlotMap<DefaultKey, i32> = serde_json::from_str(&ser).unwrap();
1441
1442 de.insert(0);
1443 de.insert(1);
1444 de.insert(2);
1445 assert_eq!(de.len(), 3);
1446 }
1447}