#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SlotId(pub(crate) usize);
impl SlotId {
pub fn index(self) -> usize {
self.0
}
}
#[derive(Debug)]
pub struct SlotArena<T> {
slots: Vec<Option<T>>,
free_list: Vec<usize>,
len: usize,
}
impl<T> SlotArena<T> {
pub fn new() -> Self {
Self {
slots: Vec::new(),
free_list: Vec::new(),
len: 0,
}
}
pub fn with_capacity(capacity: usize) -> Self {
Self {
slots: Vec::with_capacity(capacity),
free_list: Vec::new(),
len: 0,
}
}
#[inline]
pub fn insert(&mut self, value: T) -> SlotId {
let idx = if let Some(idx) = self.free_list.pop() {
self.slots[idx] = Some(value);
idx
} else {
self.slots.push(Some(value));
self.slots.len() - 1
};
self.len += 1;
SlotId(idx)
}
#[inline]
pub fn remove(&mut self, id: SlotId) -> Option<T> {
let slot = self.slots.get_mut(id.0)?;
let value = slot.take()?;
self.free_list.push(id.0);
self.len -= 1;
Some(value)
}
#[inline]
pub fn get(&self, id: SlotId) -> Option<&T> {
self.slots.get(id.0).and_then(|slot| slot.as_ref())
}
#[inline]
pub fn get_mut(&mut self, id: SlotId) -> Option<&mut T> {
self.slots.get_mut(id.0).and_then(|slot| slot.as_mut())
}
#[inline]
pub fn contains(&self, id: SlotId) -> bool {
self.slots
.get(id.0)
.map(|slot| slot.is_some())
.unwrap_or(false)
}
pub fn contains_index(&self, index: usize) -> bool {
self.slots
.get(index)
.map(|slot| slot.is_some())
.unwrap_or(false)
}
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub fn capacity(&self) -> usize {
self.slots.capacity()
}
pub fn reserve_slots(&mut self, additional: usize) {
self.slots.reserve(additional);
}
pub fn shrink_to_fit(&mut self) {
self.slots.shrink_to_fit();
self.free_list.shrink_to_fit();
}
pub fn clear_shrink(&mut self) {
self.clear();
self.shrink_to_fit();
}
pub fn clear(&mut self) {
self.slots.clear();
self.free_list.clear();
self.len = 0;
}
pub fn iter(&self) -> impl Iterator<Item = (SlotId, &T)> {
self.slots
.iter()
.enumerate()
.filter_map(|(idx, slot)| slot.as_ref().map(|value| (SlotId(idx), value)))
}
pub fn iter_ids(&self) -> impl Iterator<Item = SlotId> + '_ {
self.slots
.iter()
.enumerate()
.filter_map(|(idx, slot)| slot.as_ref().map(|_| SlotId(idx)))
}
#[cfg(any(test, debug_assertions))]
pub fn debug_snapshot(&self) -> SlotArenaSnapshot {
let mut free_list = self.free_list.clone();
free_list.sort_unstable();
let mut live_ids: Vec<_> = self.iter_ids().collect();
live_ids.sort_by_key(|id| id.index());
SlotArenaSnapshot {
len: self.len,
slots_len: self.slots.len(),
free_list,
live_ids,
}
}
pub fn approx_bytes(&self) -> usize {
std::mem::size_of::<Self>()
+ self.slots.capacity() * std::mem::size_of::<Option<T>>()
+ self.free_list.capacity() * std::mem::size_of::<usize>()
}
#[cfg(any(test, debug_assertions))]
pub fn debug_validate_invariants(&self) {
let live_count = self.slots.iter().filter(|slot| slot.is_some()).count();
assert_eq!(self.len, live_count);
assert!(self.len <= self.slots.len());
let mut seen_free = std::collections::HashSet::new();
for &idx in &self.free_list {
assert!(idx < self.slots.len());
assert!(self.slots[idx].is_none());
assert!(seen_free.insert(idx));
}
assert_eq!(self.slots.len(), self.free_list.len() + self.len);
}
}
#[cfg(any(test, debug_assertions))]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SlotArenaSnapshot {
pub len: usize,
pub slots_len: usize,
pub free_list: Vec<usize>,
pub live_ids: Vec<SlotId>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ShardedSlotId {
shard: usize,
slot: SlotId,
}
impl ShardedSlotId {
pub fn shard(self) -> usize {
self.shard
}
pub fn slot(self) -> SlotId {
self.slot
}
}
#[derive(Debug)]
pub struct ShardedSlotArena<T> {
shards: Vec<parking_lot::RwLock<SlotArena<T>>>,
selector: crate::ds::ShardSelector,
next_shard: std::sync::atomic::AtomicUsize,
}
impl<T> ShardedSlotArena<T> {
pub fn new(shards: usize) -> Self {
let shards = shards.max(1);
let mut arenas = Vec::with_capacity(shards);
for _ in 0..shards {
arenas.push(parking_lot::RwLock::new(SlotArena::new()));
}
Self {
shards: arenas,
selector: crate::ds::ShardSelector::new(shards, 0),
next_shard: std::sync::atomic::AtomicUsize::new(0),
}
}
pub fn with_capacity(shards: usize, capacity: usize) -> Self {
let shards = shards.max(1);
let mut arenas = Vec::with_capacity(shards);
for _ in 0..shards {
arenas.push(parking_lot::RwLock::new(SlotArena::with_capacity(capacity)));
}
Self {
shards: arenas,
selector: crate::ds::ShardSelector::new(shards, 0),
next_shard: std::sync::atomic::AtomicUsize::new(0),
}
}
pub fn with_shards(shards: usize, capacity_per_shard: usize) -> Self {
Self::with_capacity(shards, capacity_per_shard)
}
pub fn with_shards_seed(shards: usize, capacity_per_shard: usize, seed: u64) -> Self {
let shards = shards.max(1);
let mut arenas = Vec::with_capacity(shards);
for _ in 0..shards {
arenas.push(parking_lot::RwLock::new(SlotArena::with_capacity(
capacity_per_shard,
)));
}
Self {
shards: arenas,
selector: crate::ds::ShardSelector::new(shards, seed),
next_shard: std::sync::atomic::AtomicUsize::new(0),
}
}
pub fn shard_for_key<K: std::hash::Hash>(&self, key: &K) -> usize {
self.selector.shard_for_key(key)
}
pub fn shard_count(&self) -> usize {
self.shards.len()
}
pub fn insert(&self, value: T) -> ShardedSlotId {
let shard = self
.next_shard
.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
% self.shards.len();
let mut arena = self.shards[shard].write();
let slot = arena.insert(value);
ShardedSlotId { shard, slot }
}
pub fn remove(&self, id: ShardedSlotId) -> Option<T> {
let mut arena = self.shards.get(id.shard)?.write();
arena.remove(id.slot)
}
pub fn get_with<R>(&self, id: ShardedSlotId, f: impl FnOnce(&T) -> R) -> Option<R> {
let arena = self.shards.get(id.shard)?.read();
arena.get(id.slot).map(f)
}
pub fn get_mut_with<R>(&self, id: ShardedSlotId, f: impl FnOnce(&mut T) -> R) -> Option<R> {
let mut arena = self.shards.get(id.shard)?.write();
arena.get_mut(id.slot).map(f)
}
pub fn contains(&self, id: ShardedSlotId) -> bool {
self.shards
.get(id.shard)
.map(|arena| arena.read().contains(id.slot))
.unwrap_or(false)
}
pub fn len(&self) -> usize {
self.shards.iter().map(|arena| arena.read().len()).sum()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn reserve_slots(&self, additional: usize) {
for arena in &self.shards {
arena.write().reserve_slots(additional);
}
}
pub fn shrink_to_fit(&self) {
for arena in &self.shards {
arena.write().shrink_to_fit();
}
}
pub fn clear_shrink(&self) {
for arena in &self.shards {
arena.write().clear_shrink();
}
}
pub fn approx_bytes(&self) -> usize {
self.shards
.iter()
.map(|arena| arena.read().approx_bytes())
.sum()
}
}
impl<T> Default for SlotArena<T> {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub struct ConcurrentSlotArena<T> {
inner: parking_lot::RwLock<SlotArena<T>>,
}
impl<T> ConcurrentSlotArena<T> {
pub fn new() -> Self {
Self {
inner: parking_lot::RwLock::new(SlotArena::new()),
}
}
pub fn with_capacity(capacity: usize) -> Self {
Self {
inner: parking_lot::RwLock::new(SlotArena::with_capacity(capacity)),
}
}
pub fn insert(&self, value: T) -> SlotId {
let mut arena = self.inner.write();
arena.insert(value)
}
pub fn remove(&self, id: SlotId) -> Option<T> {
let mut arena = self.inner.write();
arena.remove(id)
}
pub fn get_with<R>(&self, id: SlotId, f: impl FnOnce(&T) -> R) -> Option<R> {
let arena = self.inner.read();
arena.get(id).map(f)
}
pub fn try_get_with<R>(&self, id: SlotId, f: impl FnOnce(&T) -> R) -> Option<R> {
let arena = self.inner.try_read()?;
arena.get(id).map(f)
}
pub fn get_mut_with<R>(&self, id: SlotId, f: impl FnOnce(&mut T) -> R) -> Option<R> {
let mut arena = self.inner.write();
arena.get_mut(id).map(f)
}
pub fn try_get_mut_with<R>(&self, id: SlotId, f: impl FnOnce(&mut T) -> R) -> Option<R> {
let mut arena = self.inner.try_write()?;
arena.get_mut(id).map(f)
}
pub fn contains(&self, id: SlotId) -> bool {
let arena = self.inner.read();
arena.contains(id)
}
pub fn len(&self) -> usize {
let arena = self.inner.read();
arena.len()
}
pub fn is_empty(&self) -> bool {
let arena = self.inner.read();
arena.is_empty()
}
pub fn capacity(&self) -> usize {
let arena = self.inner.read();
arena.capacity()
}
pub fn reserve_slots(&self, additional: usize) {
let mut arena = self.inner.write();
arena.reserve_slots(additional);
}
pub fn shrink_to_fit(&self) {
let mut arena = self.inner.write();
arena.shrink_to_fit();
}
pub fn clear_shrink(&self) {
let mut arena = self.inner.write();
arena.clear_shrink();
}
pub fn clear(&self) {
let mut arena = self.inner.write();
arena.clear();
}
pub fn approx_bytes(&self) -> usize {
let arena = self.inner.read();
arena.approx_bytes()
}
}
impl<T> Default for ConcurrentSlotArena<T> {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn slot_arena_insert_remove_reuse() {
let mut arena = SlotArena::new();
let id1 = arena.insert("a");
let id2 = arena.insert("b");
assert_eq!(arena.len(), 2);
assert_eq!(arena.get(id1), Some(&"a"));
assert_eq!(arena.get(id2), Some(&"b"));
assert_eq!(arena.remove(id1), Some("a"));
assert_eq!(arena.len(), 1);
let id3 = arena.insert("c");
assert_eq!(arena.len(), 2);
assert_eq!(arena.get(id3), Some(&"c"));
assert_eq!(id1.index(), id3.index());
}
#[test]
fn concurrent_slot_arena_basic_ops() {
let arena = ConcurrentSlotArena::new();
let id = arena.insert(10);
assert_eq!(arena.get_with(id, |v| *v), Some(10));
assert!(arena.contains(id));
assert_eq!(arena.len(), 1);
arena.get_mut_with(id, |v| *v = 20);
assert_eq!(arena.get_with(id, |v| *v), Some(20));
assert_eq!(arena.remove(id), Some(20));
assert!(!arena.contains(id));
assert!(arena.is_empty());
}
#[test]
fn slot_arena_remove_invalid_id_is_none() {
let mut arena: SlotArena<i32> = SlotArena::new();
let id = SlotId(0);
assert_eq!(arena.remove(id), None);
assert!(!arena.contains(id));
assert!(arena.is_empty());
}
#[test]
fn slot_arena_clear_resets_state() {
let mut arena = SlotArena::with_capacity(4);
let a = arena.insert("a");
let b = arena.insert("b");
assert_eq!(arena.len(), 2);
assert!(arena.contains(a));
assert!(arena.contains(b));
arena.clear();
assert_eq!(arena.len(), 0);
assert!(arena.is_empty());
assert!(!arena.contains(a));
assert!(!arena.contains(b));
assert_eq!(arena.iter().count(), 0);
}
#[test]
fn slot_arena_iter_skips_empty_slots() {
let mut arena = SlotArena::new();
let a = arena.insert("a");
let b = arena.insert("b");
let c = arena.insert("c");
assert_eq!(arena.remove(b), Some("b"));
let mut values: Vec<_> = arena.iter().map(|(_, v)| *v).collect();
values.sort();
assert_eq!(values, vec!["a", "c"]);
assert!(arena.contains(a));
assert!(arena.contains(c));
}
#[test]
fn slot_arena_get_mut_updates_value() {
let mut arena = SlotArena::new();
let id = arena.insert(1);
if let Some(value) = arena.get_mut(id) {
*value = 2;
}
assert_eq!(arena.get(id), Some(&2));
}
#[test]
fn slot_arena_capacity_tracking() {
let arena: SlotArena<i32> = SlotArena::with_capacity(16);
assert!(arena.capacity() >= 16);
assert_eq!(arena.len(), 0);
}
#[test]
fn slot_arena_contains_index_and_iter_ids() {
let mut arena = SlotArena::new();
let a = arena.insert("a");
let b = arena.insert("b");
assert!(arena.contains_index(a.index()));
assert!(arena.contains_index(b.index()));
arena.remove(a);
assert!(!arena.contains_index(a.index()));
let ids: Vec<_> = arena.iter_ids().collect();
assert_eq!(ids, vec![b]);
}
#[test]
fn slot_arena_reserve_slots_grows_capacity() {
let mut arena: SlotArena<i32> = SlotArena::new();
let before = arena.capacity();
arena.reserve_slots(32);
assert!(arena.capacity() >= before + 32);
}
#[test]
fn slot_arena_debug_snapshot() {
let mut arena = SlotArena::new();
let a = arena.insert("a");
let b = arena.insert("b");
arena.remove(a);
let snapshot = arena.debug_snapshot();
assert_eq!(snapshot.len, 1);
assert!(snapshot.live_ids.contains(&b));
assert!(snapshot.free_list.contains(&a.index()));
}
#[test]
fn sharded_slot_arena_basic_ops() {
let arena = ShardedSlotArena::new(2);
let a = arena.insert(1);
let b = arena.insert(2);
assert!(arena.contains(a));
assert!(arena.contains(b));
assert_eq!(arena.get_with(a, |v| *v), Some(1));
assert_eq!(arena.remove(b), Some(2));
assert!(!arena.contains(b));
assert_eq!(arena.len(), 1);
}
#[test]
fn sharded_slot_arena_shard_for_key() {
let arena: ShardedSlotArena<i32> = ShardedSlotArena::new(4);
let shard = arena.shard_for_key(&"alpha");
assert!(shard < arena.shard_count());
}
#[test]
fn sharded_slot_arena_with_seed() {
let arena: ShardedSlotArena<i32> = ShardedSlotArena::with_shards_seed(4, 0, 99);
let shard = arena.shard_for_key(&"alpha");
assert!(shard < arena.shard_count());
}
#[test]
fn concurrent_slot_arena_try_ops() {
let arena = ConcurrentSlotArena::new();
let id = arena.insert(1);
assert_eq!(arena.try_get_with(id, |v| *v), Some(1));
arena.try_get_mut_with(id, |v| *v = 2);
assert_eq!(arena.get_with(id, |v| *v), Some(2));
}
#[test]
fn concurrent_slot_arena_clear_and_accessors() {
let arena = ConcurrentSlotArena::new();
let a = arena.insert(1);
let b = arena.insert(2);
assert_eq!(arena.get_with(a, |v| *v), Some(1));
assert_eq!(arena.get_with(b, |v| *v), Some(2));
arena.clear();
assert!(arena.is_empty());
assert!(!arena.contains(a));
assert!(!arena.contains(b));
assert_eq!(arena.get_with(a, |v| *v), None);
}
#[test]
fn concurrent_slot_arena_get_mut_with_updates_value() {
let arena = ConcurrentSlotArena::new();
let id = arena.insert(5);
arena.get_mut_with(id, |v| *v = 10);
assert_eq!(arena.get_with(id, |v| *v), Some(10));
}
#[test]
fn slot_arena_debug_invariants_hold() {
let mut arena = SlotArena::new();
let a = arena.insert(1);
let b = arena.insert(2);
arena.remove(a);
let _ = arena.insert(3);
assert!(arena.contains(b));
arena.debug_validate_invariants();
}
}