use std::{
fmt::{Debug, Display},
hash::Hash,
marker::PhantomData,
ops::{Deref, DerefMut, Index, IndexMut},
slice, vec,
};
use crate::intern::Intern;
use serde::{Deserialize, Deserializer, Serialize, Serializer, ser::SerializeSeq};
pub trait Identifier: Copy + Hash + From<usize> + Into<usize> + Eq + Ord + std::fmt::Debug {}
impl Identifier for usize {}
pub struct Identified<Id: Identifier, T> {
pub id: Id,
pub inner: T,
}
impl<Id: Identifier, T> Identified<Id, T> {
pub fn new(id: Id, data: T) -> Self {
Self { id, inner: data }
}
}
impl<Id: Identifier, T> Deref for Identified<Id, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
&self.inner
}
}
impl<Id: Identifier, T> DerefMut for Identified<Id, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.inner
}
}
impl<'a, Id: Identifier, T> Identified<Id, &'a mut T> {
pub fn immutable(self) -> Identified<Id, &'a T> {
Identified::new(self.id, &*self.inner)
}
}
impl<Id: Identifier, T: Display> Display for Identified<Id, T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.inner.fmt(f)
}
}
impl<Id: Identifier, T: Debug> Debug for Identified<Id, T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.inner.fmt(f)
}
}
pub struct Registry<Id: Identifier, T> {
chunks: Vec<Vec<T>>,
len: usize,
_marker: PhantomData<Id>,
}
#[inline]
fn locate(n: usize) -> (usize, usize) {
let m = n + 1;
let k = (usize::BITS - 1 - m.leading_zeros()) as usize;
(k, m - (1 << k))
}
impl<Id: Identifier, T: Serialize> Serialize for Registry<Id, T> {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut seq = serializer.serialize_seq(Some(self.len))?;
for e in self.chunks.iter().flatten() {
seq.serialize_element(e)?;
}
seq.end()
}
}
impl<'de, Id: Identifier, T: Deserialize<'de>> Deserialize<'de> for Registry<Id, T> {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
Ok(Vec::<T>::deserialize(deserializer)?.into_iter().collect())
}
}
impl<Id: Identifier, T: Clone> Clone for Registry<Id, T> {
fn clone(&self) -> Self {
self.chunks.iter().flatten().cloned().collect()
}
}
impl<Id: Identifier, T: Debug> Debug for Registry<Id, T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_list()
.entries(self.chunks.iter().flatten())
.finish()
}
}
impl<Id: Identifier, T: PartialEq> PartialEq for Registry<Id, T> {
fn eq(&self, other: &Self) -> bool {
self.len == other.len
&& self
.chunks
.iter()
.flatten()
.eq(other.chunks.iter().flatten())
}
}
impl<Id: Identifier, T: Eq> Eq for Registry<Id, T> {}
impl<Id: Identifier, T> Registry<Id, T> {
pub fn push(&mut self, e: T) -> Id {
let n = self.len;
let (k, offset) = locate(n);
if offset == 0 && self.chunks.len() == k {
self.chunks.push(Vec::with_capacity(1 << k));
}
self.chunks[k].push(e);
self.len += 1;
n.into()
}
pub fn len(&self) -> usize {
self.len
}
pub fn truncate(&mut self, len: usize) {
while self.len > len {
let (k, offset) = locate(self.len - 1);
self.chunks[k].truncate(offset);
self.len -= 1;
}
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
#[track_caller]
fn at(&self, n: usize) -> &T {
debug_assert!(
n < self.len,
"registry index {n} out of bounds for length {}",
self.len
);
let (k, offset) = locate(n);
&self.chunks[k][offset]
}
#[inline]
#[track_caller]
fn at_mut(&mut self, n: usize) -> &mut T {
debug_assert!(
n < self.len,
"registry index {n} out of bounds for length {}",
self.len
);
let (k, offset) = locate(n);
&mut self.chunks[k][offset]
}
#[track_caller]
pub fn replace(&mut self, id: Id, value: T) -> T {
std::mem::replace(self.at_mut(id.into()), value)
}
#[track_caller]
pub fn get(&self, id: Id) -> Identified<Id, &T> {
Identified::new(id, self.at(id.into()))
}
#[track_caller]
pub fn get_mut(&mut self, id: Id) -> Identified<Id, &mut T> {
Identified::new(id, self.at_mut(id.into()))
}
pub fn iter(&self) -> Iter<'_, Id, T> {
Iter {
iter: self.chunks.iter().flatten(),
index: 0,
_marker: PhantomData,
}
}
pub fn iter_mut(&mut self) -> IterMut<'_, Id, T> {
IterMut {
iter: self.chunks.iter_mut().flatten(),
index: 0,
_marker: PhantomData,
}
}
pub fn select_mut(&mut self, ids: &[Id]) -> Vec<&mut T> {
let mut want: std::collections::HashMap<usize, usize> =
std::collections::HashMap::with_capacity(ids.len());
for (pos, id) in ids.iter().enumerate() {
let n: usize = (*id).into();
assert!(
n < self.len,
"select_mut: id index {n} out of bounds (len {})",
self.len
);
let prev = want.insert(n, pos);
assert!(prev.is_none(), "select_mut: duplicate id index {n}");
}
let mut slots: Vec<Option<&mut T>> = (0..ids.len()).map(|_| None).collect();
for item in self.iter_mut() {
let n: usize = item.id.into();
if let Some(&pos) = want.get(&n) {
slots[pos] = Some(item.inner);
}
}
slots
.into_iter()
.map(|slot| slot.expect("select_mut: requested id had no backing slot"))
.collect()
}
}
pub struct Iter<'a, Id: Identifier, T> {
iter: std::iter::Flatten<slice::Iter<'a, Vec<T>>>,
index: usize,
_marker: PhantomData<Id>,
}
impl<'a, Id: Identifier, T> Iterator for Iter<'a, Id, T> {
type Item = Identified<Id, &'a T>;
fn next(&mut self) -> Option<Self::Item> {
let value = self.iter.next()?;
let id = Id::from(self.index);
self.index += 1;
Some(Identified::new(id, value))
}
}
pub struct IterMut<'a, Id: Identifier, T> {
iter: std::iter::Flatten<slice::IterMut<'a, Vec<T>>>,
index: usize,
_marker: PhantomData<Id>,
}
impl<'a, Id: Identifier, T> Iterator for IterMut<'a, Id, T> {
type Item = Identified<Id, &'a mut T>;
fn next(&mut self) -> Option<Self::Item> {
let value = self.iter.next()?;
let id = Id::from(self.index);
self.index += 1;
Some(Identified::new(id, value))
}
}
pub struct IntoIter<Id: Identifier, T> {
iter: std::iter::Flatten<vec::IntoIter<Vec<T>>>,
index: usize,
_marker: PhantomData<Id>,
}
impl<Id: Identifier, T> Iterator for IntoIter<Id, T> {
type Item = Identified<Id, T>;
fn next(&mut self) -> Option<Self::Item> {
let value = self.iter.next()?;
let id = Id::from(self.index);
self.index += 1;
Some(Identified::new(id, value))
}
}
impl<Id: Identifier, T> Default for Registry<Id, T> {
fn default() -> Self {
Self {
chunks: Vec::new(),
len: 0,
_marker: PhantomData,
}
}
}
impl<'a, Id: Identifier, T> IntoIterator for &'a Registry<Id, T> {
type Item = Identified<Id, &'a T>;
type IntoIter = Iter<'a, Id, T>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'a, Id: Identifier, T> IntoIterator for &'a mut Registry<Id, T> {
type Item = Identified<Id, &'a mut T>;
type IntoIter = IterMut<'a, Id, T>;
fn into_iter(self) -> Self::IntoIter {
self.iter_mut()
}
}
impl<Id: Identifier, T> IntoIterator for Registry<Id, T> {
type Item = Identified<Id, T>;
type IntoIter = IntoIter<Id, T>;
fn into_iter(self) -> Self::IntoIter {
IntoIter {
iter: self.chunks.into_iter().flatten(),
index: 0,
_marker: PhantomData,
}
}
}
impl<Id: Identifier, T> FromIterator<T> for Registry<Id, T> {
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
let mut reg = Self::default();
for e in iter {
reg.push(e);
}
reg
}
}
impl<Id: Identifier, T> Index<Id> for Registry<Id, T> {
type Output = T;
#[track_caller]
fn index(&self, index: Id) -> &Self::Output {
self.at(index.into())
}
}
impl<Id: Identifier, T> IndexMut<Id> for Registry<Id, T> {
#[track_caller]
fn index_mut(&mut self, index: Id) -> &mut Self::Output {
self.at_mut(index.into())
}
}
impl<Id: Identifier, T: Intern> Intern for Registry<Id, T> {
type Static = Registry<Id, T::Static>;
fn intern(self, pool: &mut super::intern::StringPool) -> Self::Static {
self.into_iter()
.map(|item| item.inner.intern(pool))
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use jstd_derive::Identifier;
#[derive(Identifier)]
struct Id(usize);
#[test]
fn truncate_keeps_earlier_ids_and_chunk_capacity() {
let mut registry: Registry<Id, u32> = (0..7).collect();
let capacity: Vec<usize> = registry.chunks.iter().map(Vec::capacity).collect();
registry.truncate(2);
assert_eq!(registry.len(), 2);
assert_eq!(registry[Id(1)], 1);
assert_eq!(registry.chunks[0].capacity(), capacity[0]);
assert_eq!(registry.chunks[1].capacity(), capacity[1]);
assert!(registry.chunks[2].is_empty());
assert_eq!(registry.push(9), Id(2));
assert_eq!(registry.push(10), Id(3));
assert_eq!(registry.chunks[2].capacity(), capacity[2]);
registry.truncate(0);
assert!(registry.is_empty());
registry.truncate(5);
assert!(registry.is_empty());
}
#[test]
fn pushing_after_truncate_reuses_the_emptied_chunks() {
let mut registry: Registry<Id, u32> = (0..7).collect();
let chunks = registry.chunks.len();
let capacity: Vec<usize> = registry.chunks.iter().map(Vec::capacity).collect();
for epoch in 0..100 {
registry.truncate(0);
for i in 0..7 {
assert_eq!(registry.push(i + epoch), Id(i as usize));
}
assert_eq!(registry.chunks.len(), chunks, "epoch {epoch}");
let now: Vec<usize> = registry.chunks.iter().map(Vec::capacity).collect();
assert_eq!(now, capacity, "epoch {epoch}");
assert_eq!(registry[Id(6)], 6 + epoch);
}
for i in 7..15 {
registry.push(i);
}
assert_eq!(registry.chunks.len(), chunks + 1);
assert_eq!(registry[Id(14)], 14);
}
#[test]
fn locate_matches_doubling_layout() {
assert_eq!(locate(0), (0, 0)); assert_eq!(locate(1), (1, 0)); assert_eq!(locate(2), (1, 1));
assert_eq!(locate(3), (2, 0)); assert_eq!(locate(6), (2, 3));
assert_eq!(locate(7), (3, 0)); }
#[test]
fn push_index_iter_across_chunks() {
let mut reg = Registry::<Id, usize>::default();
let ids: Vec<Id> = (0..1000).map(|v| reg.push(v)).collect();
assert_eq!(reg.len(), 1000);
for (i, &id) in ids.iter().enumerate() {
assert_eq!(usize::from(id), i);
assert_eq!(reg[id], i);
}
let seen: Vec<usize> = reg.iter().map(|item| *item.inner).collect();
assert_eq!(seen, (0..1000).collect::<Vec<_>>());
}
#[test]
fn element_address_is_stable_across_pushes() {
let mut reg = Registry::<Id, usize>::default();
let first = reg.push(42);
let addr = ®[first] as *const usize;
for v in 0..10_000 {
reg.push(v);
}
assert_eq!(
®[first] as *const usize, addr,
"address moved after growth"
);
assert_eq!(reg[first], 42);
}
#[test]
fn replace_swaps_in_place() {
let mut reg = Registry::<Id, i32>::default();
let a = reg.push(1);
let b = reg.push(2);
let addr_b = ®[b] as *const i32;
let old = reg.replace(b, 99);
assert_eq!(old, 2);
assert_eq!(reg[b], 99);
assert_eq!(reg[a], 1, "other ids untouched");
assert_eq!(®[b] as *const i32, addr_b, "slot address stable");
assert_eq!(reg.len(), 2, "len unchanged");
}
#[test]
#[cfg(debug_assertions)]
#[should_panic(expected = "registry index 0 out of bounds for length 0")]
fn indexing_reports_requested_index_and_length() {
let reg = Registry::<Id, usize>::default();
let _ = reg[Id::from(0)];
}
#[test]
fn select_mut_disjoint_in_input_order() {
let mut reg = Registry::<Id, usize>::default();
let ids: Vec<Id> = (0..100).map(|v| reg.push(v)).collect();
let picked = [ids[7], ids[0], ids[63], ids[64], ids[2]];
let refs = reg.select_mut(&picked);
assert_eq!(refs.len(), picked.len());
assert_eq!(
refs.iter().map(|r| **r).collect::<Vec<_>>(),
vec![7, 0, 63, 64, 2]
);
for r in refs {
*r += 1000;
}
for &id in &picked {
assert_eq!(reg[id], usize::from(id) + 1000);
}
assert_eq!(reg[ids[1]], 1);
}
#[test]
#[should_panic(expected = "duplicate id")]
fn select_mut_rejects_duplicates() {
let mut reg = Registry::<Id, usize>::default();
let a = reg.push(1);
reg.push(2);
let _ = reg.select_mut(&[a, a]);
}
#[test]
fn rebuild_from_flat_sequence() {
let reg: Registry<Id, i32> = (0..300).collect();
let flat: Vec<i32> = reg.iter().map(|item| *item.inner).collect();
let back: Registry<Id, i32> = flat.into_iter().collect();
assert_eq!(reg, back);
assert_eq!(back.len(), 300);
}
#[test]
fn clone_preserves_stability() {
let mut reg = Registry::<Id, usize>::default();
for v in 0..5 {
reg.push(v); }
let mut cloned = reg.clone();
assert_eq!(reg, cloned);
let last = Id::from(4);
let addr = &cloned[last] as *const usize;
cloned.push(99); assert_eq!(
&cloned[last] as *const usize, addr,
"clone's chunk reallocated"
);
}
#[test]
fn public_views_iterators_and_serialization_preserve_values() {
let mut reg = Registry::<Id, i32>::default();
assert!(reg.is_empty());
let first = reg.push(10);
let second = reg.push(20);
assert_eq!(reg.get(first).id, first);
assert_eq!(**reg.get(second), 20);
let mut mutable = reg.get_mut(first);
**mutable += 5;
let immutable = mutable.immutable();
assert_eq!(immutable.id, first);
assert_eq!(**immutable, 15);
assert_eq!(format!("{immutable}"), "15");
assert_eq!(format!("{immutable:?}"), "15");
let mut borrowed = (®).into_iter();
assert_eq!(borrowed.size_hint(), (0, None));
assert_eq!(borrowed.next().unwrap().id, first);
assert_eq!(borrowed.next().unwrap().id, second);
assert!(borrowed.next().is_none());
let owned: Vec<_> = reg.clone().into_iter().map(|entry| entry.inner).collect();
assert_eq!(owned, [15, 20]);
assert_eq!(format!("{reg:?}"), "[15, 20]");
let bytes = bincode::serde::encode_to_vec(®, bincode::config::standard()).unwrap();
let (decoded, used): (Registry<Id, i32>, _) =
bincode::serde::decode_from_slice(&bytes, bincode::config::standard()).unwrap();
assert_eq!(used, bytes.len());
assert_eq!(decoded, reg);
}
}