#![allow(
clippy::doc_markdown,
clippy::many_single_char_names,
clippy::type_complexity
)]
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::ops::Bound;
const ORDER: usize = 8;
const MAX_ENTRIES: usize = ORDER - 1; const MAX_CHILDREN: usize = ORDER;
#[derive(Debug)]
enum BNode<K, V> {
Leaf {
entries: Vec<(K, V)>,
},
Internal {
entries: Vec<(K, V)>,
children: Vec<Arc<BNode<K, V>>>,
},
}
impl<K: Clone, V: Clone> Clone for BNode<K, V> {
fn clone(&self) -> Self {
match self {
Self::Leaf { entries } => Self::Leaf {
entries: entries.clone(),
},
Self::Internal { entries, children } => Self::Internal {
entries: entries.clone(),
children: children.clone(),
},
}
}
}
#[derive(Debug)]
pub struct PersistentBTreeMap<K, V> {
root: Arc<BNode<K, V>>,
len: usize,
}
impl<K, V> Default for PersistentBTreeMap<K, V> {
fn default() -> Self {
Self::new()
}
}
impl<K, V> Clone for PersistentBTreeMap<K, V> {
fn clone(&self) -> Self {
Self {
root: self.root.clone(),
len: self.len,
}
}
}
impl<K: PartialEq, V: PartialEq> PartialEq for PersistentBTreeMap<K, V>
where
K: Ord,
{
fn eq(&self, other: &Self) -> bool {
self.len == other.len && self.iter().eq(other.iter())
}
}
impl<K: Eq + Ord, V: Eq> Eq for PersistentBTreeMap<K, V> {}
impl<K, V> PersistentBTreeMap<K, V> {
#[must_use]
pub fn new() -> Self {
Self {
root: Arc::new(BNode::Leaf {
entries: Vec::new(),
}),
len: 0,
}
}
#[must_use]
pub const fn len(&self) -> usize {
self.len
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.len == 0
}
}
enum FoundOrDescend {
Found(usize),
Descend(usize),
}
#[inline]
fn linear_find_entry<K: Ord, V>(entries: &[(K, V)], key: &K) -> Option<usize> {
for (i, (k, _)) in entries.iter().enumerate() {
match k.cmp(key) {
core::cmp::Ordering::Equal => return Some(i),
core::cmp::Ordering::Greater => return None,
core::cmp::Ordering::Less => continue,
}
}
None
}
#[inline]
fn linear_position_internal<K: Ord, V>(entries: &[(K, V)], key: &K) -> FoundOrDescend {
for (i, (k, _)) in entries.iter().enumerate() {
match k.cmp(key) {
core::cmp::Ordering::Equal => return FoundOrDescend::Found(i),
core::cmp::Ordering::Greater => return FoundOrDescend::Descend(i),
core::cmp::Ordering::Less => continue,
}
}
FoundOrDescend::Descend(entries.len())
}
impl<K: Ord, V> PersistentBTreeMap<K, V> {
pub fn get(&self, key: &K) -> Option<&V> {
let mut node: &Arc<BNode<K, V>> = &self.root;
loop {
match &**node {
BNode::Leaf { entries } => {
return linear_find_entry(entries, key).map(|i| &entries[i].1);
}
BNode::Internal { entries, children } => {
match linear_position_internal(entries, key) {
FoundOrDescend::Found(i) => return Some(&entries[i].1),
FoundOrDescend::Descend(i) => {
node = &children[i];
}
}
}
}
}
}
pub fn predecessor(&self, key: &K) -> Option<(&K, &V)> {
let mut node: &Arc<BNode<K, V>> = &self.root;
let mut best: Option<(&K, &V)> = None;
loop {
match &**node {
BNode::Leaf { entries } => {
let i = entries.partition_point(|e| &e.0 < key);
if i > 0 {
let (k, v) = &entries[i - 1];
best = Some((k, v));
}
return best;
}
BNode::Internal { entries, children } => {
let i = entries.partition_point(|e| &e.0 < key);
if i > 0 {
let (k, v) = &entries[i - 1];
best = Some((k, v));
}
node = &children[i];
}
}
}
}
#[must_use]
pub fn from_sorted(entries: Vec<(K, V)>) -> Self
where
K: Ord + Clone,
V: Clone,
{
#[cfg(debug_assertions)]
for w in entries.windows(2) {
debug_assert!(w[0].0 < w[1].0, "from_sorted requires strictly sorted keys");
}
let len = entries.len();
if len <= MAX_ENTRIES {
return Self {
root: Arc::new(BNode::Leaf { entries }),
len,
};
}
let mut nodes: Vec<Arc<BNode<K, V>>> = Vec::new();
let mut seps: Vec<(K, V)> = Vec::new();
{
let g = (len + 1).div_ceil(MAX_ENTRIES + 1);
let mut it = entries.into_iter();
let mut remaining = len;
for gi in 0..g {
let groups_left = g - gi;
let seps_left = groups_left - 1;
let take = (remaining - seps_left).div_ceil(groups_left);
let leaf: Vec<(K, V)> = (&mut it).take(take).collect();
remaining -= leaf.len();
nodes.push(Arc::new(BNode::Leaf { entries: leaf }));
if gi + 1 < g {
let sep = it.next().expect("separator exists");
remaining -= 1;
seps.push(sep);
}
}
}
while nodes.len() > 1 {
let g = nodes.len().div_ceil(MAX_CHILDREN);
let per = nodes.len().div_ceil(g);
let mut up_nodes: Vec<Arc<BNode<K, V>>> = Vec::with_capacity(g);
let mut up_seps: Vec<(K, V)> = Vec::with_capacity(g - 1);
let mut node_it = nodes.into_iter();
let mut sep_it = seps.into_iter();
let mut children: Vec<Arc<BNode<K, V>>> = Vec::with_capacity(per);
let mut inner: Vec<(K, V)> = Vec::with_capacity(per - 1);
loop {
match node_it.next() {
Some(n) => {
if !children.is_empty() {
let sep = sep_it.next().expect("separator per boundary");
if children.len() < per {
inner.push(sep);
} else {
up_nodes.push(Arc::new(BNode::Internal {
entries: core::mem::take(&mut inner),
children: core::mem::take(&mut children),
}));
up_seps.push(sep);
}
}
children.push(n);
}
None => {
up_nodes.push(Arc::new(BNode::Internal {
entries: inner,
children,
}));
break;
}
}
}
nodes = up_nodes;
seps = up_seps;
}
Self {
root: nodes.pop().expect("one root"),
len,
}
}
pub fn iter(&self) -> Iter<'_, K, V> {
let mut stack: Vec<(&Arc<BNode<K, V>>, usize)> = Vec::with_capacity(8);
stack.push((&self.root, 0));
Iter { stack }
}
pub fn iter_rev(&self) -> IterRev<'_, K, V> {
let mut stack: Vec<(&Arc<BNode<K, V>>, usize)> = Vec::with_capacity(8);
stack.push((&self.root, 1));
IterRev { stack }
}
pub fn range<'a>(&'a self, lo: Bound<&K>, hi: Bound<&K>) -> RangeIter<'a, K, V>
where
K: Clone,
{
let mut stack: Vec<(&'a Arc<BNode<K, V>>, usize)> = Vec::with_capacity(8);
let mut node = &self.root;
loop {
match &**node {
BNode::Leaf { entries } => {
stack.push((node, lower_index(entries, lo)));
break;
}
BNode::Internal { entries, children } => {
let i = lower_index(entries, lo);
stack.push((node, 2 * i + 1));
node = &children[i];
}
}
}
let (hi_key, hi_incl) = match hi {
Bound::Unbounded => (None, false),
Bound::Included(k) => (Some(k.clone()), true),
Bound::Excluded(k) => (Some(k.clone()), false),
};
RangeIter {
inner: Iter { stack },
hi_key,
hi_incl,
done: false,
}
}
}
fn lower_index<K: Ord, V>(entries: &[(K, V)], lo: Bound<&K>) -> usize {
match lo {
Bound::Unbounded => 0,
Bound::Included(k) => entries.partition_point(|e| &e.0 < k),
Bound::Excluded(k) => entries.partition_point(|e| &e.0 <= k),
}
}
impl<K: Ord + Clone, V: Clone> PersistentBTreeMap<K, V> {
#[must_use]
pub fn insert(&self, key: K, value: V) -> (Self, Option<V>) {
let (new_left, split, prev_v) = insert_helper(&self.root, key, value);
let new_root = if let Some((right, median)) = split {
Arc::new(BNode::Internal {
entries: alloc::vec![median],
children: alloc::vec![new_left, right],
})
} else {
new_left
};
let new_len = if prev_v.is_none() {
self.len + 1
} else {
self.len
};
(
Self {
root: new_root,
len: new_len,
},
prev_v,
)
}
pub fn get_by<Q>(&self, key: &Q) -> Option<&V>
where
K: core::borrow::Borrow<Q>,
Q: Ord + ?Sized,
{
let mut node: &Arc<BNode<K, V>> = &self.root;
loop {
match &**node {
BNode::Leaf { entries } => {
return linear_find_entry_by(entries, key).map(|i| &entries[i].1);
}
BNode::Internal { entries, children } => {
match linear_position_internal_by(entries, key) {
FoundOrDescend::Found(i) => return Some(&entries[i].1),
FoundOrDescend::Descend(i) => node = &children[i],
}
}
}
}
}
pub fn get_mut_by<Q>(&mut self, key: &Q) -> Option<&mut V>
where
K: core::borrow::Borrow<Q> + Clone,
V: Clone,
Q: Ord + ?Sized,
{
get_mut_by_helper(&mut self.root, key)
}
pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
get_mut_helper(&mut self.root, key)
}
pub fn insert_mut(&mut self, key: K, value: V) -> Option<V> {
let (split, prev_v) = insert_transient_helper(&mut self.root, key, value);
if let Some((right, median)) = split {
let old_root = core::mem::replace(
&mut self.root,
Arc::new(BNode::Leaf {
entries: Vec::new(),
}),
);
self.root = Arc::new(BNode::Internal {
entries: alloc::vec![median],
children: alloc::vec![old_root, right],
});
}
if prev_v.is_none() {
self.len += 1;
}
prev_v
}
pub fn remove_mut(&mut self, key: &K) -> Option<V> {
let removed = remove_transient_helper(&mut self.root, key)?;
let collapse = match self.root.as_ref() {
BNode::Internal { entries, children } if entries.is_empty() => {
debug_assert_eq!(children.len(), 1);
children.first().cloned()
}
_ => None,
};
if let Some(only) = collapse {
self.root = only;
}
self.len -= 1;
Some(removed)
}
#[must_use]
pub fn remove(&self, key: &K) -> (Self, Option<V>) {
let mut next = self.clone();
let prev = next.remove_mut(key);
(next, prev)
}
}
const MIN_ENTRIES: usize = ORDER / 2 - 1;
impl<K, V> BNode<K, V> {
fn entry_count(&self) -> usize {
match self {
BNode::Leaf { entries } | BNode::Internal { entries, .. } => entries.len(),
}
}
}
fn remove_transient_helper<K: Ord + Clone, V: Clone>(
node: &mut Arc<BNode<K, V>>,
key: &K,
) -> Option<V> {
let (found, idx) = match node.as_ref() {
BNode::Leaf { entries } | BNode::Internal { entries, .. } => {
match entries.binary_search_by(|(ek, _)| ek.cmp(key)) {
Ok(i) => (true, i),
Err(i) => (false, i),
}
}
};
if !found && matches!(node.as_ref(), BNode::Leaf { .. }) {
return None;
}
let inner = Arc::make_mut(node);
match inner {
BNode::Leaf { entries } => Some(entries.remove(idx).1),
BNode::Internal { entries, children } => {
if found {
let pred = remove_max(&mut children[idx]);
let old = core::mem::replace(&mut entries[idx], pred);
fix_child(entries, children, idx);
Some(old.1)
} else {
let removed = remove_transient_helper(&mut children[idx], key)?;
fix_child(entries, children, idx);
Some(removed)
}
}
}
}
fn remove_max<K: Ord + Clone, V: Clone>(node: &mut Arc<BNode<K, V>>) -> (K, V) {
let inner = Arc::make_mut(node);
match inner {
BNode::Leaf { entries } => entries.pop().expect("a B-tree leaf is never empty"),
BNode::Internal { entries, children } => {
let last = children.len() - 1;
let kv = remove_max(&mut children[last]);
fix_child(entries, children, last);
kv
}
}
}
fn fix_child<K: Ord + Clone, V: Clone>(
entries: &mut Vec<(K, V)>,
children: &mut Vec<Arc<BNode<K, V>>>,
i: usize,
) {
if children[i].entry_count() >= MIN_ENTRIES {
return;
}
if i > 0 && children[i - 1].entry_count() > MIN_ENTRIES {
rotate_from_left(entries, children, i);
} else if i + 1 < children.len() && children[i + 1].entry_count() > MIN_ENTRIES {
rotate_from_right(entries, children, i);
} else if i > 0 {
merge_children(entries, children, i - 1);
} else {
merge_children(entries, children, i);
}
}
fn rotate_from_left<K: Ord + Clone, V: Clone>(
entries: &mut [(K, V)],
children: &mut [Arc<BNode<K, V>>],
i: usize,
) {
let (moved_entry, moved_child) = match Arc::make_mut(&mut children[i - 1]) {
BNode::Leaf { entries: le } => (le.pop().expect("sibling has entries to spare"), None),
BNode::Internal {
entries: le,
children: lc,
} => (
le.pop().expect("sibling has entries to spare"),
Some(
lc.pop()
.expect("internal node has entries.len()+1 children"),
),
),
};
let separator = core::mem::replace(&mut entries[i - 1], moved_entry);
match Arc::make_mut(&mut children[i]) {
BNode::Leaf { entries: ce } => {
debug_assert!(moved_child.is_none());
ce.insert(0, separator);
}
BNode::Internal {
entries: ce,
children: cc,
} => {
ce.insert(0, separator);
cc.insert(
0,
moved_child.expect("sibling of an internal node is internal"),
);
}
}
}
fn rotate_from_right<K: Ord + Clone, V: Clone>(
entries: &mut [(K, V)],
children: &mut [Arc<BNode<K, V>>],
i: usize,
) {
let (moved_entry, moved_child) = match Arc::make_mut(&mut children[i + 1]) {
BNode::Leaf { entries: re } => (re.remove(0), None),
BNode::Internal {
entries: re,
children: rc,
} => (re.remove(0), Some(rc.remove(0))),
};
let separator = core::mem::replace(&mut entries[i], moved_entry);
match Arc::make_mut(&mut children[i]) {
BNode::Leaf { entries: ce } => {
debug_assert!(moved_child.is_none());
ce.push(separator);
}
BNode::Internal {
entries: ce,
children: cc,
} => {
ce.push(separator);
cc.push(moved_child.expect("sibling of an internal node is internal"));
}
}
}
fn merge_children<K: Ord + Clone, V: Clone>(
entries: &mut Vec<(K, V)>,
children: &mut Vec<Arc<BNode<K, V>>>,
sep: usize,
) {
let separator = entries.remove(sep);
let right = children.remove(sep + 1);
let right = Arc::try_unwrap(right).unwrap_or_else(|shared| (*shared).clone());
match (Arc::make_mut(&mut children[sep]), right) {
(BNode::Leaf { entries: le }, BNode::Leaf { entries: re }) => {
le.push(separator);
le.extend(re);
}
(
BNode::Internal {
entries: le,
children: lc,
},
BNode::Internal {
entries: re,
children: rc,
},
) => {
le.push(separator);
le.extend(re);
lc.extend(rc);
}
(BNode::Leaf { entries: le }, BNode::Internal { .. })
| (BNode::Internal { entries: le, .. }, BNode::Leaf { .. }) => {
debug_assert!(false, "B-tree siblings must be at the same depth");
le.push(separator);
}
}
}
fn linear_find_entry_by<K, V, Q>(entries: &[(K, V)], key: &Q) -> Option<usize>
where
K: core::borrow::Borrow<Q>,
Q: Ord + ?Sized,
{
for (i, (k, _)) in entries.iter().enumerate() {
match k.borrow().cmp(key) {
core::cmp::Ordering::Equal => return Some(i),
core::cmp::Ordering::Greater => return None,
core::cmp::Ordering::Less => continue,
}
}
None
}
fn linear_position_internal_by<K, V, Q>(entries: &[(K, V)], key: &Q) -> FoundOrDescend
where
K: core::borrow::Borrow<Q>,
Q: Ord + ?Sized,
{
for (i, (k, _)) in entries.iter().enumerate() {
match k.borrow().cmp(key) {
core::cmp::Ordering::Equal => return FoundOrDescend::Found(i),
core::cmp::Ordering::Greater => return FoundOrDescend::Descend(i),
core::cmp::Ordering::Less => continue,
}
}
FoundOrDescend::Descend(entries.len())
}
fn get_mut_by_helper<'a, K, V, Q>(node: &'a mut Arc<BNode<K, V>>, key: &Q) -> Option<&'a mut V>
where
K: core::borrow::Borrow<Q> + Clone,
V: Clone,
Q: Ord + ?Sized,
{
match Arc::make_mut(node) {
BNode::Leaf { entries } => {
let i = linear_find_entry_by(entries, key)?;
Some(&mut entries[i].1)
}
BNode::Internal { entries, children } => match linear_position_internal_by(entries, key) {
FoundOrDescend::Found(i) => Some(&mut entries[i].1),
FoundOrDescend::Descend(i) => get_mut_by_helper(&mut children[i], key),
},
}
}
fn get_mut_helper<'a, K: Ord + Clone, V: Clone>(
node: &'a mut Arc<BNode<K, V>>,
key: &K,
) -> Option<&'a mut V> {
match Arc::make_mut(node) {
BNode::Leaf { entries } => {
let i = linear_find_entry(entries, key)?;
Some(&mut entries[i].1)
}
BNode::Internal { entries, children } => match linear_position_internal(entries, key) {
FoundOrDescend::Found(i) => Some(&mut entries[i].1),
FoundOrDescend::Descend(i) => get_mut_helper(&mut children[i], key),
},
}
}
fn insert_transient_helper<K: Ord + Clone, V: Clone>(
node: &mut Arc<BNode<K, V>>,
k: K,
v: V,
) -> (Option<(Arc<BNode<K, V>>, (K, V))>, Option<V>) {
let inner = Arc::make_mut(node);
match inner {
BNode::Leaf { entries } => {
let pos = entries.binary_search_by(|(ek, _)| ek.cmp(&k));
let prev_v = match pos {
Ok(idx) => Some(core::mem::replace(&mut entries[idx].1, v)),
Err(idx) => {
entries.insert(idx, (k, v));
None
}
};
if entries.len() <= MAX_ENTRIES {
return (None, prev_v);
}
let mid = entries.len() / 2;
let right_entries = entries.split_off(mid + 1);
let median = entries.pop().expect("mid was in-bounds");
let right = Arc::new(BNode::Leaf {
entries: right_entries,
});
(Some((right, median)), prev_v)
}
BNode::Internal { entries, children } => {
let pos = entries.binary_search_by(|(ek, _)| ek.cmp(&k));
match pos {
Ok(idx) => {
let prev_v = core::mem::replace(&mut entries[idx].1, v);
(None, Some(prev_v))
}
Err(idx) => {
let (split, prev_v) = insert_transient_helper(&mut children[idx], k, v);
if let Some((right_sibling, median)) = split {
entries.insert(idx, median);
children.insert(idx + 1, right_sibling);
}
if children.len() <= MAX_CHILDREN {
return (None, prev_v);
}
let mid = entries.len() / 2;
let right_entries = entries.split_off(mid + 1);
let median = entries.pop().expect("mid was in-bounds");
let right_children = children.split_off(mid + 1);
let right = Arc::new(BNode::Internal {
entries: right_entries,
children: right_children,
});
(Some((right, median)), prev_v)
}
}
}
}
}
fn insert_helper<K: Ord + Clone, V: Clone>(
node: &Arc<BNode<K, V>>,
k: K,
v: V,
) -> (
Arc<BNode<K, V>>,
Option<(Arc<BNode<K, V>>, (K, V))>,
Option<V>,
) {
match &**node {
BNode::Leaf { entries } => {
let pos = entries.binary_search_by(|(ek, _)| ek.cmp(&k));
let mut new_entries = entries.clone();
let prev_v = match pos {
Ok(idx) => Some(core::mem::replace(&mut new_entries[idx].1, v)),
Err(idx) => {
new_entries.insert(idx, (k, v));
None
}
};
if new_entries.len() <= MAX_ENTRIES {
return (
Arc::new(BNode::Leaf {
entries: new_entries,
}),
None,
prev_v,
);
}
let mid = new_entries.len() / 2; let right_entries = new_entries.split_off(mid + 1);
let median = new_entries.pop().expect("mid was in-bounds");
let left = Arc::new(BNode::Leaf {
entries: new_entries,
});
let right = Arc::new(BNode::Leaf {
entries: right_entries,
});
(left, Some((right, median)), prev_v)
}
BNode::Internal { entries, children } => {
let pos = entries.binary_search_by(|(ek, _)| ek.cmp(&k));
match pos {
Ok(idx) => {
let mut new_entries = entries.clone();
let prev_v = core::mem::replace(&mut new_entries[idx].1, v);
(
Arc::new(BNode::Internal {
entries: new_entries,
children: children.clone(),
}),
None,
Some(prev_v),
)
}
Err(idx) => {
let (new_child, split, prev_v) = insert_helper(&children[idx], k, v);
let mut new_entries = entries.clone();
let mut new_children = children.clone();
new_children[idx] = new_child;
if let Some((right_sibling, median)) = split {
new_entries.insert(idx, median);
new_children.insert(idx + 1, right_sibling);
}
if new_children.len() <= MAX_CHILDREN {
return (
Arc::new(BNode::Internal {
entries: new_entries,
children: new_children,
}),
None,
prev_v,
);
}
let mid = new_entries.len() / 2; let right_entries = new_entries.split_off(mid + 1);
let median = new_entries.pop().expect("mid was in-bounds");
let right_children = new_children.split_off(mid + 1);
let left = Arc::new(BNode::Internal {
entries: new_entries,
children: new_children,
});
let right = Arc::new(BNode::Internal {
entries: right_entries,
children: right_children,
});
(left, Some((right, median)), prev_v)
}
}
}
}
}
#[derive(Debug)]
pub struct Iter<'a, K, V> {
stack: Vec<(&'a Arc<BNode<K, V>>, usize)>,
}
impl<'a, K, V> Iterator for Iter<'a, K, V> {
type Item = (&'a K, &'a V);
fn next(&mut self) -> Option<(&'a K, &'a V)> {
loop {
let (node, idx) = *self.stack.last()?;
match &**node {
BNode::Leaf { entries } => {
if idx < entries.len() {
let (k, v) = &entries[idx];
self.stack.last_mut().unwrap().1 = idx + 1;
return Some((k, v));
}
self.stack.pop();
}
BNode::Internal { entries, children } => {
let phase = idx & 1;
let slot = idx >> 1;
if phase == 0 {
if slot < children.len() {
self.stack.last_mut().unwrap().1 = idx + 1;
self.stack.push((&children[slot], 0));
continue;
}
self.stack.pop();
} else {
if slot < entries.len() {
let (k, v) = &entries[slot];
self.stack.last_mut().unwrap().1 = idx + 1;
return Some((k, v));
}
self.stack.pop();
}
}
}
}
}
}
#[derive(Debug)]
pub struct RangeIter<'a, K, V> {
inner: Iter<'a, K, V>,
hi_key: Option<K>,
hi_incl: bool,
done: bool,
}
impl<'a, K: Ord, V> Iterator for RangeIter<'a, K, V> {
type Item = (&'a K, &'a V);
fn next(&mut self) -> Option<(&'a K, &'a V)> {
if self.done {
return None;
}
let (k, v) = self.inner.next()?;
if let Some(h) = &self.hi_key {
let past = if self.hi_incl { k > h } else { k >= h };
if past {
self.done = true;
return None;
}
}
Some((k, v))
}
}
#[derive(Debug)]
pub struct IterRev<'a, K, V> {
stack: Vec<(&'a Arc<BNode<K, V>>, usize)>,
}
impl<'a, K, V> Iterator for IterRev<'a, K, V> {
type Item = (&'a K, &'a V);
fn next(&mut self) -> Option<(&'a K, &'a V)> {
loop {
let (node, pos) = *self.stack.last()?;
match &**node {
BNode::Leaf { entries } => {
if pos <= entries.len() {
let i = entries.len() - pos;
self.stack.last_mut().unwrap().1 = pos + 1;
let (k, v) = &entries[i];
return Some((k, v));
}
self.stack.pop();
}
BNode::Internal { entries, children } => {
let n_steps = 2 * entries.len() + 1;
if pos <= n_steps {
self.stack.last_mut().unwrap().1 = pos + 1;
if pos % 2 == 1 {
let child_idx = entries.len() - (pos - 1) / 2;
self.stack.push((&children[child_idx], 1));
continue;
}
let entry_idx = entries.len() - pos / 2;
let (k, v) = &entries[entry_idx];
return Some((k, v));
}
self.stack.pop();
}
}
}
}
}
impl<'a, K: Ord, V> IntoIterator for &'a PersistentBTreeMap<K, V> {
type Item = (&'a K, &'a V);
type IntoIter = Iter<'a, K, V>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
#[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::*;
use alloc::collections::BTreeMap;
use alloc::vec;
fn check_node<K: Ord + core::fmt::Debug, V>(
node: &BNode<K, V>,
is_root: bool,
depth: usize,
) -> usize {
let entries = match node {
BNode::Leaf { entries } | BNode::Internal { entries, .. } => entries,
};
assert!(
entries.len() <= MAX_ENTRIES,
"node overfull: {} entries",
entries.len()
);
if !is_root {
assert!(
entries.len() >= MIN_ENTRIES,
"non-root node underfull: {} entries",
entries.len()
);
}
for w in entries.windows(2) {
assert!(w[0].0 < w[1].0, "entries out of order inside a node");
}
match node {
BNode::Leaf { .. } => depth,
BNode::Internal { entries, children } => {
assert_eq!(
children.len(),
entries.len() + 1,
"internal node must have entries.len()+1 children"
);
let mut seen: Option<usize> = None;
for c in children {
let d = check_node(c, false, depth + 1);
match seen {
None => seen = Some(d),
Some(prev) => assert_eq!(prev, d, "leaves at different depths"),
}
}
seen.expect("internal node has children")
}
}
}
fn check_map<K: Ord + Clone + core::fmt::Debug, V: Clone>(m: &PersistentBTreeMap<K, V>) {
check_node(m.root.as_ref(), true, 0);
let keys: Vec<&K> = m.iter().map(|(k, _)| k).collect();
for w in keys.windows(2) {
assert!(w[0] < w[1], "iteration order is not globally sorted");
}
assert_eq!(keys.len(), m.len(), "len disagrees with iteration");
}
fn lcg(state: &mut u64) -> u64 {
*state = state
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1);
*state >> 33
}
#[test]
fn round465_remove_absent_key_is_a_no_op() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..50_i64 {
pb.insert_mut(i * 2, i);
}
let before = pb.len();
assert_eq!(pb.remove_mut(&7), None);
assert_eq!(pb.remove_mut(&-1), None);
assert_eq!(pb.remove_mut(&1000), None);
assert_eq!(pb.len(), before);
check_map(&pb);
}
#[test]
fn round465_remove_every_key_empties_the_map() {
for n in [1_i64, 7, 8, 9, 100, 500] {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..n {
pb.insert_mut(i, i * 3);
}
for i in 0..n {
assert_eq!(pb.remove_mut(&i), Some(i * 3), "n={n} i={i}");
assert_eq!(pb.len() as i64, n - i - 1);
check_map(&pb);
for j in (i + 1)..n {
assert_eq!(
pb.get(&j),
Some(&(j * 3)),
"n={n} lost {j} after removing {i}"
);
}
}
assert!(pb.is_empty());
}
}
#[test]
fn round465_remove_in_reverse_order_empties_the_map() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..300_i64 {
pb.insert_mut(i, i);
}
for i in (0..300_i64).rev() {
assert_eq!(pb.remove_mut(&i), Some(i));
check_map(&pb);
}
assert!(pb.is_empty());
}
#[test]
fn round465_matches_btreemap_under_mixed_traffic() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
let mut model: BTreeMap<i64, i64> = BTreeMap::new();
let mut seed = 0x5eed_1234_u64;
for step in 0..4000 {
let k = (lcg(&mut seed) % 300) as i64;
if lcg(&mut seed) % 3 == 0 {
assert_eq!(
pb.remove_mut(&k),
model.remove(&k),
"step {step} remove {k}"
);
} else {
let v = (lcg(&mut seed) % 1000) as i64;
assert_eq!(
pb.insert_mut(k, v),
model.insert(k, v),
"step {step} insert {k}"
);
}
assert_eq!(pb.len(), model.len(), "step {step}");
if step % 97 == 0 {
check_map(&pb);
let got: Vec<(i64, i64)> = pb.iter().map(|(k, v)| (*k, *v)).collect();
let want: Vec<(i64, i64)> = model.iter().map(|(k, v)| (*k, *v)).collect();
assert_eq!(got, want, "step {step}");
}
}
check_map(&pb);
let got: Vec<(i64, i64)> = pb.iter().map(|(k, v)| (*k, *v)).collect();
let want: Vec<(i64, i64)> = model.iter().map(|(k, v)| (*k, *v)).collect();
assert_eq!(got, want);
}
#[test]
fn round465_remove_leaves_a_shared_snapshot_untouched() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..200_i64 {
pb.insert_mut(i, i * 5);
}
let snapshot = pb.clone();
for i in 0..100_i64 {
pb.remove_mut(&(i * 2));
}
check_map(&pb);
check_map(&snapshot);
assert_eq!(snapshot.len(), 200);
for i in 0..200_i64 {
assert_eq!(snapshot.get(&i), Some(&(i * 5)), "snapshot lost {i}");
}
for i in 0..100_i64 {
assert_eq!(pb.get(&(i * 2)), None);
assert_eq!(pb.get(&(i * 2 + 1)), Some(&(i * 10 + 5)));
}
}
#[test]
fn round465_immutable_remove_does_not_touch_the_receiver() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..60_i64 {
pb.insert_mut(i, i);
}
let (next, prev) = pb.remove(&30);
assert_eq!(prev, Some(30));
assert_eq!(pb.get(&30), Some(&30), "receiver must be untouched");
assert_eq!(next.get(&30), None);
assert_eq!(pb.len(), 60);
assert_eq!(next.len(), 59);
check_map(&pb);
check_map(&next);
}
#[test]
fn round465_range_and_predecessor_still_work_after_removals() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..400_i64 {
pb.insert_mut(i, i);
}
for i in 0..400_i64 {
if i % 3 == 0 {
pb.remove_mut(&i);
}
}
check_map(&pb);
let in_range: Vec<i64> = pb
.range(Bound::Included(&100), Bound::Excluded(&120))
.map(|(k, _)| *k)
.collect();
let want: Vec<i64> = (100..120).filter(|i| i % 3 != 0).collect();
assert_eq!(in_range, want);
assert_eq!(pb.predecessor(&100).map(|(k, _)| *k), Some(98));
let rev: Vec<i64> = pb.iter_rev().map(|(k, _)| *k).take(3).collect();
assert_eq!(rev, vec![398, 397, 395]);
}
#[test]
fn empty_map_is_empty() {
let pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
assert_eq!(pb.len(), 0);
assert!(pb.is_empty());
assert!(pb.get(&42).is_none());
}
#[test]
fn insert_single_into_empty_works() {
let (pb, prev) = PersistentBTreeMap::<i64, i64>::new().insert(1, 100);
assert_eq!(prev, None);
assert_eq!(pb.len(), 1);
assert_eq!(pb.get(&1), Some(&100));
assert_eq!(pb.get(&2), None);
}
#[test]
fn insert_replace_returns_prev_keeps_len() {
let (pb, p1) = PersistentBTreeMap::<i64, i64>::new().insert(7, 10);
assert_eq!(p1, None);
let (pb, p2) = pb.insert(7, 99);
assert_eq!(p2, Some(10));
assert_eq!(pb.len(), 1);
assert_eq!(pb.get(&7), Some(&99));
}
#[test]
fn insert_crosses_leaf_split_boundary() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..20_i64 {
pb = pb.insert(i, i * 7).0;
}
for i in 0..20_i64 {
assert_eq!(pb.get(&i), Some(&(i * 7)));
}
assert!(pb.get(&20).is_none());
assert_eq!(pb.len(), 20);
}
#[test]
fn insert_grows_through_multiple_internal_splits() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..200_i64 {
pb = pb.insert(i, i * 11).0;
}
for i in 0..200_i64 {
assert_eq!(pb.get(&i), Some(&(i * 11)));
}
assert_eq!(pb.len(), 200);
}
#[test]
fn clone_then_insert_preserves_original() {
let mut a: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..100_i64 {
a = a.insert(i, i).0;
}
let b = a.clone();
let (b, _) = b.insert(999, 999);
assert_eq!(a.len(), 100);
assert!(a.get(&999).is_none());
assert_eq!(b.len(), 101);
assert_eq!(b.get(&999), Some(&999));
for i in 0..100_i64 {
assert_eq!(a.get(&i), Some(&i), "A drift at {i}");
assert_eq!(b.get(&i), Some(&i), "B drift at {i}");
}
}
#[test]
fn iter_yields_sorted_order() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for &k in &[7_i64, 3, 11, 1, 9, 5, 14, 2, 8, 12, 4, 6, 10, 13] {
pb = pb.insert(k, k * 2).0;
}
let collected: Vec<(i64, i64)> = pb.iter().map(|(k, v)| (*k, *v)).collect();
let expected: Vec<(i64, i64)> = (1..=14).map(|k| (k, k * 2)).collect();
assert_eq!(collected, expected);
}
#[test]
fn iter_handles_taller_tree() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..500_i64 {
pb = pb.insert(i, i).0;
}
let collected: Vec<i64> = pb.iter().map(|(k, _)| *k).collect();
let expected: Vec<i64> = (0..500).collect();
assert_eq!(collected, expected);
}
#[test]
fn range_basic_bounds() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..100_i64 {
pb = pb.insert(i, i * 10).0;
}
let keys = |lo: Bound<&i64>, hi: Bound<&i64>| -> Vec<i64> {
pb.range(lo, hi).map(|(k, _)| *k).collect()
};
assert_eq!(
keys(Bound::Included(&20), Bound::Included(&24)),
vec![20, 21, 22, 23, 24]
);
assert_eq!(
keys(Bound::Excluded(&20), Bound::Excluded(&24)),
vec![21, 22, 23]
);
assert_eq!(keys(Bound::Unbounded, Bound::Excluded(&3)), vec![0, 1, 2]);
assert_eq!(
keys(Bound::Included(&97), Bound::Unbounded),
vec![97, 98, 99]
);
assert!(keys(Bound::Included(&50), Bound::Included(&49)).is_empty());
assert_eq!(
keys(Bound::Included(&-5), Bound::Included(&2)),
vec![0, 1, 2]
);
assert_eq!(
keys(Bound::Included(&200), Bound::Unbounded),
Vec::<i64>::new()
);
}
#[test]
fn fuzz_range_against_btreemap() {
let mut rng = Splitmix::new(0x5EED_1234_u64);
const KEY_RANGE: i64 = 512;
for &n_inserts in &[0usize, 5, 40, 300, 2000] {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
let mut oracle: BTreeMap<i64, i64> = BTreeMap::new();
for _ in 0..n_inserts {
let key = (rng.next() as i64).rem_euclid(KEY_RANGE);
let val = rng.next() as i64;
pb = pb.insert(key, val).0;
oracle.insert(key, val);
}
for _ in 0..2000 {
let a = (rng.next() as i64).rem_euclid(KEY_RANGE + 40) - 20;
let b = (rng.next() as i64).rem_euclid(KEY_RANGE + 40) - 20;
let (lo_raw, hi_raw) = if a <= b { (a, b) } else { (b, a) };
let mk = |raw: i64, sel: u64| -> Bound<i64> {
match sel % 3 {
0 => Bound::Included(raw),
1 => Bound::Excluded(raw),
_ => Bound::Unbounded,
}
};
let lo = mk(lo_raw, rng.next());
let hi = mk(hi_raw, rng.next());
if lo_raw == hi_raw
&& matches!(lo, Bound::Excluded(_))
&& matches!(hi, Bound::Excluded(_))
{
continue;
}
let lo_ref = match &lo {
Bound::Included(k) => Bound::Included(k),
Bound::Excluded(k) => Bound::Excluded(k),
Bound::Unbounded => Bound::Unbounded,
};
let hi_ref = match &hi {
Bound::Included(k) => Bound::Included(k),
Bound::Excluded(k) => Bound::Excluded(k),
Bound::Unbounded => Bound::Unbounded,
};
let got: Vec<(i64, i64)> =
pb.range(lo_ref, hi_ref).map(|(k, v)| (*k, *v)).collect();
let want: Vec<(i64, i64)> = oracle.range((lo, hi)).map(|(k, v)| (*k, *v)).collect();
assert_eq!(
got, want,
"range drift n={n_inserts} lo={lo_raw:?} hi={hi_raw:?}"
);
}
}
}
#[test]
fn fuzz_predecessor_against_btreemap() {
let mut rng = Splitmix::new(0x9E37_79B9_7F4A_7C15_u64);
const KEY_RANGE: i64 = 512;
for &n_inserts in &[0usize, 5, 40, 300, 2000] {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
let mut oracle: BTreeMap<i64, i64> = BTreeMap::new();
for _ in 0..n_inserts {
let key = (rng.next() as i64).rem_euclid(KEY_RANGE);
let val = rng.next() as i64;
pb = pb.insert(key, val).0;
oracle.insert(key, val);
}
for _ in 0..2000 {
let probe = (rng.next() as i64).rem_euclid(KEY_RANGE + 40) - 20;
let got = pb.predecessor(&probe).map(|(k, v)| (*k, *v));
let want = oracle.range(..probe).next_back().map(|(k, v)| (*k, *v));
assert_eq!(got, want, "predecessor drift n={n_inserts} probe={probe}");
}
}
}
#[test]
fn iter_rev_yields_descending() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for &k in &[7_i64, 3, 11, 1, 9, 5, 14, 2, 8, 12, 4, 6, 10, 13] {
pb = pb.insert(k, k * 2).0;
}
let collected: Vec<(i64, i64)> = pb.iter_rev().map(|(k, v)| (*k, *v)).collect();
let expected: Vec<(i64, i64)> = (1..=14).rev().map(|k| (k, k * 2)).collect();
assert_eq!(collected, expected);
}
#[test]
fn iter_rev_handles_taller_tree() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..500_i64 {
pb = pb.insert(i, i).0;
}
let collected: Vec<i64> = pb.iter_rev().map(|(k, _)| *k).collect();
let expected: Vec<i64> = (0..500).rev().collect();
assert_eq!(collected, expected);
}
#[test]
fn iter_rev_empty_map_returns_nothing() {
let pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
assert_eq!(pb.iter_rev().count(), 0);
}
#[test]
fn iter_rev_lazy_stops_at_take() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for i in 0..10000_i64 {
pb = pb.insert(i, i).0;
}
let top5: Vec<i64> = pb.iter_rev().take(5).map(|(k, _)| *k).collect();
assert_eq!(top5, vec![9999, 9998, 9997, 9996, 9995]);
}
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_std_btreemap() {
let mut pb: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
let mut oracle: BTreeMap<i64, i64> = BTreeMap::new();
let mut rng = Splitmix::new(0xC0FFEE_u64);
const STEPS: usize = 100_000;
const KEY_RANGE: i64 = 4096;
for step in 0..STEPS {
let op = rng.next() % 3; let key = (rng.next() as i64) % KEY_RANGE;
match op {
0 | 1 => {
let val = rng.next() as i64;
let (new_pb, prev_pb) = pb.insert(key, val);
let prev_oracle = oracle.insert(key, val);
assert_eq!(prev_pb, prev_oracle, "prev drift @ step {step}, key {key}");
pb = new_pb;
assert_eq!(pb.len(), oracle.len(), "len drift @ step {step}");
}
2 => {
let pb_v = pb.get(&key).copied();
let oracle_v = oracle.get(&key).copied();
assert_eq!(pb_v, oracle_v, "get drift @ step {step}, key {key}");
}
_ => unreachable!(),
}
}
for (k, v) in &oracle {
assert_eq!(pb.get(k), Some(v), "final drift at key {k}");
}
let pb_collected: Vec<(i64, i64)> = pb.iter().map(|(k, v)| (*k, *v)).collect();
let oracle_collected: Vec<(i64, i64)> = oracle.iter().map(|(k, v)| (*k, *v)).collect();
assert_eq!(pb_collected, oracle_collected);
}
#[test]
fn fuzz_oracle_clone_isolation() {
let mut a: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
let mut oracle_a: BTreeMap<i64, i64> = BTreeMap::new();
let mut rng = Splitmix::new(0xDECAFBAD_u64);
for _ in 0..1_000 {
let k = (rng.next() as i64) % 1000;
let v = rng.next() as i64;
a = a.insert(k, v).0;
oracle_a.insert(k, 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 k = (rng.next() as i64) % 2000;
let v = rng.next() as i64;
b = b.insert(k, v).0;
oracle_b.insert(k, v);
}
for _ in 0..300 {
let k = (rng.next() as i64) % 500;
let v = rng.next() as i64;
c = c.insert(k, v).0;
oracle_c.insert(k, v);
}
for (k, v) in &oracle_a {
assert_eq!(a.get(k), Some(v), "A drift at {k}");
}
for (k, v) in &oracle_b {
assert_eq!(b.get(k), Some(v), "B drift at {k}");
}
for (k, v) in &oracle_c {
assert_eq!(c.get(k), Some(v), "C drift at {k}");
}
}
#[test]
fn partial_eq_compares_by_elements() {
let mut a: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
let mut b: PersistentBTreeMap<i64, i64> = PersistentBTreeMap::new();
for &k in &[5_i64, 2, 8, 1, 7, 3, 6, 4] {
a = a.insert(k, k * 10).0;
}
for &k in &[1_i64, 2, 3, 4, 5, 6, 7, 8] {
b = b.insert(k, k * 10).0;
}
assert_eq!(a, b);
let (a, _) = a.insert(9, 90);
assert_ne!(a, b);
}
}