#[derive(Debug)]
pub struct Store<T> {
sparse: Vec<u32>,
dense: Vec<u32>,
values: Vec<T>,
}
const ABSENT: u32 = u32::MAX;
impl<T> Default for Store<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> Store<T> {
#[must_use]
pub const fn new() -> Self {
Self {
sparse: Vec::new(),
dense: Vec::new(),
values: Vec::new(),
}
}
#[must_use]
pub fn len(&self) -> usize {
self.values.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.values.is_empty()
}
#[must_use]
pub fn contains(&self, slot: u32) -> bool {
self.position(slot).is_some()
}
fn position(&self, slot: u32) -> Option<usize> {
match self.sparse.get(slot as usize).copied() {
Some(ABSENT) | None => None,
Some(position) => Some(position as usize),
}
}
#[must_use]
pub fn get(&self, slot: u32) -> Option<&T> {
self.values.get(self.position(slot)?)
}
pub fn get_mut(&mut self, slot: u32) -> Option<&mut T> {
let position = self.position(slot)?;
self.values.get_mut(position)
}
pub fn insert(&mut self, slot: u32, value: T) -> Option<T> {
if let Some(position) = self.position(slot) {
let existing = self.values.get_mut(position)?;
return Some(core::mem::replace(existing, value));
}
let needed = (slot as usize).checked_add(1)?;
if self.sparse.len() < needed {
self.sparse.resize(needed, ABSENT);
}
let position = u32::try_from(self.dense.len()).ok()?;
if let Some(entry) = self.sparse.get_mut(slot as usize) {
*entry = position;
}
self.dense.push(slot);
self.values.push(value);
None
}
pub fn remove(&mut self, slot: u32) -> Option<T> {
let position = self.position(slot)?;
let last = self.dense.len().checked_sub(1)?;
self.dense.swap(position, last);
self.values.swap(position, last);
if let Some(moved) = self.dense.get(position).copied()
&& position != last
&& let Some(entry) = self.sparse.get_mut(moved as usize)
{
*entry = u32::try_from(position).unwrap_or(ABSENT);
}
if let Some(entry) = self.sparse.get_mut(slot as usize) {
*entry = ABSENT;
}
self.dense.pop();
self.values.pop()
}
pub fn iter(&self) -> impl Iterator<Item = (u32, &T)> + '_ {
self.sparse
.iter()
.enumerate()
.filter(|(_, position)| **position != ABSENT)
.filter_map(move |(slot, position)| {
let value = self.values.get(*position as usize)?;
Some((u32::try_from(slot).ok()?, value))
})
}
pub(crate) fn scan_len(&self) -> usize {
self.sparse.len()
}
pub fn for_each_mut(&mut self, mut visit: impl FnMut(u32, &mut T)) {
let order: Vec<(u32, u32)> = self
.sparse
.iter()
.enumerate()
.filter(|(_, position)| **position != ABSENT)
.filter_map(|(slot, position)| Some((u32::try_from(slot).ok()?, *position)))
.collect();
for (slot, position) in order {
if let Some(value) = self.values.get_mut(position as usize) {
visit(slot, value);
}
}
}
pub fn iter_unordered(&self) -> impl Iterator<Item = &T> + '_ {
self.values.iter()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_new_store_is_empty() {
let store: Store<u32> = Store::new();
assert!(store.is_empty());
assert_eq!(store.len(), 0);
assert!(store.get(0).is_none());
assert!(!store.contains(0));
}
#[test]
fn default_and_new_agree() {
let made: Store<u32> = Store::default();
assert!(made.is_empty());
assert_eq!(made.len(), Store::<u32>::new().len());
}
#[test]
fn insert_then_get_returns_the_value() {
let mut store = Store::new();
assert!(store.insert(3, "three").is_none());
assert_eq!(store.get(3), Some(&"three"));
assert!(store.contains(3));
assert_eq!(store.len(), 1);
assert!(!store.contains(0));
}
#[test]
fn inserting_twice_replaces_and_returns_the_old_value() {
let mut store = Store::new();
store.insert(1, 10);
assert_eq!(store.insert(1, 20), Some(10));
assert_eq!(store.get(1), Some(&20));
assert_eq!(store.len(), 1, "replacing must not grow the store");
}
#[test]
fn get_mut_edits_in_place() {
let mut store = Store::new();
store.insert(2, 5);
if let Some(value) = store.get_mut(2) {
*value += 1;
}
assert_eq!(store.get(2), Some(&6));
assert!(store.get_mut(9).is_none());
}
#[test]
fn removing_from_the_middle_keeps_every_other_lookup_correct() {
let mut store = Store::new();
for slot in 0..5 {
store.insert(slot, slot * 100);
}
assert_eq!(store.remove(1), Some(100));
assert_eq!(store.len(), 4);
assert!(!store.contains(1));
for slot in [0u32, 2, 3, 4] {
assert_eq!(store.get(slot), Some(&(slot * 100)), "slot {slot}");
}
}
#[test]
fn removing_the_last_element_is_also_correct() {
let mut store = Store::new();
store.insert(0, 'a');
store.insert(1, 'b');
assert_eq!(store.remove(1), Some('b'));
assert_eq!(store.get(0), Some(&'a'));
assert!(!store.contains(1));
assert_eq!(store.remove(0), Some('a'));
assert!(store.is_empty());
}
#[test]
fn removing_what_is_not_there_returns_nothing() {
let mut store: Store<u8> = Store::new();
assert!(store.remove(7).is_none());
store.insert(0, 1);
assert!(store.remove(7).is_none());
assert_eq!(store.len(), 1);
}
#[test]
fn iteration_is_by_slot_whatever_the_churn() {
let mut store = Store::new();
for slot in [5u32, 1, 9, 3, 7] {
store.insert(slot, slot);
}
store.remove(3);
store.insert(2, 2);
store.remove(9);
let seen: Vec<u32> = store.iter().map(|(slot, _)| slot).collect();
assert_eq!(seen, vec![1, 2, 5, 7]);
let dense: Vec<u32> = store.iter_unordered().copied().collect();
assert_ne!(
dense, seen,
"dense order happened to match; pick harsher churn"
);
}
#[test]
fn for_each_mut_visits_in_slot_order_and_can_edit() {
let mut store = Store::new();
for slot in [4u32, 0, 2] {
store.insert(slot, slot);
}
let mut order = Vec::new();
store.for_each_mut(|slot, value| {
order.push(slot);
*value += 1;
});
assert_eq!(order, vec![0, 2, 4]);
assert_eq!(store.get(0), Some(&1));
assert_eq!(store.get(4), Some(&5));
}
#[test]
fn the_store_survives_a_long_churn() {
let mut store = Store::new();
for round in 0..50u32 {
for slot in 0..20u32 {
store.insert(slot, round * 100 + slot);
}
for slot in (0..20u32).step_by(3) {
store.remove(slot);
}
}
let seen: Vec<u32> = store.iter().map(|(slot, _)| slot).collect();
let expected: Vec<u32> = (0..20).filter(|slot| slot % 3 != 0).collect();
assert_eq!(seen, expected);
assert_eq!(store.len(), expected.len());
}
}