use alloc::sync::Arc;
use alloc::vec::Vec;
use core::ops::Index;
const SHIFT: u32 = 5;
const BRANCH: usize = 1 << SHIFT; const MASK: usize = BRANCH - 1;
#[derive(Debug, Clone)]
enum Node<T> {
Internal(Vec<Arc<Node<T>>>),
Leaf(Vec<T>),
}
#[derive(Debug)]
pub struct PersistentVec<T> {
root: Arc<Node<T>>,
tail: Arc<Vec<T>>,
len: usize,
shift: u32,
}
impl<T> Default for PersistentVec<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> Clone for PersistentVec<T> {
fn clone(&self) -> Self {
Self {
root: self.root.clone(),
tail: self.tail.clone(),
len: self.len,
shift: self.shift,
}
}
}
impl<T: PartialEq> PartialEq for PersistentVec<T> {
fn eq(&self, other: &Self) -> bool {
self.len == other.len && self.iter().eq(other.iter())
}
}
impl<T: Eq> Eq for PersistentVec<T> {}
impl<T> PersistentVec<T> {
#[must_use]
pub fn new() -> Self {
Self {
root: Arc::new(Node::Internal(Vec::new())),
tail: Arc::new(Vec::new()),
len: 0,
shift: SHIFT,
}
}
#[must_use]
pub const fn len(&self) -> usize {
self.len
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.len == 0
}
pub fn get(&self, i: usize) -> Option<&T> {
if i >= self.len {
return None;
}
let trie_size = self.len - self.tail.len();
if i >= trie_size {
return self.tail.get(i - trie_size);
}
let mut node: &Arc<Node<T>> = &self.root;
let mut shift = self.shift;
loop {
match &**node {
Node::Leaf(elems) => {
return elems.get(i & MASK);
}
Node::Internal(children) => {
let sub_idx = (i >> shift) & MASK;
node = children.get(sub_idx)?;
shift = shift.saturating_sub(SHIFT);
}
}
}
}
pub fn iter(&self) -> Iter<'_, T> {
Iter { pv: self, pos: 0 }
}
}
impl<'a, T> IntoIterator for &'a PersistentVec<T> {
type Item = &'a T;
type IntoIter = Iter<'a, T>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<T> Index<usize> for PersistentVec<T> {
type Output = T;
fn index(&self, i: usize) -> &T {
self.get(i).expect("PersistentVec index out of bounds")
}
}
impl<T: Clone> PersistentVec<T> {
#[must_use]
pub fn push(&self, x: T) -> Self {
if self.tail.len() < BRANCH {
let mut new_tail = (*self.tail).clone();
new_tail.push(x);
return Self {
root: self.root.clone(),
tail: Arc::new(new_tail),
len: self.len + 1,
shift: self.shift,
};
}
let leaf: Arc<Node<T>> = Arc::new(Node::Leaf((*self.tail).clone()));
let old_trie_size = self.len - BRANCH; let trie_capacity: usize = 1usize << (self.shift + SHIFT);
let needs_grow = old_trie_size + BRANCH > trie_capacity;
let (new_root, new_shift) = if needs_grow {
let internal_levels_above_leaf = self.shift / SHIFT;
let new_branch = new_path(internal_levels_above_leaf, leaf);
let new_root = Arc::new(Node::Internal(alloc::vec![self.root.clone(), new_branch]));
(new_root, self.shift + SHIFT)
} else {
(
push_leaf_into_node(&self.root, self.shift, old_trie_size, leaf),
self.shift,
)
};
Self {
root: new_root,
tail: Arc::new(alloc::vec![x]),
len: self.len + 1,
shift: new_shift,
}
}
pub fn push_mut(&mut self, x: T) {
if self.tail.len() < BRANCH {
let tail = Arc::make_mut(&mut self.tail);
tail.push(x);
self.len += 1;
return;
}
let old_tail_arc = core::mem::replace(&mut self.tail, Arc::new(alloc::vec![x]));
let old_tail_vec: Vec<T> =
Arc::try_unwrap(old_tail_arc).unwrap_or_else(|arc| (*arc).clone());
let leaf: Arc<Node<T>> = Arc::new(Node::Leaf(old_tail_vec));
let old_trie_size = self.len - BRANCH;
let trie_capacity: usize = 1usize << (self.shift + SHIFT);
let needs_grow = old_trie_size + BRANCH > trie_capacity;
if needs_grow {
let internal_levels = self.shift / SHIFT;
let new_branch = new_path(internal_levels, leaf);
self.root = Arc::new(Node::Internal(alloc::vec![self.root.clone(), new_branch]));
self.shift += SHIFT;
} else {
self.root = push_leaf_into_node(&self.root, self.shift, old_trie_size, leaf);
}
self.len += 1;
}
#[must_use]
pub fn set(&self, i: usize, x: T) -> Option<Self> {
if i >= self.len {
return None;
}
let trie_size = self.len - self.tail.len();
if i >= trie_size {
let mut new_tail: Vec<T> = (*self.tail).clone();
new_tail[i - trie_size] = x;
return Some(Self {
root: self.root.clone(),
tail: Arc::new(new_tail),
len: self.len,
shift: self.shift,
});
}
let new_root = set_in_trie(&self.root, self.shift, i, x);
Some(Self {
root: new_root,
tail: self.tail.clone(),
len: self.len,
shift: self.shift,
})
}
pub fn get_mut(&mut self, i: usize) -> Option<&mut T> {
if i >= self.len {
return None;
}
let trie_size = self.len - self.tail.len();
if i >= trie_size {
let tail = Arc::make_mut(&mut self.tail);
return tail.get_mut(i - trie_size);
}
get_mut_in_trie(&mut self.root, self.shift, i)
}
}
fn push_leaf_into_node<T: Clone>(
node: &Arc<Node<T>>,
shift: u32,
trie_index: usize,
leaf: Arc<Node<T>>,
) -> Arc<Node<T>> {
let sub_idx = (trie_index >> shift) & MASK;
let Node::Internal(children) = &**node else {
debug_assert!(false, "push_leaf_into_node hit a Leaf — shift bug");
return node.clone();
};
let mut new_children: Vec<Arc<Node<T>>> = children.clone();
if shift == SHIFT {
debug_assert!(
sub_idx == new_children.len(),
"leaves are pushed sequentially; sub_idx {} != next slot {}",
sub_idx,
new_children.len()
);
new_children.push(leaf);
} else {
let child: Arc<Node<T>> = if sub_idx < new_children.len() {
push_leaf_into_node(&new_children[sub_idx], shift - SHIFT, trie_index, leaf)
} else {
let internal_levels_above_leaf = (shift / SHIFT) - 1;
new_path(internal_levels_above_leaf, leaf)
};
if sub_idx < new_children.len() {
new_children[sub_idx] = child;
} else {
new_children.push(child);
}
}
Arc::new(Node::Internal(new_children))
}
fn new_path<T>(internal_levels: u32, leaf: Arc<Node<T>>) -> Arc<Node<T>> {
let mut node = leaf;
for _ in 0..internal_levels {
node = Arc::new(Node::Internal(alloc::vec![node]));
}
node
}
fn set_in_trie<T: Clone>(node: &Arc<Node<T>>, shift: u32, i: usize, x: T) -> Arc<Node<T>> {
match &**node {
Node::Leaf(elems) => {
let mut new_elems = elems.clone();
new_elems[i & MASK] = x;
Arc::new(Node::Leaf(new_elems))
}
Node::Internal(children) => {
let sub_idx = (i >> shift) & MASK;
let new_child = set_in_trie(&children[sub_idx], shift - SHIFT, i, x);
let mut new_children = children.clone();
new_children[sub_idx] = new_child;
Arc::new(Node::Internal(new_children))
}
}
}
fn get_mut_in_trie<T: Clone>(node: &mut Arc<Node<T>>, shift: u32, i: usize) -> Option<&mut T> {
match Arc::make_mut(node) {
Node::Leaf(elems) => elems.get_mut(i & MASK),
Node::Internal(children) => {
let sub_idx = (i >> shift) & MASK;
let child = children.get_mut(sub_idx)?;
get_mut_in_trie(child, shift - SHIFT, i)
}
}
}
#[derive(Debug)]
pub struct Iter<'a, T> {
pv: &'a PersistentVec<T>,
pos: usize,
}
impl<'a, T> Iterator for Iter<'a, T> {
type Item = &'a T;
fn next(&mut self) -> Option<&'a T> {
if self.pos >= self.pv.len {
return None;
}
let v = self.pv.get(self.pos)?;
self.pos += 1;
Some(v)
}
fn size_hint(&self) -> (usize, Option<usize>) {
let remaining = self.pv.len.saturating_sub(self.pos);
(remaining, Some(remaining))
}
}
impl<T> ExactSizeIterator for Iter<'_, T> {}
#[cfg(test)]
impl<T> PersistentVec<T> {
pub(crate) fn shares_storage_with(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.root, &other.root) && Arc::ptr_eq(&self.tail, &other.tail)
}
}
#[cfg(test)]
#[allow(
clippy::cast_possible_truncation,
clippy::cast_possible_wrap,
clippy::cast_sign_loss,
clippy::cast_lossless,
clippy::needless_range_loop,
clippy::items_after_statements,
clippy::manual_range_patterns,
clippy::unreadable_literal,
clippy::similar_names
)]
mod tests {
use super::*;
#[test]
fn empty_vec_is_empty() {
let pv: PersistentVec<u64> = PersistentVec::new();
assert_eq!(pv.len(), 0);
assert!(pv.is_empty());
assert!(pv.get(0).is_none());
}
#[test]
fn push_single_fits_in_tail() {
let pv: PersistentVec<u64> = PersistentVec::new().push(42);
assert_eq!(pv.len(), 1);
assert_eq!(pv.get(0), Some(&42));
assert!(pv.get(1).is_none());
}
#[test]
fn push_fills_tail_then_incorporates() {
let mut pv: PersistentVec<u64> = PersistentVec::new();
for i in 0..40_u64 {
pv = pv.push(i);
}
for i in 0..40_u64 {
assert_eq!(pv.get(i as usize), Some(&i), "mismatch at {i}");
}
assert!(pv.get(40).is_none());
}
#[test]
fn push_crosses_root_overflow_boundary() {
let mut pv: PersistentVec<u64> = PersistentVec::new();
for i in 0..1100_u64 {
pv = pv.push(i);
}
for i in 0..1100_u64 {
assert_eq!(pv.get(i as usize), Some(&i), "mismatch at {i}");
}
}
#[test]
fn push_crosses_second_grow_boundary() {
let mut pv: PersistentVec<u64> = PersistentVec::new();
for i in 0..33_000_u64 {
pv = pv.push(i);
}
let probes = [0_usize, 1, 31, 32, 1023, 1024, 1056, 32_767, 32_768, 32_999];
for &p in &probes {
assert_eq!(pv.get(p), Some(&(p as u64)), "mismatch at {p}");
}
assert!(pv.get(33_000).is_none());
}
#[test]
fn clone_then_push_preserves_original() {
let mut a: PersistentVec<u64> = PersistentVec::new();
for i in 0..50_u64 {
a = a.push(i);
}
let b = a.clone();
let b = b.push(999);
assert_eq!(a.len(), 50);
assert_eq!(b.len(), 51);
assert_eq!(a.get(50), None);
assert_eq!(b.get(50), Some(&999));
for i in 0..50_usize {
assert_eq!(a.get(i), Some(&(i as u64)));
assert_eq!(b.get(i), Some(&(i as u64)));
}
}
#[test]
fn set_rewrites_element_in_tail() {
let pv: PersistentVec<u64> = PersistentVec::new()
.push(10)
.push(20)
.push(30)
.set(1, 200)
.unwrap();
assert_eq!(pv.get(0), Some(&10));
assert_eq!(pv.get(1), Some(&200));
assert_eq!(pv.get(2), Some(&30));
}
#[test]
fn set_rewrites_element_in_trie() {
let mut pv: PersistentVec<u64> = PersistentVec::new();
for i in 0..40_u64 {
pv = pv.push(i);
}
let pv2 = pv.set(0, 9999).unwrap();
assert_eq!(pv2.get(0), Some(&9999));
assert_eq!(pv.get(0), Some(&0), "set must not mutate original");
assert_eq!(pv2.get(39), Some(&39));
}
#[test]
fn set_out_of_bounds_is_none() {
let pv: PersistentVec<u64> = PersistentVec::new().push(1);
assert!(pv.set(5, 99).is_none());
}
#[test]
fn iter_matches_get_for_full_walk() {
let mut pv: PersistentVec<u64> = PersistentVec::new();
for i in 0..200_u64 {
pv = pv.push(i * 7);
}
let via_iter: Vec<u64> = pv.iter().copied().collect();
let via_get: Vec<u64> = (0..pv.len()).map(|i| *pv.get(i).unwrap()).collect();
assert_eq!(via_iter, via_get);
assert_eq!(via_iter.len(), 200);
assert_eq!(via_iter[199], 199 * 7);
}
#[test]
fn iter_size_hint_exact() {
let mut pv: PersistentVec<u64> = PersistentVec::new();
for i in 0..15_u64 {
pv = pv.push(i);
}
let it = pv.iter();
assert_eq!(it.size_hint(), (15, Some(15)));
assert_eq!(it.count(), 15);
}
struct Splitmix(u64);
impl Splitmix {
fn new(seed: u64) -> Self {
Self(seed)
}
fn next(&mut self) -> u64 {
self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
let mut x = self.0;
x = (x ^ (x >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
x = (x ^ (x >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
x ^ (x >> 31)
}
}
#[test]
fn fuzz_oracle_against_vec_u64() {
let mut pv: PersistentVec<u64> = PersistentVec::new();
let mut oracle: Vec<u64> = Vec::new();
let mut rng = Splitmix::new(0xC0FFEE_u64);
const STEPS: usize = 100_000;
for step in 0..STEPS {
let r = rng.next();
let op = r % 4; match (op, oracle.len()) {
(0 | 1 | 2, _) | (_, 0) => {
let val = rng.next();
pv = pv.push(val);
oracle.push(val);
}
(3, n) => {
let idx = (rng.next() as usize) % n;
let val = rng.next();
pv = pv.set(idx, val).expect("in-bounds set");
oracle[idx] = val;
}
_ => unreachable!(),
}
assert_eq!(pv.len(), oracle.len(), "len drift @ step {step}");
if !oracle.is_empty() {
assert_eq!(pv.get(0), oracle.first(), "head drift @ step {step}");
assert_eq!(
pv.get(oracle.len() - 1),
oracle.last(),
"tail drift @ step {step}"
);
let probe = (rng.next() as usize) % oracle.len();
assert_eq!(
pv.get(probe),
Some(&oracle[probe]),
"interior drift @ step {step}, probe {probe}"
);
}
}
for i in 0..oracle.len() {
assert_eq!(pv.get(i), Some(&oracle[i]), "final drift at {i}");
}
let via_iter: Vec<u64> = pv.iter().copied().collect();
assert_eq!(via_iter, oracle, "iter drift");
}
#[test]
fn fuzz_oracle_clone_isolation() {
let mut a: PersistentVec<u64> = PersistentVec::new();
let mut oracle_a: Vec<u64> = Vec::new();
let mut rng = Splitmix::new(0xDECAFBAD_u64);
for _ in 0..2_000 {
let v = rng.next();
a = a.push(v);
oracle_a.push(v);
}
let mut b = a.clone();
let mut oracle_b = oracle_a.clone();
let mut c = a.clone();
let mut oracle_c = oracle_a.clone();
for _ in 0..500 {
let v = rng.next();
b = b.push(v);
oracle_b.push(v);
}
for _ in 0..300 {
let idx = (rng.next() as usize) % oracle_c.len();
let v = rng.next();
c = c.set(idx, v).expect("in-bounds");
oracle_c[idx] = v;
}
for (i, &want) in oracle_a.iter().enumerate() {
assert_eq!(a.get(i), Some(&want), "A drift at {i}");
}
for (i, &want) in oracle_b.iter().enumerate() {
assert_eq!(b.get(i), Some(&want), "B drift at {i}");
}
for (i, &want) in oracle_c.iter().enumerate() {
assert_eq!(c.get(i), Some(&want), "C drift at {i}");
}
assert_eq!(a.len(), oracle_a.len());
assert_eq!(b.len(), oracle_b.len());
assert_eq!(c.len(), oracle_c.len());
}
#[test]
fn fuzz_oracle_push_mut_against_vec_u64() {
let mut pv: PersistentVec<u64> = PersistentVec::new();
let mut oracle: Vec<u64> = Vec::new();
let mut rng = Splitmix::new(0xFEEDFACE_u64);
const STEPS: usize = 100_000;
for step in 0..STEPS {
let val = rng.next();
pv.push_mut(val);
oracle.push(val);
assert_eq!(pv.len(), oracle.len(), "len drift @ step {step}");
if step % 1024 == 0 {
let probe = (rng.next() as usize) % oracle.len();
assert_eq!(
pv.get(probe),
Some(&oracle[probe]),
"interior drift @ step {step}, probe {probe}"
);
}
}
for i in 0..oracle.len() {
assert_eq!(pv.get(i), Some(&oracle[i]), "final drift at {i}");
}
}
#[test]
fn push_mut_does_not_disturb_cloned_handle() {
let mut a: PersistentVec<u64> = PersistentVec::new();
for i in 0..200_u64 {
a.push_mut(i);
}
let b = a.clone();
for i in 200_u64..500 {
a.push_mut(i);
}
assert_eq!(b.len(), 200);
for i in 0..200_u64 {
assert_eq!(b.get(i as usize), Some(&i), "B drift at {i}");
}
assert!(b.get(200).is_none());
assert_eq!(a.len(), 500);
for i in 0..500_u64 {
assert_eq!(a.get(i as usize), Some(&i), "A drift at {i}");
}
}
#[test]
fn push_clone_arc_count_stays_constant_in_old_handle() {
let mut a: PersistentVec<u64> = PersistentVec::new();
for i in 0..200_u64 {
a = a.push(i);
}
let snapshots: Vec<PersistentVec<u64>> = (0..5).map(|_| a.clone()).collect();
drop(snapshots);
for i in 0..200_u64 {
assert_eq!(a.get(i as usize), Some(&i));
}
assert_eq!(a.len(), 200);
}
}