From a9005d7a599dc66f635639b52ce2e5e430184845 Mon Sep 17 00:00:00 2001 From: Ryan Date: Sat, 7 May 2022 01:05:23 -0700 Subject: [PATCH] Make slotmap keys unique This lets us use slotmap keys between different instances of the slotmap without unspecified behavior happening. Also, we no longer have to worry about version overflow. The downside is that the keys are now 12 bytes instead of 8, but this is an acceptable trade if it means fewer bugs. --- src/slotmap.rs | 196 ++++++++++++++++++++++--------------------------- 1 file changed, 87 insertions(+), 109 deletions(-) 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()); + } }