diff --git a/src/slotmap.rs b/src/slotmap.rs index e68426b..d080e01 100644 --- a/src/slotmap.rs +++ b/src/slotmap.rs @@ -2,7 +2,8 @@ use std::iter::FusedIterator; use std::mem; -use std::num::NonZeroU32; +use std::num::{NonZeroU32, NonZeroU64}; +use std::sync::atomic::{AtomicU64, Ordering}; use rayon::iter::{ IndexedParallelIterator, IntoParallelRefIterator, IntoParallelRefMutIterator, ParallelIterator, @@ -11,6 +12,7 @@ use rayon::iter::{ #[derive(Clone, Debug)] pub struct SlotMap { slots: Vec>, + /// Top of the free stack. next_free_head: u32, /// The number of occupied slots. count: u32, @@ -18,30 +20,47 @@ pub struct SlotMap { #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)] pub struct Key { - version: NonZeroU32, index: u32, + // Split the u64 version into two u32 fields so that the key is 12 bytes on 64 bit systems. + version_high: NonZeroU32, + version_low: u32, } impl Key { - pub fn version(self) -> NonZeroU32 { - self.version + fn new(index: u32, version: NonZeroU64) -> Self { + Self { + index, + version_high: NonZeroU32::new((version.get() >> 32) as u32) + .expect("versions <= 0x00000000ffffffff are illegal"), + version_low: version.get() as u32, + } + } + + fn new_unique(index: u32) -> Self { + static NEXT: AtomicU64 = AtomicU64::new(u64::MAX); + + let version = NEXT.fetch_sub(1, Ordering::SeqCst); + Self { + index, + version_high: ((version >> 32) as u32).try_into().unwrap(), + version_low: version as u32, + } } pub fn index(self) -> u32 { self.index } + + pub fn version(self) -> NonZeroU64 { + let n = (self.version_high.get() as u64) << 32 | self.version_low as u64; + NonZeroU64::new(n).expect("version should be nonzero") + } } #[derive(Clone, Debug)] -struct Slot { - version: NonZeroU32, - item: Item, -} - -#[derive(Clone, Debug)] -enum Item { - Occupied(T), - Vacant { next_free: u32 }, +enum Slot { + Occupied { value: T, version: NonZeroU64 }, + Free { next_free: u32 }, } impl SlotMap { @@ -58,55 +77,51 @@ impl SlotMap { } pub fn insert(&mut self, val: T) -> Key { + self.insert_with(|_| val) + } + + pub fn insert_with(&mut self, f: impl FnOnce(Key) -> T) -> Key { assert!(self.count < u32::MAX, "SlotMap: too many items inserted"); if self.next_free_head == self.slots.len() as u32 { - self.slots.push(Slot { - version: ONE, - item: Item::Occupied(val), - }); - self.count += 1; self.next_free_head += 1; - Key { - version: ONE, - index: self.next_free_head - 1, - } + + let key = Key::new_unique(self.next_free_head - 1); + + self.slots.push(Slot::Occupied { + value: f(key), + version: key.version(), + }); + key } else { let slot = &mut self.slots[self.next_free_head as usize]; - slot.version = match NonZeroU32::new(slot.version.get().wrapping_add(1)) { - Some(n) => n, - None => { - log::debug!("SlotMap: version overflow at idx = {}", self.next_free_head); - ONE - } + + let next_free = match slot { + Slot::Occupied { .. } => unreachable!("corrupt free list"), + Slot::Free { next_free } => *next_free, }; - let next_free = match slot.item { - Item::Occupied(_) => unreachable!("corrupt free list"), - Item::Vacant { next_free } => next_free, - }; + let key = Key::new_unique(self.next_free_head); - let key = Key { - version: slot.version, - index: self.next_free_head, + *slot = Slot::Occupied { + value: f(key), + version: key.version(), }; self.next_free_head = next_free; self.count += 1; - slot.item = Item::Occupied(val); key } } pub fn remove(&mut self, key: Key) -> Option { - let Slot { version, item } = self.slots.get_mut(key.index as usize)?; - - match item { - Item::Occupied(_) if *version == key.version => { - let item = mem::replace( - item, - Item::Vacant { + let slot = self.slots.get_mut(key.index as usize)?; + match slot { + Slot::Occupied { version, .. } if *version == key.version() => { + let old_slot = mem::replace( + slot, + Slot::Free { next_free: self.next_free_head, }, ); @@ -114,9 +129,9 @@ impl SlotMap { self.next_free_head = key.index; self.count -= 1; - match item { - Item::Occupied(val) => Some(val), - Item::Vacant { next_free } => unreachable!(), + match old_slot { + Slot::Occupied { value, .. } => Some(value), + Slot::Free { .. } => unreachable!(), } } _ => None, @@ -125,31 +140,18 @@ impl SlotMap { pub fn get(&self, key: Key) -> Option<&T> { match self.slots.get(key.index as usize)? { - Slot { - version, - item: Item::Occupied(val), - } if *version == key.version => Some(val), + Slot::Occupied { value, version } if *version == key.version() => Some(value), _ => None, } } pub fn get_mut(&mut self, key: Key) -> Option<&mut T> { match self.slots.get_mut(key.index as usize)? { - Slot { - version, - item: Item::Occupied(val), - } if *version == key.version => Some(val), + Slot::Occupied { value, version } if *version == key.version() => Some(value), _ => None, } } - pub fn key_at_index(&self, idx: usize) -> Option { - Some(Key { - version: self.slots.get(idx)?.version, - index: idx as u32, - }) - } - pub fn clear(&mut self) { self.slots.clear(); self.next_free_head = 0; @@ -158,13 +160,11 @@ impl SlotMap { pub fn retain(&mut self, mut f: impl FnMut(Key, &mut T) -> bool) { for (i, slot) in self.slots.iter_mut().enumerate() { - if let Item::Occupied(val) = &mut slot.item { - let key = Key { - version: slot.version, - index: i as u32, - }; - if !f(key, val) { - slot.item = Item::Vacant { + if let Slot::Occupied { value, version } = &mut slot { + let key = Key::new(i as u32, *version); + + if !f(key, value) { + *slot = Slot::Free { next_free: self.next_free_head, }; @@ -179,15 +179,9 @@ impl SlotMap { self.slots .iter() .enumerate() - .filter_map(|(i, slot)| match &slot.item { - Item::Occupied(val) => Some(( - Key { - version: slot.version, - index: i as u32, - }, - val, - )), - Item::Vacant { .. } => None, + .filter_map(|(i, slot)| match &slot { + Slot::Occupied { value, version } => Some((Key::new(i as u32, *version), value)), + Slot::Free { .. } => None, }) } @@ -195,15 +189,9 @@ impl SlotMap { self.slots .iter_mut() .enumerate() - .filter_map(|(i, slot)| match &mut slot.item { - Item::Occupied(val) => Some(( - Key { - version: slot.version, - index: i as u32, - }, - val, - )), - Item::Vacant { .. } => None, + .filter_map(|(i, slot)| match slot { + Slot::Occupied { value, version } => Some((Key::new(i as u32, *version), value)), + Slot::Free { .. } => None, }) } } @@ -213,15 +201,9 @@ impl SlotMap { self.slots .par_iter() .enumerate() - .filter_map(|(i, slot)| match &slot.item { - Item::Occupied(val) => Some(( - Key { - version: slot.version, - index: i as u32, - }, - val, - )), - Item::Vacant { .. } => None, + .filter_map(|(i, slot)| match &slot { + Slot::Occupied { value, version } => Some((Key::new(i as u32, *version), value)), + Slot::Free { .. } => None, }) } } @@ -231,15 +213,9 @@ impl SlotMap { self.slots .par_iter_mut() .enumerate() - .filter_map(|(i, slot)| match &mut slot.item { - Item::Occupied(val) => Some(( - Key { - version: slot.version, - index: i as u32, - }, - val, - )), - Item::Vacant { .. } => None, + .filter_map(|(i, slot)| match slot { + Slot::Occupied { value, version } => Some((Key::new(i as u32, *version), value)), + Slot::Free { .. } => None, }) } } @@ -250,11 +226,6 @@ impl Default for SlotMap { } } -const ONE: NonZeroU32 = match NonZeroU32::new(1) { - Some(n) => n, - None => unreachable!(), -}; - #[cfg(test)] mod tests { use super::*; @@ -279,6 +250,7 @@ mod tests { assert_eq!(sm.count(), 0); } + #[test] fn retain() { let mut sm = SlotMap::new(); @@ -294,4 +266,10 @@ mod tests { assert_eq!(sm.get(k0), None); assert_eq!(sm.get(k2), None); } + + #[test] + #[should_panic] + fn bad_key() { + let _ = Key::new(0, NonZeroU64::new(0x00000000ffffffff).unwrap()); + } }