orchid/orchid-base/src/id_store.rs

use std::ops::{Index, IndexMut};

use itertools::Itertools;

enum Rec<T> {
	Val(T),
	Next(usize),
}

/// A simple and very fast store that assigns small stable integer IDs to
/// objects. It uses a free-list for O(1) insertion, deletion and retrieval.
pub struct IdStore<T> {
	first: usize,
	values: Vec<Rec<T>>,
}
impl<T> IdStore<T> {
	pub fn new() -> Self { IdStore { first: 0, values: Vec::new() } }
	pub fn add(&mut self, value: T) -> usize {
		if self.first == 0 && self.values.is_empty() {
			self.first = 1;
			self.values.push(Rec::Val(value));
			return 0;
		}
		if self.first == self.values.len() {
			let len = self.values.len();
			self.values.extend((len..len * 2).map(|i| Rec::Next(i + 1)));
		}
		let Some(rec) = self.values.get_mut(self.first) else {
			panic!("Bounds check and growth above")
		};
		let Rec::Next(next) = rec else {
			panic!("first should always point to an empty space or one past the length")
		};
		let id = std::mem::replace(&mut self.first, *next);
		*rec = Rec::Val(value);
		id
	}
	pub fn add_with(&mut self, cb: impl FnOnce(usize) -> T) -> usize { self.add(cb(self.first)) }
	pub fn remove(&mut self, id: usize) -> T {
		let Some(rec) = self.values.get_mut(id) else { panic!("Index out of bounds") };
		let Rec::Val(val) = std::mem::replace(rec, Rec::Next(self.first)) else {
			panic!("Index vacated")
		};
		self.first = id;
		val
	}
	pub fn iter(&self) -> impl Iterator<Item = (usize, &T)> {
		(self.values.iter().enumerate())
			.filter_map(|(i, rec)| if let Rec::Val(val) = rec { Some((i, val)) } else { None })
	}
	pub fn iter_mut(&mut self) -> impl Iterator<Item = (usize, &mut T)> {
		(self.values.iter_mut().enumerate())
			.filter_map(|(i, rec)| if let Rec::Val(val) = rec { Some((i, val)) } else { None })
	}
}
#[allow(clippy::type_complexity, reason = "This is verbose enough as it is")]
pub struct IntoIter<T>(
	std::iter::FilterMap<
		std::iter::Enumerate<std::vec::IntoIter<Rec<T>>>,
		fn((usize, Rec<T>)) -> Option<(usize, T)>,
	>,
);
impl<T> Iterator for IntoIter<T> {
	type Item = (usize, T);
	fn next(&mut self) -> Option<Self::Item> { self.0.next() }
	fn size_hint(&self) -> (usize, Option<usize>) { self.0.size_hint() }
}
impl<T> IntoIterator for IdStore<T> {
	type Item = (usize, T);
	type IntoIter = IntoIter<T>;
	fn into_iter(self) -> Self::IntoIter {
		IntoIter(
			(self.values.into_iter().enumerate())
				.filter_map(|(i, rec)| if let Rec::Val(val) = rec { Some((i, val)) } else { None }),
		)
	}
}
impl<T> Index<usize> for IdStore<T> {
	type Output = T;
	fn index(&self, index: usize) -> &Self::Output {
		match self.values.get(index) {
			Some(Rec::Val(val)) => val,
			_ => panic!("Invalid or stale index"),
		}
	}
}
impl<T> IndexMut<usize> for IdStore<T> {
	fn index_mut(&mut self, index: usize) -> &mut Self::Output {
		match self.values.get_mut(index) {
			Some(Rec::Val(val)) => val,
			_ => panic!("Invalid or stale index"),
		}
	}
}

impl<T> Default for IdStore<T> {
	fn default() -> Self { Self::new() }
}
impl<A> FromIterator<A> for IdStore<A> {
	fn from_iter<T: IntoIterator<Item = A>>(iter: T) -> Self {
		let values = iter.into_iter().map(|a| Rec::Val(a)).collect_vec();
		Self { first: values.len(), values }
	}
}

#[cfg(test)]
mod test {
	use super::*;

	#[test]
	fn add_and_retrieve() {
		let mut store = IdStore::new();
		let key1 = store.add(14);
		let key2 = store.add(34);
		assert_eq!(store[key1], 14);
		assert_eq!(store[key2], 34);
		assert_eq!(store.remove(key1), 14);
		assert_eq!(store.iter().collect_vec(), vec![(key2, &34)]);
	}
}