use std::cmp::min;
use std::collections::Bound;
use std::ops::RangeBounds;
use crate::iter::LendingKeyView;
use crate::keys::KeyTrait;
use crate::mapping::{
NodeMapping,
direct_mapping::{DirectMapping, DirectMappingIter},
indexed_mapping::{IndexedMapping, IndexedMappingIter},
sorted_keyed_mapping::{SortedKeyedMapping, SortedKeyedMappingIter},
};
use crate::partials::Partial;
use crate::utils::bitset::Bitset64;
#[cfg(not(feature = "triomphe-arc"))]
use std::sync::Arc;
#[cfg(feature = "triomphe-arc")]
use triomphe::Arc;
type RemoveResult<P, V> = (Option<Arc<VersionedNode<P, V>>>, V);
type VersionedPrefixSubtreeView<'a, P, V> = (&'a VersionedNode<P, V>, Vec<&'a [u8]>, usize);
type VersionedIterEntry<'a, P, V> = (u8, &'a VersionedNode<P, V>);
enum VersionedIterFrameIter<'a, P: Partial, V> {
Plain(VersionedNodeIter<'a, P, V>),
Leading {
first: Option<VersionedIterEntry<'a, P, V>>,
rest: VersionedNodeIter<'a, P, V>,
},
}
impl<'a, P: Partial, V> Iterator for VersionedIterFrameIter<'a, P, V> {
type Item = VersionedIterEntry<'a, P, V>;
fn next(&mut self) -> Option<Self::Item> {
match self {
VersionedIterFrameIter::Plain(iter) => iter.next(),
VersionedIterFrameIter::Leading { first, rest } => first.take().or_else(|| rest.next()),
}
}
}
pub(crate) enum VersionedNodeIter<'a, P: Partial, V> {
Node4(SortedKeyedMappingIter<'a, Arc<VersionedNode<P, V>>, 4>),
Node16(SortedKeyedMappingIter<'a, Arc<VersionedNode<P, V>>, 16>),
Node48(IndexedMappingIter<'a, Arc<VersionedNode<P, V>>, 48, Bitset64<1>>),
Node256(DirectMappingIter<'a, Arc<VersionedNode<P, V>>>),
Empty,
}
impl<'a, P: Partial, V> Iterator for VersionedNodeIter<'a, P, V> {
type Item = (u8, &'a VersionedNode<P, V>);
fn next(&mut self) -> Option<Self::Item> {
match self {
VersionedNodeIter::Node4(iter) => iter.next().map(|(key, child)| (key, child.as_ref())),
VersionedNodeIter::Node16(iter) => {
iter.next().map(|(key, child)| (key, child.as_ref()))
}
VersionedNodeIter::Node48(iter) => {
iter.next().map(|(key, child)| (key, child.as_ref()))
}
VersionedNodeIter::Node256(iter) => {
iter.next().map(|(key, child)| (key, child.as_ref()))
}
VersionedNodeIter::Empty => None,
}
}
}
pub struct VersionedIter<'a, K: KeyTrait<PartialType = P>, P: Partial + 'a, V> {
inner: Box<dyn Iterator<Item = (K, &'a V)> + 'a>,
_marker: std::marker::PhantomData<(K, P)>,
}
pub struct VersionedPrefixMatchIter<'a, K: KeyTrait<PartialType = P>, P: Partial + 'a, V> {
cur_node: Option<&'a VersionedNode<P, V>>,
probe: K,
cur_key: Vec<u8>,
depth: usize,
}
struct VersionedIterInner<'a, K: KeyTrait<PartialType = P>, P: Partial + 'a, V> {
node_iter_stack: Vec<(usize, VersionedIterFrameIter<'a, P, V>)>,
cur_key: Vec<u8>,
start_bound: Option<Bound<K>>,
}
pub(crate) struct VersionedLendingIterInner<'a, P: Partial + 'a, V> {
node_iter_stack: Vec<(usize, usize, VersionedIterFrameIter<'a, P, V>)>,
cur_segments: Vec<&'a [u8]>,
cur_len: usize,
end_bound: Option<(Vec<u8>, bool)>,
}
pub struct VersionedValuesIter<'a, P: Partial + 'a, V> {
root_value: Option<&'a V>,
node_iter_stack: Vec<VersionedNodeIter<'a, P, V>>,
}
pub struct VersionedRange<'a, K: KeyTrait + 'a, V> {
iter: VersionedIter<'a, K, K::PartialType, V>,
end: Bound<K>,
}
pub struct VersionedAdaptiveRadixTree<KeyType, ValueType>
where
KeyType: KeyTrait,
ValueType: Clone,
{
root: Option<Arc<VersionedNode<KeyType::PartialType, ValueType>>>,
version: u64,
_phantom: std::marker::PhantomData<KeyType>,
}
pub struct VersionedNode<P: Partial, V> {
pub(crate) prefix: P,
pub(crate) value: Option<V>,
pub(crate) content: VersionedContent<P, V>,
pub(crate) version: u64,
}
pub(crate) enum VersionedContent<P: Partial, V> {
Empty,
Node4(Box<SortedKeyedMapping<Arc<VersionedNode<P, V>>, 4>>),
Node16(Box<SortedKeyedMapping<Arc<VersionedNode<P, V>>, 16>>),
Node48(Box<IndexedMapping<Arc<VersionedNode<P, V>>, 48, Bitset64<1>>>),
Node256(Box<DirectMapping<Arc<VersionedNode<P, V>>>>),
}
impl<KeyType: KeyTrait, ValueType: Clone> Default
for VersionedAdaptiveRadixTree<KeyType, ValueType>
{
fn default() -> Self {
Self::new()
}
}
impl<KeyType, ValueType> Clone for VersionedAdaptiveRadixTree<KeyType, ValueType>
where
KeyType: KeyTrait,
ValueType: Clone,
{
fn clone(&self) -> Self {
Self {
root: self.root.clone(),
version: self.version,
_phantom: std::marker::PhantomData,
}
}
}
impl<KeyType, ValueType> VersionedAdaptiveRadixTree<KeyType, ValueType>
where
KeyType: KeyTrait,
ValueType: Clone,
{
pub fn new() -> Self {
Self {
root: None,
version: 0,
_phantom: std::marker::PhantomData,
}
}
pub fn snapshot(&self) -> Self {
Self {
root: self.root.clone(),
version: self.version + 1,
_phantom: std::marker::PhantomData,
}
}
#[inline]
pub fn get<Key>(&self, key: Key) -> Option<&ValueType>
where
Key: Into<KeyType>,
{
self.get_k(&key.into())
}
#[inline]
pub fn get_k(&self, key: &KeyType) -> Option<&ValueType> {
let root = self.root.as_ref()?;
Self::get_iterate(root, key)
}
pub fn iter(&self) -> VersionedIter<'_, KeyType, KeyType::PartialType, ValueType> {
VersionedIter::new(self.root.as_deref())
}
pub fn for_each_view<F>(&self, on_each: F)
where
F: for<'view> FnMut(LendingKeyView<'_, 'view>, &ValueType),
{
VersionedLendingIterInner::for_each(self.root.as_deref(), on_each);
}
pub fn values_iter(&self) -> VersionedValuesIter<'_, KeyType::PartialType, ValueType> {
VersionedValuesIter::new(self.root.as_deref())
}
pub fn intersect_with<'a, F>(&'a self, other: &'a Self, mut on_match: F)
where
F: FnMut(KeyType, &'a ValueType, &'a ValueType),
{
let (Some(left_root), Some(right_root)) = (self.root.as_deref(), other.root.as_deref())
else {
return;
};
let mut key_buf = Vec::with_capacity(KeyType::MAXIMUM_SIZE.unwrap_or(64));
Self::intersect_nodes(left_root, 0, right_root, 0, &mut key_buf, &mut on_match);
}
pub fn intersect_lending_with<'a, F>(&'a self, other: &'a Self, mut on_match: F)
where
F: for<'view> FnMut(LendingKeyView<'a, 'view>, &'a ValueType, &'a ValueType),
{
let (Some(left_root), Some(right_root)) = (self.root.as_deref(), other.root.as_deref())
else {
return;
};
let mut segments = Vec::new();
let mut key_len = 0usize;
Self::intersect_nodes_lending(
left_root,
0,
right_root,
0,
&mut segments,
&mut key_len,
&mut on_match,
);
}
pub fn intersect_values_with<'a, F>(&'a self, other: &'a Self, mut on_match: F)
where
F: FnMut(&'a ValueType, &'a ValueType),
{
let (Some(left_root), Some(right_root)) = (self.root.as_deref(), other.root.as_deref())
else {
return;
};
Self::intersect_nodes_values(left_root, 0, right_root, 0, &mut on_match);
}
pub fn intersect_count(&self, other: &Self) -> usize {
let mut count = 0usize;
self.intersect_values_with(other, |_left_value, _right_value| {
count += 1;
});
count
}
#[inline]
pub fn prefix_match_iter<Key>(
&self,
key: Key,
) -> VersionedPrefixMatchIter<'_, KeyType, KeyType::PartialType, ValueType>
where
Key: Into<KeyType>,
{
VersionedPrefixMatchIter::new(self.root.as_deref(), key.into())
}
#[inline]
pub fn prefix_match_iter_k(
&self,
key: &KeyType,
) -> VersionedPrefixMatchIter<'_, KeyType, KeyType::PartialType, ValueType> {
VersionedPrefixMatchIter::new(self.root.as_deref(), key.clone())
}
#[inline]
pub fn prefix_match_for_each<Key, F>(&self, key: Key, on_match: F)
where
Key: Into<KeyType>,
F: FnMut(&[u8], &ValueType),
{
self.prefix_match_for_each_k(&key.into(), on_match)
}
#[inline]
pub fn prefix_match_for_each_k<F>(&self, key: &KeyType, on_match: F)
where
F: FnMut(&[u8], &ValueType),
{
let Some(root) = self.root.as_deref() else {
return;
};
Self::prefix_match_for_each_impl(root, key, on_match);
}
#[inline]
pub fn prefix_iter<Key>(
&self,
prefix: Key,
) -> VersionedIter<'_, KeyType, KeyType::PartialType, ValueType>
where
Key: Into<KeyType>,
{
self.prefix_iter_k(&prefix.into())
}
pub fn prefix_iter_k(
&self,
prefix: &KeyType,
) -> VersionedIter<'_, KeyType, KeyType::PartialType, ValueType> {
let Some(root) = self.root.as_deref() else {
return VersionedIter::empty();
};
let Some((subtree_root, subtree_root_key)) = Self::find_prefix_subtree(root, prefix) else {
return VersionedIter::empty();
};
VersionedIter::new_with_prefix(Some(subtree_root), subtree_root_key)
}
pub fn prefix_for_each_view<Key, F>(&self, prefix: Key, on_each: F)
where
Key: Into<KeyType>,
F: for<'view> FnMut(LendingKeyView<'_, 'view>, &ValueType),
{
self.prefix_for_each_view_k(&prefix.into(), on_each)
}
pub fn prefix_for_each_view_k<F>(&self, prefix: &KeyType, on_each: F)
where
F: for<'view> FnMut(LendingKeyView<'_, 'view>, &ValueType),
{
let Some(root) = self.root.as_deref() else {
return;
};
let Some((subtree_root, subtree_root_segments, subtree_root_len)) =
Self::find_prefix_subtree_view(root, prefix)
else {
return;
};
VersionedLendingIterInner::for_each_with_prefix(
Some(subtree_root),
subtree_root_segments,
subtree_root_len,
on_each,
);
}
pub fn range<'a, R>(&'a self, range: R) -> VersionedRange<'a, KeyType, ValueType>
where
R: RangeBounds<KeyType> + 'a,
{
let start_bound = range.start_bound().cloned();
let end_bound = range.end_bound().cloned();
let iter = match start_bound {
Bound::Unbounded => self.iter(),
_ => VersionedIter::new_with_start_bound(self.root.as_deref(), start_bound),
};
VersionedRange {
iter,
end: end_bound,
}
}
pub fn for_each_range_view<R, F>(&self, range: R, on_each: F)
where
R: RangeBounds<KeyType>,
F: for<'view> FnMut(LendingKeyView<'_, 'view>, &ValueType),
{
let start_bound = range.start_bound().cloned();
let end_bound = range.end_bound().cloned();
VersionedLendingIterInner::for_each_with_bounds(
self.root.as_deref(),
start_bound,
end_bound,
on_each,
);
}
#[inline]
pub fn insert<KV>(&mut self, key: KV, value: ValueType) -> bool
where
KV: Into<KeyType>,
{
self.insert_k(&key.into(), value)
}
pub fn insert_k(&mut self, key: &KeyType, value: ValueType) -> bool {
self.version += 1;
let Some(root) = self.root.take() else {
self.root = Some(Arc::new(VersionedNode::new_leaf(
key.to_partial(0),
value,
self.version,
)));
return false;
};
let (new_root, was_replaced) =
Self::insert_recurse(root, key, value, 0, self.version, None);
self.root = Some(new_root);
was_replaced
}
#[inline]
pub fn insert_and_replace<KV>(&mut self, key: KV, value: ValueType) -> Option<ValueType>
where
KV: Into<KeyType>,
{
self.insert_and_replace_k(&key.into(), value)
}
pub fn insert_and_replace_k(&mut self, key: &KeyType, value: ValueType) -> Option<ValueType> {
self.version += 1;
let Some(root) = self.root.take() else {
self.root = Some(Arc::new(VersionedNode::new_leaf(
key.to_partial(0),
value,
self.version,
)));
return None;
};
let mut old_value = None;
let (new_root, _was_replaced) =
Self::insert_recurse(root, key, value, 0, self.version, Some(&mut old_value));
self.root = Some(new_root);
old_value
}
pub fn remove<KV>(&mut self, key: KV) -> Option<ValueType>
where
KV: Into<KeyType>,
{
self.remove_k(&key.into())
}
pub fn remove_k(&mut self, key: &KeyType) -> Option<ValueType> {
self.get_k(key)?;
self.version += 1;
let root = self
.root
.take()
.expect("non-empty tree checked before remove mutation");
if root.is_leaf() {
if root.prefix.len() != key.length_at(0) {
self.root = Some(root);
return None;
}
match Arc::try_unwrap(root) {
Ok(mut owned_root) => {
return owned_root.value.take();
}
Err(shared_root) => {
let cloned_value = shared_root.value().cloned();
self.root = None;
return cloned_value;
}
}
}
if root.prefix.len() == key.length_at(0) {
let new_root = Self::ensure_cow_node(root, self.version);
let mut new_root = match Arc::try_unwrap(new_root) {
Ok(owned) => owned,
Err(_) => panic!("ensure_cow_node should have given us exclusive ownership"),
};
let removed = new_root.value.take();
if new_root.num_children() == 0 && new_root.value.is_none() {
self.root = None;
} else {
self.root = Some(Arc::new(new_root));
}
return removed;
}
let (new_root, removed_value) = Self::remove_recurse(root, key, 0, self.version)
.expect("prechecked key should be removable");
if let Some(root_node) = new_root {
if root_node.is_inner() && root_node.num_children() == 0 && root_node.value().is_none()
{
self.root = None;
} else {
self.root = Some(root_node);
}
} else {
self.root = None;
}
Some(removed_value)
}
pub fn is_empty(&self) -> bool {
self.root.is_none()
}
pub fn version(&self) -> u64 {
self.version
}
pub fn into_unversioned(self) -> crate::tree::AdaptiveRadixTree<KeyType, ValueType> {
use crate::tree::AdaptiveRadixTree;
let Some(root) = self.root else {
return AdaptiveRadixTree::new();
};
let converted_root = Self::convert_to_unversioned_node(root);
AdaptiveRadixTree::from_root(converted_root)
}
fn convert_to_unversioned_node(
node: Arc<VersionedNode<KeyType::PartialType, ValueType>>,
) -> crate::node::DefaultNode<KeyType::PartialType, ValueType> {
use crate::mapping::{
direct_mapping::DirectMapping, indexed_mapping::IndexedMapping,
sorted_keyed_mapping::SortedKeyedMapping,
};
use crate::node::{Content, DefaultNode};
match Arc::try_unwrap(node) {
Ok(owned_node) => {
let VersionedNode {
prefix,
value,
content,
version: _,
} = owned_node;
let unversioned_content = match content {
VersionedContent::Empty => Content::Empty,
VersionedContent::Node4(km) => {
let mut new_km = SortedKeyedMapping::new();
for (key, child) in km.into_iter() {
let converted_child = Self::convert_to_unversioned_node(child);
new_km.add_child(key, converted_child);
}
Content::Node4(Box::new(new_km))
}
VersionedContent::Node16(km) => {
let mut new_km = SortedKeyedMapping::new();
for (key, child) in km.into_iter() {
let converted_child = Self::convert_to_unversioned_node(child);
new_km.add_child(key, converted_child);
}
Content::Node16(Box::new(new_km))
}
VersionedContent::Node48(km) => {
let mut new_km = IndexedMapping::new();
for (key, child) in km.into_iter() {
let converted_child = Self::convert_to_unversioned_node(child);
new_km.add_child(key, converted_child);
}
Content::Node48(Box::new(new_km))
}
VersionedContent::Node256(km) => {
let mut new_km = DirectMapping::new();
for (key, child) in km.into_iter() {
let converted_child = Self::convert_to_unversioned_node(child);
new_km.add_child(key, converted_child);
}
Content::Node256(Box::new(new_km))
}
};
DefaultNode {
prefix,
value,
content: unversioned_content,
}
}
Err(shared_node) => {
let unversioned_content = match &shared_node.content {
VersionedContent::Empty => Content::Empty,
VersionedContent::Node4(km) => {
let mut new_km = SortedKeyedMapping::new();
for (key, child) in km.iter() {
let converted_child =
Self::convert_to_unversioned_node(Arc::clone(child));
new_km.add_child(key, converted_child);
}
Content::Node4(Box::new(new_km))
}
VersionedContent::Node16(km) => {
let mut new_km = SortedKeyedMapping::new();
for (key, child) in km.iter() {
let converted_child =
Self::convert_to_unversioned_node(Arc::clone(child));
new_km.add_child(key, converted_child);
}
Content::Node16(Box::new(new_km))
}
VersionedContent::Node48(km) => {
let mut new_km = IndexedMapping::new();
for (key, child) in km.iter() {
let converted_child =
Self::convert_to_unversioned_node(Arc::clone(child));
new_km.add_child(key, converted_child);
}
Content::Node48(Box::new(new_km))
}
VersionedContent::Node256(km) => {
let mut new_km = DirectMapping::new();
for (key, child) in km.iter() {
let converted_child =
Self::convert_to_unversioned_node(Arc::clone(child));
new_km.add_child(key, converted_child);
}
Content::Node256(Box::new(new_km))
}
};
DefaultNode {
prefix: shared_node.prefix.clone(),
value: shared_node.value.clone(),
content: unversioned_content,
}
}
}
}
}
impl<P: Partial, V> VersionedNode<P, V> {
pub fn new_leaf(prefix: P, value: V, version: u64) -> Self {
Self {
prefix,
value: Some(value),
content: VersionedContent::Empty,
version,
}
}
pub fn new_inner(prefix: P, version: u64) -> Self {
Self {
prefix,
value: None,
content: VersionedContent::Node4(Box::default()),
version,
}
}
pub fn is_leaf(&self) -> bool {
matches!(&self.content, VersionedContent::Empty)
}
pub fn is_inner(&self) -> bool {
!self.is_leaf()
}
pub fn value(&self) -> Option<&V> {
self.value.as_ref()
}
pub fn seek_child(&self, key: u8) -> Option<&Arc<VersionedNode<P, V>>> {
match &self.content {
VersionedContent::Node4(km) => km.seek_child(key),
VersionedContent::Node16(km) => km.seek_child(key),
VersionedContent::Node48(km) => km.seek_child(key),
VersionedContent::Node256(km) => km.seek_child(key),
VersionedContent::Empty => None,
}
}
pub fn num_children(&self) -> usize {
match &self.content {
VersionedContent::Node4(km) => km.num_children(),
VersionedContent::Node16(km) => km.num_children(),
VersionedContent::Node48(km) => km.num_children(),
VersionedContent::Node256(km) => km.num_children(),
VersionedContent::Empty => 0,
}
}
pub fn is_full(&self) -> bool {
match &self.content {
VersionedContent::Node4(km) => km.num_children() >= 4,
VersionedContent::Node16(km) => km.num_children() >= 16,
VersionedContent::Node48(km) => km.num_children() >= 48,
VersionedContent::Node256(_) => false, VersionedContent::Empty => false,
}
}
pub fn grow(&self, new_version: u64) -> Self
where
P: Clone,
V: Clone,
{
Self {
prefix: self.prefix.clone(),
value: self.value.clone(),
content: match &self.content {
VersionedContent::Node4(km) => {
let mut new_km = SortedKeyedMapping::new();
for (key, child) in km.iter() {
new_km.add_child(key, Arc::clone(child));
}
VersionedContent::Node16(Box::new(new_km))
}
VersionedContent::Node16(km) => {
let mut new_km = IndexedMapping::new();
for (key, child) in km.iter() {
new_km.add_child(key, Arc::clone(child));
}
VersionedContent::Node48(Box::new(new_km))
}
VersionedContent::Node48(km) => {
let mut new_km = DirectMapping::new();
for (key, child) in km.iter() {
new_km.add_child(key, Arc::clone(child));
}
VersionedContent::Node256(Box::new(new_km))
}
VersionedContent::Node256(_) => {
panic!("Node256 cannot grow further")
}
VersionedContent::Empty => {
panic!("Leaf nodes cannot grow")
}
},
version: new_version,
}
}
pub fn cow_clone_inner(&self, new_version: u64) -> Self
where
P: Clone,
V: Clone,
{
Self {
prefix: self.prefix.clone(),
value: self.value.clone(),
content: match &self.content {
VersionedContent::Empty => VersionedContent::Empty,
VersionedContent::Node4(km) => {
let mut new_km = SortedKeyedMapping::new();
for (key, child) in km.iter() {
new_km.add_child(key, Arc::clone(child));
}
VersionedContent::Node4(Box::new(new_km))
}
VersionedContent::Node16(km) => {
let mut new_km = SortedKeyedMapping::new();
for (key, child) in km.iter() {
new_km.add_child(key, Arc::clone(child));
}
VersionedContent::Node16(Box::new(new_km))
}
VersionedContent::Node48(km) => {
VersionedContent::Node48(Box::new(km.clone_mapping()))
}
VersionedContent::Node256(km) => {
VersionedContent::Node256(Box::new(km.clone_mapping()))
}
},
version: new_version,
}
}
fn add_child(&mut self, key: u8, child: Arc<VersionedNode<P, V>>)
where
P: Clone,
V: Clone,
{
if matches!(self.content, VersionedContent::Empty) {
self.content = VersionedContent::Node4(Box::default());
}
if self.is_full() {
*self = self.grow(self.version);
}
match &mut self.content {
VersionedContent::Node4(km) => km.add_child(key, child),
VersionedContent::Node16(km) => km.add_child(key, child),
VersionedContent::Node48(km) => km.add_child(key, child),
VersionedContent::Node256(km) => km.add_child(key, child),
VersionedContent::Empty => {
unreachable!("empty nodes are promoted before adding children")
}
}
}
fn delete_child(&mut self, key: u8) -> Option<Arc<VersionedNode<P, V>>> {
match &mut self.content {
VersionedContent::Node4(km) => km.delete_child(key),
VersionedContent::Node16(km) => km.delete_child(key),
VersionedContent::Node48(km) => km.delete_child(key),
VersionedContent::Node256(km) => km.delete_child(key),
VersionedContent::Empty => None,
}
}
pub(crate) fn iter(&self) -> VersionedNodeIter<'_, P, V> {
match &self.content {
VersionedContent::Node4(n) => VersionedNodeIter::Node4(n.iter()),
VersionedContent::Node16(n) => VersionedNodeIter::Node16(n.iter()),
VersionedContent::Node48(n) => VersionedNodeIter::Node48(n.iter()),
VersionedContent::Node256(n) => VersionedNodeIter::Node256(n.iter()),
VersionedContent::Empty => VersionedNodeIter::Empty,
}
}
}
impl<'a, K: KeyTrait<PartialType = P>, P: Partial + 'a, V> VersionedIterInner<'a, K, P, V> {
#[inline]
fn key_order(lhs: &K, rhs: &K) -> std::cmp::Ordering {
let lhs_len = lhs.length_at(0);
let rhs_len = rhs.length_at(0);
let common = lhs_len.min(rhs_len);
for i in 0..common {
match lhs.at(i).cmp(&rhs.at(i)) {
std::cmp::Ordering::Equal => {}
ord => return ord,
}
}
lhs_len.cmp(&rhs_len)
}
fn from_node_and_key(node: &'a VersionedNode<P, V>, cur_key: K) -> Self {
Self {
node_iter_stack: vec![(
cur_key.length_at(0),
VersionedIterFrameIter::Plain(node.iter()),
)],
cur_key: cur_key.as_ref().to_vec(),
start_bound: None,
}
}
fn new(node: &'a VersionedNode<P, V>) -> Self {
Self::from_node_and_key(node, K::new_from_partial(&node.prefix))
}
fn new_with_start_bound(node: &'a VersionedNode<P, V>, start_bound: Bound<K>) -> Self {
let seek_key = match &start_bound {
Bound::Included(key) | Bound::Excluded(key) => Some(key),
Bound::Unbounded => None,
};
if let Some(seek_key) = seek_key {
return Self {
node_iter_stack: Self::build_positioned_stack(node, seek_key, 0),
cur_key: node.prefix.as_ref().to_vec(),
start_bound: Some(start_bound),
};
}
Self {
node_iter_stack: vec![(
node.prefix.len(),
VersionedIterFrameIter::Plain(node.iter()),
)],
cur_key: node.prefix.as_ref().to_vec(),
start_bound: None,
}
}
fn build_positioned_stack(
node: &'a VersionedNode<P, V>,
seek_key: &K,
depth: usize,
) -> Vec<(usize, VersionedIterFrameIter<'a, P, V>)> {
let prefix_common = node.prefix.prefix_length_key(seek_key, depth);
if prefix_common != node.prefix.len() {
let seek_remaining = seek_key.length_at(depth);
if prefix_common >= seek_remaining {
return vec![(
node.prefix.len(),
VersionedIterFrameIter::Plain(node.iter()),
)];
}
let node_byte = node.prefix.at(prefix_common);
let seek_byte = seek_key.at(depth + prefix_common);
if node_byte < seek_byte {
return vec![];
}
return vec![(
node.prefix.len(),
VersionedIterFrameIter::Plain(node.iter()),
)];
}
if seek_key.length_at(depth) == node.prefix.len() {
return vec![(
node.prefix.len(),
VersionedIterFrameIter::Plain(node.iter()),
)];
}
let target_depth = depth + node.prefix.len();
let target_byte = seek_key.at(target_depth);
let mut iter = node.iter();
while let Some((key, child)) = iter.next() {
if key < target_byte {
continue;
}
return vec![(
node.prefix.len(),
VersionedIterFrameIter::Leading {
first: Some((key, child)),
rest: iter,
},
)];
}
vec![]
}
}
impl<'a, K: KeyTrait<PartialType = P> + 'a, P: Partial + 'a, V> VersionedIter<'a, K, P, V> {
fn empty() -> Self {
Self {
inner: Box::new(std::iter::empty()),
_marker: Default::default(),
}
}
fn from_root_and_children(
root_key: K,
root_value: Option<&'a V>,
children: VersionedIterInner<'a, K, P, V>,
) -> Self {
let inner: Box<dyn Iterator<Item = (K, &'a V)> + 'a> = match root_value {
Some(value) => Box::new(std::iter::once((root_key, value)).chain(children)),
None => Box::new(children),
};
Self {
inner,
_marker: Default::default(),
}
}
fn new(node: Option<&'a VersionedNode<P, V>>) -> Self {
let Some(root_node) = node else {
return Self::empty();
};
let root_key = K::new_from_partial(&root_node.prefix);
let root_value = root_node.value();
if root_node.is_leaf() {
return Self {
inner: Box::new(std::iter::once((
root_key,
root_value.expect("corruption: missing data at leaf node during iteration"),
))),
_marker: Default::default(),
};
}
Self::from_root_and_children(root_key, root_value, VersionedIterInner::new(root_node))
}
fn new_with_prefix(node: Option<&'a VersionedNode<P, V>>, root_key: K) -> Self {
let Some(root_node) = node else {
return Self::empty();
};
let root_value = root_node.value();
if root_node.is_leaf() {
return Self {
inner: Box::new(std::iter::once((
root_key,
root_value.expect("corruption: missing data at leaf node during iteration"),
))),
_marker: Default::default(),
};
}
Self::from_root_and_children(
root_key.clone(),
root_value,
VersionedIterInner::from_node_and_key(root_node, root_key),
)
}
fn new_with_start_bound(node: Option<&'a VersionedNode<P, V>>, start_bound: Bound<K>) -> Self {
let Some(root_node) = node else {
return Self::empty();
};
let root_key = K::new_from_partial(&root_node.prefix);
let root_value = root_node.value();
let satisfies_start = match &start_bound {
Bound::Included(start_key) => {
VersionedIterInner::<K, P, V>::key_order(&root_key, start_key)
>= std::cmp::Ordering::Equal
}
Bound::Excluded(start_key) => {
VersionedIterInner::<K, P, V>::key_order(&root_key, start_key)
> std::cmp::Ordering::Equal
}
Bound::Unbounded => true,
};
if root_node.is_leaf() {
if satisfies_start {
return Self {
inner: Box::new(std::iter::once((
root_key,
root_value.expect("corruption: missing data at leaf node during iteration"),
))),
_marker: Default::default(),
};
}
return Self::empty();
}
let children = VersionedIterInner::new_with_start_bound(root_node, start_bound);
if satisfies_start {
return Self::from_root_and_children(root_key, root_value, children);
}
Self {
inner: Box::new(children),
_marker: Default::default(),
}
}
}
impl<'a, K: KeyTrait<PartialType = P>, P: Partial + 'a, V> Iterator for VersionedIter<'a, K, P, V> {
type Item = (K, &'a V);
fn next(&mut self) -> Option<Self::Item> {
self.inner.next()
}
}
impl<'a, K: KeyTrait<PartialType = P>, P: Partial + 'a, V> VersionedPrefixMatchIter<'a, K, P, V> {
fn new(node: Option<&'a VersionedNode<P, V>>, probe: K) -> Self {
Self {
cur_node: node,
probe,
cur_key: Vec::new(),
depth: 0,
}
}
}
impl<'a, K: KeyTrait<PartialType = P>, P: Partial + 'a, V> Iterator
for VersionedPrefixMatchIter<'a, K, P, V>
{
type Item = (K, &'a V);
fn next(&mut self) -> Option<Self::Item> {
loop {
let cur_node = self.cur_node.take()?;
let remaining_len = self.probe.length_at(self.depth);
let prefix_len = cur_node.prefix.len();
let prefix_common_match = cur_node.prefix.prefix_length_key(&self.probe, self.depth);
if prefix_common_match != prefix_len {
return None;
}
self.cur_key.extend_from_slice(cur_node.prefix.as_ref());
self.depth += prefix_len;
self.cur_node = if prefix_len == remaining_len {
None
} else {
cur_node
.seek_child(self.probe.at(self.depth))
.map(std::convert::AsRef::as_ref)
};
if let Some(value) = cur_node.value() {
return Some((K::new_from_slice(&self.cur_key), value));
}
}
}
}
impl<'a, K: KeyTrait<PartialType = P>, P: Partial + 'a, V> Iterator
for VersionedIterInner<'a, K, P, V>
{
type Item = (K, &'a V);
fn next(&mut self) -> Option<Self::Item> {
loop {
let (tree_depth, last_iter) = self.node_iter_stack.last_mut()?;
let tree_depth = *tree_depth;
self.cur_key.truncate(tree_depth);
let Some((_key, node)) = last_iter.next() else {
self.node_iter_stack.pop();
if let Some((parent_depth, _)) = self.node_iter_stack.last() {
self.cur_key.truncate(*parent_depth);
}
continue;
};
self.cur_key.extend_from_slice(node.prefix.as_ref());
let is_inner = node.is_inner();
if is_inner {
self.node_iter_stack.push((
tree_depth + node.prefix.len(),
VersionedIterFrameIter::Plain(node.iter()),
));
}
if let Some(value) = node.value() {
let key = K::new_from_slice(&self.cur_key);
if let Some(start_bound) = self.start_bound.as_ref() {
let satisfies_start = match start_bound {
Bound::Included(start_key) => {
VersionedIterInner::<K, P, V>::key_order(&key, start_key)
>= std::cmp::Ordering::Equal
}
Bound::Excluded(start_key) => {
VersionedIterInner::<K, P, V>::key_order(&key, start_key)
> std::cmp::Ordering::Equal
}
Bound::Unbounded => true,
};
if !satisfies_start {
continue;
}
self.start_bound = None;
}
return Some((key, value));
}
if !is_inner {
self.cur_key.truncate(tree_depth);
}
}
}
}
impl<'a, P: Partial + 'a, V> VersionedLendingIterInner<'a, P, V> {
fn cmp_segments_to_slice(segments: &[&[u8]], len: usize, slice: &[u8]) -> std::cmp::Ordering {
let mut offset = 0usize;
for segment in segments {
let remaining = &slice[offset..];
let common = segment.len().min(remaining.len());
match segment[..common].cmp(&remaining[..common]) {
std::cmp::Ordering::Equal => {}
ord => return ord,
}
if segment.len() != common {
return std::cmp::Ordering::Greater;
}
if remaining.len() != common {
return std::cmp::Ordering::Less;
}
offset += common;
}
len.cmp(&slice.len())
}
fn within_end_bound(&self) -> bool {
let Some((end_key, inclusive)) = self.end_bound.as_ref() else {
return true;
};
match Self::cmp_segments_to_slice(&self.cur_segments, self.cur_len, end_key) {
std::cmp::Ordering::Less => true,
std::cmp::Ordering::Equal => *inclusive,
std::cmp::Ordering::Greater => false,
}
}
fn lending_view<'view>(&'view self) -> LendingKeyView<'a, 'view> {
LendingKeyView::new(&self.cur_segments, self.cur_len)
}
fn visit_each<F>(&mut self, on_each: &mut F)
where
F: for<'view> FnMut(LendingKeyView<'a, 'view>, &'a V),
{
loop {
let next = {
let (segment_depth, key_len, last_iter) = match self.node_iter_stack.last_mut() {
Some(v) => v,
None => return,
};
let segment_depth = *segment_depth;
let key_len = *key_len;
self.cur_segments.truncate(segment_depth);
self.cur_len = key_len;
let Some((_key, node)) = last_iter.next() else {
self.node_iter_stack.pop();
if let Some((parent_segment_depth, parent_key_len, _)) =
self.node_iter_stack.last()
{
self.cur_segments.truncate(*parent_segment_depth);
self.cur_len = *parent_key_len;
}
continue;
};
let segment = node.prefix.as_ref();
if !segment.is_empty() {
self.cur_segments.push(segment);
self.cur_len += segment.len();
}
let is_inner = node.is_inner();
if is_inner {
self.node_iter_stack.push((
self.cur_segments.len(),
self.cur_len,
VersionedIterFrameIter::Plain(node.iter()),
));
}
Some((segment, is_inner, node.value()))
};
let Some((segment, is_inner, value)) = next else {
continue;
};
if let Some(value) = value {
if !self.within_end_bound() {
self.node_iter_stack.clear();
self.cur_segments.clear();
self.cur_len = 0;
return;
}
on_each(self.lending_view(), value);
}
if !is_inner && !segment.is_empty() {
self.cur_segments.pop();
self.cur_len -= segment.len();
}
}
}
fn build_positioned_stack<K: KeyTrait<PartialType = P>>(
node: &'a VersionedNode<P, V>,
seek_key: &K,
depth: usize,
) -> Vec<(usize, usize, VersionedIterFrameIter<'a, P, V>)> {
let root_segment_depth = usize::from(!node.prefix.as_ref().is_empty());
let prefix_common = node.prefix.prefix_length_key(seek_key, depth);
if prefix_common != node.prefix.len() {
let seek_remaining = seek_key.length_at(depth);
if prefix_common >= seek_remaining {
return vec![(
root_segment_depth,
node.prefix.len(),
VersionedIterFrameIter::Plain(node.iter()),
)];
}
let node_byte = node.prefix.at(prefix_common);
let seek_byte = seek_key.at(depth + prefix_common);
if node_byte < seek_byte {
return vec![];
}
return vec![(
root_segment_depth,
node.prefix.len(),
VersionedIterFrameIter::Plain(node.iter()),
)];
}
if seek_key.length_at(depth) == node.prefix.len() {
return vec![(
root_segment_depth,
node.prefix.len(),
VersionedIterFrameIter::Plain(node.iter()),
)];
}
let target_depth = depth + node.prefix.len();
let target_byte = seek_key.at(target_depth);
let mut iter = node.iter();
while let Some((key, child)) = iter.next() {
if key < target_byte {
continue;
}
return vec![(
root_segment_depth,
node.prefix.len(),
VersionedIterFrameIter::Leading {
first: Some((key, child)),
rest: iter,
},
)];
}
vec![]
}
fn for_each<F>(node: Option<&'a VersionedNode<P, V>>, mut on_each: F)
where
F: for<'view> FnMut(LendingKeyView<'a, 'view>, &'a V),
{
let Some(root_node) = node else {
return;
};
let root_segments = if root_node.prefix.is_empty() {
Vec::new()
} else {
vec![root_node.prefix.as_ref()]
};
let root_len = root_node.prefix.len();
if let Some(value) = root_node.value() {
on_each(LendingKeyView::new(&root_segments, root_len), value);
}
if root_node.is_inner() {
let mut inner = Self {
node_iter_stack: vec![(
root_segments.len(),
root_len,
VersionedIterFrameIter::Plain(root_node.iter()),
)],
cur_segments: root_segments,
cur_len: root_len,
end_bound: None,
};
inner.visit_each(&mut on_each);
}
}
fn for_each_with_prefix<F>(
node: Option<&'a VersionedNode<P, V>>,
root_segments: Vec<&'a [u8]>,
root_len: usize,
mut on_each: F,
) where
F: for<'view> FnMut(LendingKeyView<'a, 'view>, &'a V),
{
let Some(root_node) = node else {
return;
};
if let Some(value) = root_node.value() {
on_each(LendingKeyView::new(&root_segments, root_len), value);
}
if root_node.is_inner() {
let mut inner = Self {
node_iter_stack: vec![(
root_segments.len(),
root_len,
VersionedIterFrameIter::Plain(root_node.iter()),
)],
cur_segments: root_segments,
cur_len: root_len,
end_bound: None,
};
inner.visit_each(&mut on_each);
}
}
fn for_each_with_bounds<K, F>(
node: Option<&'a VersionedNode<P, V>>,
start_bound: Bound<K>,
end_bound: Bound<K>,
mut on_each: F,
) where
K: KeyTrait<PartialType = P>,
F: for<'view> FnMut(LendingKeyView<'a, 'view>, &'a V),
{
let Some(root_node) = node else {
return;
};
let end_bound_vec = match end_bound {
Bound::Included(key) => Some((key.as_ref().to_vec(), true)),
Bound::Excluded(key) => Some((key.as_ref().to_vec(), false)),
Bound::Unbounded => None,
};
let root_segments = if root_node.prefix.is_empty() {
Vec::new()
} else {
vec![root_node.prefix.as_ref()]
};
let root_len = root_node.prefix.len();
let root_view = LendingKeyView::new(&root_segments, root_len);
let satisfies_start = match &start_bound {
Bound::Included(start_key) => {
root_view.cmp_slice(start_key.as_ref()) >= std::cmp::Ordering::Equal
}
Bound::Excluded(start_key) => {
root_view.cmp_slice(start_key.as_ref()) > std::cmp::Ordering::Equal
}
Bound::Unbounded => true,
};
let satisfies_end = match end_bound_vec.as_ref() {
Some((end_key, inclusive)) => match root_view.cmp_slice(end_key) {
std::cmp::Ordering::Less => true,
std::cmp::Ordering::Equal => *inclusive,
std::cmp::Ordering::Greater => false,
},
None => true,
};
if !satisfies_end {
return;
}
if let Some(value) = root_node.value()
&& satisfies_start
{
on_each(root_view, value);
}
if !root_node.is_inner() {
return;
}
let seek_key = match &start_bound {
Bound::Included(key) | Bound::Excluded(key) => Some(key),
Bound::Unbounded => None,
};
let mut inner = if let Some(seek_key) = seek_key {
Self {
node_iter_stack: Self::build_positioned_stack(root_node, seek_key, 0),
cur_segments: root_segments,
cur_len: root_len,
end_bound: end_bound_vec,
}
} else {
Self {
node_iter_stack: vec![(
root_segments.len(),
root_len,
VersionedIterFrameIter::Plain(root_node.iter()),
)],
cur_segments: root_segments,
cur_len: root_len,
end_bound: end_bound_vec,
}
};
inner.visit_each(&mut on_each);
}
}
impl<'a, P: Partial + 'a, V> VersionedValuesIter<'a, P, V> {
fn new(node: Option<&'a VersionedNode<P, V>>) -> Self {
let Some(root_node) = node else {
return Self {
root_value: None,
node_iter_stack: Vec::new(),
};
};
Self {
root_value: root_node.value(),
node_iter_stack: vec![root_node.iter()],
}
}
}
impl<'a, P: Partial + 'a, V> Iterator for VersionedValuesIter<'a, P, V> {
type Item = &'a V;
fn next(&mut self) -> Option<Self::Item> {
if let Some(value) = self.root_value.take() {
return Some(value);
}
loop {
let last_iter = self.node_iter_stack.last_mut()?;
let Some((_key, node)) = last_iter.next() else {
self.node_iter_stack.pop();
continue;
};
if node.is_inner() {
self.node_iter_stack.push(node.iter());
}
if let Some(value) = node.value() {
return Some(value);
}
}
}
}
impl<'a, K: KeyTrait + 'a, V> Iterator for VersionedRange<'a, K, V> {
type Item = (K, &'a V);
fn next(&mut self) -> Option<Self::Item> {
let next = self.iter.next()?;
match &self.end {
Bound::Included(end_key) => match next.0.cmp(end_key) {
std::cmp::Ordering::Less | std::cmp::Ordering::Equal => Some(next),
std::cmp::Ordering::Greater => None,
},
Bound::Excluded(end_key) => match next.0.cmp(end_key) {
std::cmp::Ordering::Less => Some(next),
std::cmp::Ordering::Equal | std::cmp::Ordering::Greater => None,
},
Bound::Unbounded => Some(next),
}
}
}
impl<KeyType, ValueType> VersionedAdaptiveRadixTree<KeyType, ValueType>
where
KeyType: KeyTrait,
ValueType: Clone,
{
fn get_iterate<'a>(
cur_node: &'a VersionedNode<KeyType::PartialType, ValueType>,
key: &KeyType,
) -> Option<&'a ValueType> {
let mut cur_node = cur_node;
let mut depth = 0;
loop {
let prefix_common_match = cur_node.prefix.prefix_length_key(key, depth);
if prefix_common_match != cur_node.prefix.len() {
return None;
}
if cur_node.prefix.len() == key.length_at(depth) {
return cur_node.value();
}
let k = key.at(depth + cur_node.prefix.len());
depth += cur_node.prefix.len();
cur_node = cur_node.seek_child(k)?.as_ref();
}
}
fn prefix_match_for_each_impl<'a, F>(
cur_node: &'a VersionedNode<KeyType::PartialType, ValueType>,
key: &KeyType,
mut on_match: F,
) where
F: FnMut(&[u8], &'a ValueType),
{
let mut cur_node = cur_node;
let mut depth = 0;
let key_bytes = key.as_ref();
loop {
let prefix_common_match = cur_node.prefix.prefix_length_key(key, depth);
if prefix_common_match != cur_node.prefix.len() {
return;
}
let matched_len = depth + cur_node.prefix.len();
if let Some(value) = cur_node.value() {
on_match(&key_bytes[..matched_len], value);
}
if cur_node.prefix.len() == key.length_at(depth) {
return;
}
let k = key.at(depth + cur_node.prefix.len());
depth += cur_node.prefix.len();
let Some(child) = cur_node.seek_child(k) else {
return;
};
cur_node = child.as_ref();
}
}
fn intersect_nodes<'a, F>(
left: &'a VersionedNode<KeyType::PartialType, ValueType>,
mut left_offset: usize,
right: &'a VersionedNode<KeyType::PartialType, ValueType>,
mut right_offset: usize,
key_buf: &mut Vec<u8>,
on_match: &mut F,
) where
F: FnMut(KeyType, &'a ValueType, &'a ValueType),
{
let restore_len = key_buf.len();
let left_prefix = left.prefix.as_ref();
let right_prefix = right.prefix.as_ref();
while left_offset < left_prefix.len() && right_offset < right_prefix.len() {
let left_byte = left_prefix[left_offset];
let right_byte = right_prefix[right_offset];
if left_byte != right_byte {
key_buf.truncate(restore_len);
return;
}
key_buf.push(left_byte);
left_offset += 1;
right_offset += 1;
}
if left_offset < left_prefix.len() {
if !right.is_inner() {
key_buf.truncate(restore_len);
return;
}
let edge = left_prefix[left_offset];
let Some(right_child) = right.seek_child(edge) else {
key_buf.truncate(restore_len);
return;
};
key_buf.push(edge);
Self::intersect_nodes(
left,
left_offset + 1,
right_child.as_ref(),
1,
key_buf,
on_match,
);
key_buf.truncate(restore_len);
return;
}
if right_offset < right_prefix.len() {
if !left.is_inner() {
key_buf.truncate(restore_len);
return;
}
let edge = right_prefix[right_offset];
let Some(left_child) = left.seek_child(edge) else {
key_buf.truncate(restore_len);
return;
};
key_buf.push(edge);
Self::intersect_nodes(
left_child.as_ref(),
1,
right,
right_offset + 1,
key_buf,
on_match,
);
key_buf.truncate(restore_len);
return;
}
if let (Some(left_value), Some(right_value)) = (left.value(), right.value()) {
on_match(
KeyType::new_from_slice(key_buf.as_slice()),
left_value,
right_value,
);
}
if left.is_inner() && right.is_inner() {
if left.num_children() <= right.num_children() {
for (edge, left_child) in left.iter() {
if let Some(right_child) = right.seek_child(edge) {
Self::intersect_nodes(
left_child,
0,
right_child.as_ref(),
0,
key_buf,
on_match,
);
}
}
} else {
for (edge, right_child) in right.iter() {
if let Some(left_child) = left.seek_child(edge) {
Self::intersect_nodes(
left_child.as_ref(),
0,
right_child,
0,
key_buf,
on_match,
);
}
}
}
}
key_buf.truncate(restore_len);
}
fn intersect_nodes_lending<'a, F>(
left: &'a VersionedNode<KeyType::PartialType, ValueType>,
mut left_offset: usize,
right: &'a VersionedNode<KeyType::PartialType, ValueType>,
mut right_offset: usize,
key_segments: &mut Vec<&'a [u8]>,
key_len: &mut usize,
on_match: &mut F,
) where
F: for<'view> FnMut(LendingKeyView<'a, 'view>, &'a ValueType, &'a ValueType),
{
let restore_segments = key_segments.len();
let restore_len = *key_len;
let left_prefix = left.prefix.as_ref();
let right_prefix = right.prefix.as_ref();
let matched_left_start = left_offset;
while left_offset < left_prefix.len() && right_offset < right_prefix.len() {
if left_prefix[left_offset] != right_prefix[right_offset] {
key_segments.truncate(restore_segments);
*key_len = restore_len;
return;
}
left_offset += 1;
right_offset += 1;
}
if left_offset > matched_left_start {
let matched = &left_prefix[matched_left_start..left_offset];
key_segments.push(matched);
*key_len += matched.len();
}
if left_offset < left_prefix.len() {
if !right.is_inner() {
key_segments.truncate(restore_segments);
*key_len = restore_len;
return;
}
let edge = left_prefix[left_offset];
let Some(right_child) = right.seek_child(edge) else {
key_segments.truncate(restore_segments);
*key_len = restore_len;
return;
};
let edge_segment = &left_prefix[left_offset..left_offset + 1];
key_segments.push(edge_segment);
*key_len += 1;
Self::intersect_nodes_lending(
left,
left_offset + 1,
right_child.as_ref(),
1,
key_segments,
key_len,
on_match,
);
key_segments.truncate(restore_segments);
*key_len = restore_len;
return;
}
if right_offset < right_prefix.len() {
if !left.is_inner() {
key_segments.truncate(restore_segments);
*key_len = restore_len;
return;
}
let edge = right_prefix[right_offset];
let Some(left_child) = left.seek_child(edge) else {
key_segments.truncate(restore_segments);
*key_len = restore_len;
return;
};
let edge_segment = &right_prefix[right_offset..right_offset + 1];
key_segments.push(edge_segment);
*key_len += 1;
Self::intersect_nodes_lending(
left_child.as_ref(),
1,
right,
right_offset + 1,
key_segments,
key_len,
on_match,
);
key_segments.truncate(restore_segments);
*key_len = restore_len;
return;
}
if let (Some(left_value), Some(right_value)) = (left.value(), right.value()) {
on_match(
LendingKeyView::new(key_segments, *key_len),
left_value,
right_value,
);
}
if left.is_inner() && right.is_inner() {
if left.num_children() <= right.num_children() {
for (edge, left_child) in left.iter() {
if let Some(right_child) = right.seek_child(edge) {
Self::intersect_nodes_lending(
left_child,
0,
right_child.as_ref(),
0,
key_segments,
key_len,
on_match,
);
}
}
} else {
for (edge, right_child) in right.iter() {
if let Some(left_child) = left.seek_child(edge) {
Self::intersect_nodes_lending(
left_child.as_ref(),
0,
right_child,
0,
key_segments,
key_len,
on_match,
);
}
}
}
}
key_segments.truncate(restore_segments);
*key_len = restore_len;
}
fn intersect_nodes_values<'a, F>(
left: &'a VersionedNode<KeyType::PartialType, ValueType>,
mut left_offset: usize,
right: &'a VersionedNode<KeyType::PartialType, ValueType>,
mut right_offset: usize,
on_match: &mut F,
) where
F: FnMut(&'a ValueType, &'a ValueType),
{
let left_prefix = left.prefix.as_ref();
let right_prefix = right.prefix.as_ref();
while left_offset < left_prefix.len() && right_offset < right_prefix.len() {
if left_prefix[left_offset] != right_prefix[right_offset] {
return;
}
left_offset += 1;
right_offset += 1;
}
if left_offset < left_prefix.len() {
if !right.is_inner() {
return;
}
let edge = left_prefix[left_offset];
let Some(right_child) = right.seek_child(edge) else {
return;
};
Self::intersect_nodes_values(left, left_offset + 1, right_child.as_ref(), 1, on_match);
return;
}
if right_offset < right_prefix.len() {
if !left.is_inner() {
return;
}
let edge = right_prefix[right_offset];
let Some(left_child) = left.seek_child(edge) else {
return;
};
Self::intersect_nodes_values(left_child.as_ref(), 1, right, right_offset + 1, on_match);
return;
}
if let (Some(left_value), Some(right_value)) = (left.value(), right.value()) {
on_match(left_value, right_value);
}
if left.is_inner() && right.is_inner() {
if left.num_children() <= right.num_children() {
for (edge, left_child) in left.iter() {
if let Some(right_child) = right.seek_child(edge) {
Self::intersect_nodes_values(
left_child,
0,
right_child.as_ref(),
0,
on_match,
);
}
}
} else {
for (edge, right_child) in right.iter() {
if let Some(left_child) = left.seek_child(edge) {
Self::intersect_nodes_values(
left_child.as_ref(),
0,
right_child,
0,
on_match,
);
}
}
}
}
}
fn find_prefix_subtree<'a>(
cur_node: &'a VersionedNode<KeyType::PartialType, ValueType>,
prefix: &KeyType,
) -> Option<(&'a VersionedNode<KeyType::PartialType, ValueType>, KeyType)> {
let mut cur_node = cur_node;
let mut cur_key = cur_node.prefix.as_ref().to_vec();
let mut depth = 0;
loop {
let prefix_common_match = cur_node.prefix.prefix_length_key(prefix, depth);
if prefix_common_match != cur_node.prefix.len() {
if prefix_common_match == prefix.length_at(depth) {
return Some((cur_node, KeyType::new_from_slice(&cur_key)));
}
return None;
}
if cur_node.prefix.len() == prefix.length_at(depth) {
return Some((cur_node, KeyType::new_from_slice(&cur_key)));
}
let key = prefix.at(depth + cur_node.prefix.len());
depth += cur_node.prefix.len();
cur_node = cur_node.seek_child(key)?.as_ref();
cur_key.extend_from_slice(cur_node.prefix.as_ref());
}
}
fn find_prefix_subtree_view<'a>(
cur_node: &'a VersionedNode<KeyType::PartialType, ValueType>,
prefix: &KeyType,
) -> Option<VersionedPrefixSubtreeView<'a, KeyType::PartialType, ValueType>> {
let mut cur_node = cur_node;
let mut cur_segments = if cur_node.prefix.is_empty() {
Vec::new()
} else {
vec![cur_node.prefix.as_ref()]
};
let mut cur_len = cur_node.prefix.len();
let mut depth = 0;
loop {
let prefix_common_match = cur_node.prefix.prefix_length_key(prefix, depth);
if prefix_common_match != cur_node.prefix.len() {
if prefix_common_match == prefix.length_at(depth) {
return Some((cur_node, cur_segments, cur_len));
}
return None;
}
if cur_node.prefix.len() == prefix.length_at(depth) {
return Some((cur_node, cur_segments, cur_len));
}
let key = prefix.at(depth + cur_node.prefix.len());
depth += cur_node.prefix.len();
cur_node = cur_node.seek_child(key)?.as_ref();
let segment = cur_node.prefix.as_ref();
if !segment.is_empty() {
cur_segments.push(segment);
cur_len += segment.len();
}
}
}
fn ensure_cow_node(
node: Arc<VersionedNode<KeyType::PartialType, ValueType>>,
target_version: u64,
) -> Arc<VersionedNode<KeyType::PartialType, ValueType>> {
if node.version == target_version {
node
} else {
match Arc::try_unwrap(node) {
Ok(mut owned_node) => {
owned_node.version = target_version;
Arc::new(owned_node)
}
Err(shared_node) => {
Arc::new(shared_node.cow_clone_inner(target_version))
}
}
}
}
fn insert_recurse(
cur_node: Arc<VersionedNode<KeyType::PartialType, ValueType>>,
key: &KeyType,
value: ValueType,
depth: usize,
version: u64,
old_value_out: Option<&mut Option<ValueType>>,
) -> (Arc<VersionedNode<KeyType::PartialType, ValueType>>, bool) {
let longest_common_prefix = cur_node.prefix.prefix_length_key(key, depth);
let is_prefix_match =
min(cur_node.prefix.len(), key.length_at(depth)) == longest_common_prefix;
if is_prefix_match && cur_node.prefix.len() == key.length_at(depth) {
let new_node = Self::ensure_cow_node(cur_node, version);
let mut new_node = match Arc::try_unwrap(new_node) {
Ok(owned) => owned,
Err(_) => panic!("ensure_cow_node should have given us exclusive ownership"),
};
let old_value = new_node.value.replace(value);
let was_replaced = old_value.is_some();
if let (Some(old_value_out), Some(old_value)) = (old_value_out, old_value) {
*old_value_out = Some(old_value);
}
return (Arc::new(new_node), was_replaced);
}
if is_prefix_match && cur_node.prefix.len() > key.length_at(depth) {
let mut existing_node = cur_node.cow_clone_inner(version);
let old_prefix = existing_node.prefix.clone();
existing_node.prefix = old_prefix.partial_after(longest_common_prefix);
let mut new_parent =
VersionedNode::new_inner(old_prefix.partial_before(longest_common_prefix), version);
new_parent.value = Some(value);
let edge = old_prefix.at(longest_common_prefix);
new_parent.add_child(edge, Arc::new(existing_node));
return (Arc::new(new_parent), false);
}
if !is_prefix_match {
let mut new_inner = VersionedNode::new_inner(
cur_node.prefix.partial_before(longest_common_prefix),
version,
);
let k1 = cur_node.prefix.at(longest_common_prefix);
let k2 = key.at(depth + longest_common_prefix);
let mut existing_node_clone = cur_node.cow_clone_inner(version);
existing_node_clone.prefix = cur_node.prefix.partial_after(longest_common_prefix);
let existing_arc = Arc::new(existing_node_clone);
let new_leaf = Arc::new(VersionedNode::new_leaf(
key.to_partial(depth + longest_common_prefix),
value,
version,
));
match &mut new_inner.content {
VersionedContent::Node4(km) => {
km.add_child(k1, existing_arc);
km.add_child(k2, new_leaf);
}
_ => unreachable!(),
}
return (Arc::new(new_inner), false);
}
if cur_node.is_leaf() {
let edge = key.at(depth + longest_common_prefix);
let new_leaf = Arc::new(VersionedNode::new_leaf(
key.to_partial(depth + longest_common_prefix),
value,
version,
));
let new_node = Self::ensure_cow_node(cur_node, version);
let mut new_node = match Arc::try_unwrap(new_node) {
Ok(owned) => owned,
Err(_) => panic!("ensure_cow_node should have given us exclusive ownership"),
};
new_node.add_child(edge, new_leaf);
return (Arc::new(new_node), false);
}
let k = key.at(depth + cur_node.prefix.len());
let prefix_len = cur_node.prefix.len();
let new_node = Self::ensure_cow_node(cur_node, version);
let mut new_node_mut = match Arc::try_unwrap(new_node) {
Ok(owned) => owned,
Err(_) => panic!("ensure_cow_node should have given us exclusive ownership"),
};
if let Some(child) = new_node_mut.delete_child(k) {
let (new_child, was_replaced) = Self::insert_recurse(
child,
key,
value,
depth + prefix_len,
version,
old_value_out,
);
new_node_mut.add_child(k, new_child);
(Arc::new(new_node_mut), was_replaced)
} else {
let new_leaf = Arc::new(VersionedNode::new_leaf(
key.to_partial(depth + prefix_len),
value,
version,
));
new_node_mut.add_child(k, new_leaf);
(Arc::new(new_node_mut), false)
}
}
fn remove_recurse(
cur_node: Arc<VersionedNode<KeyType::PartialType, ValueType>>,
key: &KeyType,
depth: usize,
version: u64,
) -> Option<RemoveResult<KeyType::PartialType, ValueType>> {
let prefix_common_match = cur_node.prefix.prefix_length_key(key, depth);
if prefix_common_match != cur_node.prefix.len() {
return None;
}
if cur_node.prefix.len() == key.length_at(depth) {
if cur_node.is_leaf() {
let removed_value = cur_node.value()?.clone();
return Some((None, removed_value));
}
let new_node = Self::ensure_cow_node(cur_node, version);
let mut new_node = match Arc::try_unwrap(new_node) {
Ok(owned) => owned,
Err(_) => panic!("ensure_cow_node should have given us exclusive ownership"),
};
let removed_value = new_node.value.take()?;
if new_node.num_children() == 0 && new_node.value.is_none() {
return Some((None, removed_value));
}
return Some((Some(Arc::new(new_node)), removed_value));
}
if cur_node.is_leaf() {
return None;
}
let k = key.at(depth + cur_node.prefix.len());
let prefix_len = cur_node.prefix.len();
let new_node = Self::ensure_cow_node(cur_node, version);
let mut new_node_mut = match Arc::try_unwrap(new_node) {
Ok(owned) => owned,
Err(_) => panic!("ensure_cow_node should have given us exclusive ownership"),
};
let child = new_node_mut.delete_child(k)?;
let (new_child_opt, removed_value) =
Self::remove_recurse(child, key, depth + prefix_len, version)
.expect("prechecked key should be removable");
if let Some(new_child) = new_child_opt {
new_node_mut.add_child(k, new_child);
}
if new_node_mut.num_children() == 0 && new_node_mut.value.is_none() {
return Some((None, removed_value));
}
Some((Some(Arc::new(new_node_mut)), removed_value))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::keys::{array_key::ArrayKey, overflow_key::OverflowKey};
use proptest::prelude::*;
use std::mem::size_of;
type DenseSequentialKey = OverflowKey<8, 4>;
fn dense_sequential_key(i: u32) -> DenseSequentialKey {
let mut bytes = [0; 5];
bytes[0] = 8;
bytes[1..].copy_from_slice(&i.to_be_bytes());
<DenseSequentialKey as crate::keys::KeyTrait>::new_from_slice(&bytes)
}
#[test]
fn boxed_versioned_content_keeps_node_header_small() {
type P = <ArrayKey<16> as crate::keys::KeyTrait>::PartialType;
let content_size = size_of::<VersionedContent<P, u64>>();
let node_size = size_of::<VersionedNode<P, u64>>();
println!("VersionedContent<P, u64> = {content_size} bytes");
println!("VersionedNode<P, u64> = {node_size} bytes");
assert!(content_size <= 2 * size_of::<usize>());
assert!(node_size <= 80);
}
#[derive(Clone, Debug)]
enum VersionedOp {
Get {
key: u8,
},
Insert {
key: u8,
value: u16,
},
Remove {
key: u8,
},
Snapshot,
SnapshotInsert {
snapshot_idx: u8,
key: u8,
value: u16,
},
SnapshotRemove {
snapshot_idx: u8,
key: u8,
},
}
#[derive(Clone, Debug)]
enum DenseVersionedOp {
Insert {
key_idx: u16,
value: u16,
},
Remove {
key_idx: u16,
},
Snapshot,
SnapshotInsert {
snapshot_idx: u8,
key_idx: u16,
value: u16,
},
SnapshotRemove {
snapshot_idx: u8,
key_idx: u16,
},
}
fn versioned_op_strategy() -> impl Strategy<Value = VersionedOp> {
prop_oneof![
any::<u8>().prop_map(|key| VersionedOp::Get { key }),
(any::<u8>(), any::<u16>()).prop_map(|(key, value)| VersionedOp::Insert { key, value }),
any::<u8>().prop_map(|key| VersionedOp::Remove { key }),
Just(VersionedOp::Snapshot),
(any::<u8>(), any::<u8>(), any::<u16>()).prop_map(|(snapshot_idx, key, value)| {
VersionedOp::SnapshotInsert {
snapshot_idx,
key,
value,
}
}),
(any::<u8>(), any::<u8>())
.prop_map(|(snapshot_idx, key)| VersionedOp::SnapshotRemove { snapshot_idx, key }),
]
}
fn dense_versioned_op_strategy() -> impl Strategy<Value = DenseVersionedOp> {
prop_oneof![
(0u16..512, any::<u16>())
.prop_map(|(key_idx, value)| { DenseVersionedOp::Insert { key_idx, value } }),
(0u16..512).prop_map(|key_idx| DenseVersionedOp::Remove { key_idx }),
Just(DenseVersionedOp::Snapshot),
(any::<u8>(), 0u16..512, any::<u16>()).prop_map(|(snapshot_idx, key_idx, value)| {
DenseVersionedOp::SnapshotInsert {
snapshot_idx,
key_idx,
value,
}
},),
(any::<u8>(), 0u16..512).prop_map(|(snapshot_idx, key_idx)| {
DenseVersionedOp::SnapshotRemove {
snapshot_idx,
key_idx,
}
}),
]
}
fn assert_versioned_tree_matches_map(
tree: &VersionedAdaptiveRadixTree<ArrayKey<16>, u16>,
map: &std::collections::BTreeMap<u8, u16>,
) {
for key in 0u8..=u8::MAX {
assert_eq!(
tree.get(key).copied(),
map.get(&key).copied(),
"mismatch at key {key}"
);
}
}
fn assert_dense_tree_matches_map(
tree: &VersionedAdaptiveRadixTree<DenseSequentialKey, u16>,
map: &std::collections::BTreeMap<u16, u16>,
key_limit: u16,
) {
for key_idx in 0..key_limit {
let key = dense_sequential_key(key_idx as u32);
assert_eq!(
tree.get_k(&key).copied(),
map.get(&key_idx).copied(),
"mismatch at dense sequential key index {key_idx}"
);
}
}
proptest! {
#[test]
fn prop_snapshot_operations_match_reference_model(
ops in proptest::collection::vec(versioned_op_strategy(), 0..96)
) {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, u16>::new();
let mut map = std::collections::BTreeMap::<u8, u16>::new();
let mut snapshots = Vec::new();
let mut snapshot_maps = Vec::new();
for op in ops {
match op {
VersionedOp::Get { key } => {
prop_assert_eq!(tree.get(key).copied(), map.get(&key).copied());
}
VersionedOp::Insert { key, value } => {
let expected_replaced = map.insert(key, value).is_some();
let actual_replaced = tree.insert(key, value);
prop_assert_eq!(actual_replaced, expected_replaced);
prop_assert_eq!(tree.get(key).copied(), map.get(&key).copied());
}
VersionedOp::Remove { key } => {
let expected_removed = map.remove(&key);
let actual_removed = tree.remove(key);
prop_assert_eq!(actual_removed, expected_removed);
prop_assert_eq!(tree.get(key).copied(), map.get(&key).copied());
}
VersionedOp::Snapshot => {
snapshots.push(tree.snapshot());
snapshot_maps.push(map.clone());
}
VersionedOp::SnapshotInsert {
snapshot_idx,
key,
value,
} => {
if !snapshots.is_empty() {
let idx = snapshot_idx as usize % snapshots.len();
let expected_replaced = snapshot_maps[idx].insert(key, value).is_some();
let actual_replaced = snapshots[idx].insert(key, value);
prop_assert_eq!(actual_replaced, expected_replaced);
prop_assert_eq!(
snapshots[idx].get(key).copied(),
snapshot_maps[idx].get(&key).copied()
);
prop_assert_eq!(tree.get(key).copied(), map.get(&key).copied());
}
}
VersionedOp::SnapshotRemove { snapshot_idx, key } => {
if !snapshots.is_empty() {
let idx = snapshot_idx as usize % snapshots.len();
let expected_removed = snapshot_maps[idx].remove(&key);
let actual_removed = snapshots[idx].remove(key);
prop_assert_eq!(actual_removed, expected_removed);
prop_assert_eq!(
snapshots[idx].get(key).copied(),
snapshot_maps[idx].get(&key).copied()
);
prop_assert_eq!(tree.get(key).copied(), map.get(&key).copied());
}
}
}
}
assert_versioned_tree_matches_map(&tree, &map);
for (snapshot, snapshot_map) in snapshots.iter().zip(snapshot_maps.iter()) {
assert_versioned_tree_matches_map(snapshot, snapshot_map);
}
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(64))]
#[test]
fn prop_dense_sequential_snapshots_remain_isolated(
ops in proptest::collection::vec(dense_versioned_op_strategy(), 0..96)
) {
let mut tree = VersionedAdaptiveRadixTree::<DenseSequentialKey, u16>::new();
let mut map = std::collections::BTreeMap::<u16, u16>::new();
let mut snapshots = Vec::new();
let mut snapshot_maps = Vec::new();
for key_idx in 0u16..256 {
let key = dense_sequential_key(key_idx as u32);
prop_assert!(!tree.insert_k(&key, key_idx));
map.insert(key_idx, key_idx);
}
for op in ops {
match op {
DenseVersionedOp::Insert { key_idx, value } => {
let key = dense_sequential_key(key_idx as u32);
let expected_replaced = map.insert(key_idx, value).is_some();
let actual_replaced = tree.insert_k(&key, value);
prop_assert_eq!(actual_replaced, expected_replaced);
prop_assert_eq!(tree.get_k(&key).copied(), map.get(&key_idx).copied());
}
DenseVersionedOp::Remove { key_idx } => {
let key = dense_sequential_key(key_idx as u32);
let expected_removed = map.remove(&key_idx);
let actual_removed = tree.remove_k(&key);
prop_assert_eq!(actual_removed, expected_removed);
prop_assert_eq!(tree.get_k(&key).copied(), map.get(&key_idx).copied());
}
DenseVersionedOp::Snapshot => {
snapshots.push(tree.snapshot());
snapshot_maps.push(map.clone());
}
DenseVersionedOp::SnapshotInsert {
snapshot_idx,
key_idx,
value,
} => {
if !snapshots.is_empty() {
let idx = snapshot_idx as usize % snapshots.len();
let key = dense_sequential_key(key_idx as u32);
let expected_replaced =
snapshot_maps[idx].insert(key_idx, value).is_some();
let actual_replaced = snapshots[idx].insert_k(&key, value);
prop_assert_eq!(actual_replaced, expected_replaced);
prop_assert_eq!(
snapshots[idx].get_k(&key).copied(),
snapshot_maps[idx].get(&key_idx).copied()
);
prop_assert_eq!(tree.get_k(&key).copied(), map.get(&key_idx).copied());
}
}
DenseVersionedOp::SnapshotRemove {
snapshot_idx,
key_idx,
} => {
if !snapshots.is_empty() {
let idx = snapshot_idx as usize % snapshots.len();
let key = dense_sequential_key(key_idx as u32);
let expected_removed = snapshot_maps[idx].remove(&key_idx);
let actual_removed = snapshots[idx].remove_k(&key);
prop_assert_eq!(actual_removed, expected_removed);
prop_assert_eq!(
snapshots[idx].get_k(&key).copied(),
snapshot_maps[idx].get(&key_idx).copied()
);
prop_assert_eq!(tree.get_k(&key).copied(), map.get(&key_idx).copied());
}
}
}
}
assert_dense_tree_matches_map(&tree, &map, 512);
for (snapshot, snapshot_map) in snapshots.iter().zip(snapshot_maps.iter()) {
assert_dense_tree_matches_map(snapshot, snapshot_map, 512);
}
}
}
#[test]
fn test_basic_snapshot() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
tree.insert("key1", 1);
assert_eq!(tree.get("key1"), Some(&1));
let snapshot = tree.snapshot();
assert_eq!(snapshot.get("key1"), Some(&1));
assert_eq!(snapshot.version(), tree.version() + 1);
}
#[test]
fn test_independent_mutations() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
tree.insert("key1", 1);
let mut snapshot = tree.snapshot();
tree.insert("key2", 2);
snapshot.insert("key3", 3);
assert_eq!(tree.get("key2"), Some(&2));
assert_eq!(tree.get("key3"), None);
assert_eq!(snapshot.get("key2"), None);
assert_eq!(snapshot.get("key3"), Some(&3));
assert_eq!(tree.get("key1"), Some(&1));
assert_eq!(snapshot.get("key1"), Some(&1));
}
#[test]
fn test_node_growth() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
for i in 0..10 {
let key = format!("key{i:02}");
tree.insert(key, i);
}
for i in 0..10 {
let key = format!("key{i:02}");
assert_eq!(tree.get(&key), Some(&i));
}
let snapshot = tree.snapshot();
for i in 10..20 {
let key = format!("key{i:02}");
tree.insert(key, i);
}
for i in 10..20 {
let key = format!("key{i:02}");
assert_eq!(snapshot.get(&key), None);
assert_eq!(tree.get(&key), Some(&i));
}
for i in 0..10 {
let key = format!("key{i:02}");
assert_eq!(snapshot.get(&key), Some(&i));
}
}
#[test]
fn test_structural_sharing() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
for i in 0..20 {
let key = format!("shared_key_{i:02}");
tree.insert(key, i);
}
let snapshot1 = tree.snapshot();
let snapshot2 = tree.snapshot();
let snapshot3 = tree.snapshot();
if let Some(root) = &tree.root {
let strong_count = Arc::strong_count(root);
assert_eq!(
strong_count, 4,
"Root should be shared between original and 3 snapshots"
);
}
tree.insert("new_key", 999);
if let (Some(s1_root), Some(s2_root), Some(s3_root)) =
(&snapshot1.root, &snapshot2.root, &snapshot3.root)
{
assert!(
Arc::ptr_eq(s1_root, s2_root),
"Snapshot1 and Snapshot2 should share root"
);
assert!(
Arc::ptr_eq(s2_root, s3_root),
"Snapshot2 and Snapshot3 should share root"
);
let shared_count = Arc::strong_count(s1_root);
assert_eq!(
shared_count, 3,
"Shared root should have 3 references after original tree CoW"
);
}
if let Some(orig_root) = &tree.root {
let orig_count = Arc::strong_count(orig_root);
assert_eq!(
orig_count, 1,
"Original tree should have exclusive ownership of new root"
);
}
assert_eq!(snapshot1.get("new_key"), None);
assert_eq!(snapshot2.get("new_key"), None);
assert_eq!(snapshot3.get("new_key"), None);
assert_eq!(tree.get("new_key"), Some(&999));
for i in 0..20 {
let key = format!("shared_key_{i:02}");
assert_eq!(tree.get(&key), Some(&i));
assert_eq!(snapshot1.get(&key), Some(&i));
assert_eq!(snapshot2.get(&key), Some(&i));
assert_eq!(snapshot3.get(&key), Some(&i));
}
}
#[test]
fn test_snapshot_cleanup() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
for i in 0..10i32 {
tree.insert(i, i * 10);
}
let initial_root_refs = if let Some(root) = &tree.root {
Arc::strong_count(root)
} else {
panic!("Tree should have a root");
};
let snapshot1 = tree.snapshot();
let snapshot2 = tree.snapshot();
{
let _snapshot3 = tree.snapshot(); }
let with_snapshots_refs = if let Some(root) = &tree.root {
Arc::strong_count(root)
} else {
panic!("Tree should have a root");
};
assert!(
with_snapshots_refs > initial_root_refs,
"Root should have more references with snapshots"
);
drop(snapshot2);
let after_drops_refs = if let Some(root) = &tree.root {
Arc::strong_count(root)
} else {
panic!("Tree should have a root");
};
assert!(
after_drops_refs < with_snapshots_refs,
"Root should have fewer references after dropping snapshots"
);
assert_eq!(
after_drops_refs, 2,
"Should have exactly 2 references: tree + snapshot1"
);
assert!(snapshot1.get(0).is_some());
assert!(snapshot1.get(5).is_some());
assert!(snapshot1.get(9).is_some());
drop(snapshot1);
let final_refs = if let Some(root) = &tree.root {
Arc::strong_count(root)
} else {
panic!("Tree should have a root");
};
assert_eq!(
final_refs, 1,
"Tree should have exclusive ownership after all snapshots dropped"
);
for i in 0..10i32 {
assert_eq!(tree.get(i), Some(&(i * 10)));
}
}
#[test]
fn test_signed_integer_keys() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
tree.insert(-5i32, -50);
tree.insert(0i32, 0);
tree.insert(1i32, 10);
tree.insert(8i32, 80);
tree.insert(-1i32, -10);
let snapshot = tree.snapshot();
assert_eq!(tree.get(-5i32), Some(&-50));
assert_eq!(tree.get(-1i32), Some(&-10));
assert_eq!(tree.get(0i32), Some(&0));
assert_eq!(tree.get(1i32), Some(&10));
assert_eq!(tree.get(8i32), Some(&80));
assert_eq!(snapshot.get(-5i32), Some(&-50));
assert_eq!(snapshot.get(-1i32), Some(&-10));
assert_eq!(snapshot.get(0i32), Some(&0));
assert_eq!(snapshot.get(1i32), Some(&10));
assert_eq!(snapshot.get(8i32), Some(&80));
assert_eq!(tree.get(99i32), None);
assert_eq!(snapshot.get(99i32), None);
}
#[test]
fn test_deep_structural_sharing() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<32>, i32>::new();
let prefixes = [
"user",
"user_profile",
"user_settings",
"system",
"system_config",
];
for (i, prefix) in prefixes.iter().enumerate() {
for j in 0..5 {
let key = format!("{prefix}_{j:02}");
tree.insert(key, (i * 100 + j) as i32);
}
}
let snapshot = tree.snapshot();
tree.insert("user_00", 999);
assert_eq!(tree.get("user_00"), Some(&999));
assert_eq!(snapshot.get("user_00"), Some(&0));
for (i, prefix) in prefixes.iter().enumerate() {
for j in 0..5 {
let key = format!("{prefix}_{j:02}");
if key != "user_00" {
let expected_value = (i * 100 + j) as i32;
assert_eq!(tree.get(&key), Some(&expected_value));
assert_eq!(snapshot.get(&key), Some(&expected_value));
}
}
}
tree.insert("new_branch_00", 777);
assert_eq!(tree.get("new_branch_00"), Some(&777));
assert_eq!(snapshot.get("new_branch_00"), None);
for (i, prefix) in prefixes.iter().enumerate() {
for j in 0..5 {
let key = format!("{prefix}_{j:02}");
let expected_in_tree = if key == "user_00" {
999
} else {
(i * 100 + j) as i32
};
let expected_in_snapshot = (i * 100 + j) as i32;
assert_eq!(tree.get(&key), Some(&expected_in_tree));
assert_eq!(snapshot.get(&key), Some(&expected_in_snapshot));
}
}
}
#[test]
fn test_into_unversioned_fast_path() {
let mut vtree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
vtree.insert("key1", 10);
vtree.insert("key2", 20);
vtree.insert("key3", 30);
vtree.insert("apple", 100);
vtree.insert("application", 200);
let tree = vtree.into_unversioned();
assert_eq!(tree.get("key1"), Some(&10));
assert_eq!(tree.get("key2"), Some(&20));
assert_eq!(tree.get("key3"), Some(&30));
assert_eq!(tree.get("apple"), Some(&100));
assert_eq!(tree.get("application"), Some(&200));
assert_eq!(tree.get("nonexistent"), None);
}
#[test]
fn test_into_unversioned_slow_path() {
let mut vtree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
vtree.insert("key1", 10);
vtree.insert("key2", 20);
vtree.insert("key3", 30);
let snapshot = vtree.snapshot();
vtree.insert("key4", 40);
vtree.insert("key5", 50);
let tree = vtree.into_unversioned();
assert_eq!(tree.get("key1"), Some(&10));
assert_eq!(tree.get("key2"), Some(&20));
assert_eq!(tree.get("key3"), Some(&30));
assert_eq!(tree.get("key4"), Some(&40));
assert_eq!(tree.get("key5"), Some(&50));
assert_eq!(snapshot.get("key1"), Some(&10));
assert_eq!(snapshot.get("key2"), Some(&20));
assert_eq!(snapshot.get("key3"), Some(&30));
assert_eq!(snapshot.get("key4"), None); assert_eq!(snapshot.get("key5"), None); }
#[test]
fn test_into_unversioned_empty_tree() {
let vtree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
let tree = vtree.into_unversioned();
assert!(tree.is_empty());
assert_eq!(tree.get("anything"), None);
}
#[test]
fn test_into_unversioned_single_element() {
let mut vtree = VersionedAdaptiveRadixTree::<ArrayKey<16>, String>::new();
vtree.insert("only_key", "only_value".to_string());
let tree = vtree.into_unversioned();
assert_eq!(tree.get("only_key"), Some(&"only_value".to_string()));
assert_eq!(tree.get("other"), None);
}
#[test]
fn test_into_unversioned_with_node_growth() {
let mut vtree = VersionedAdaptiveRadixTree::<ArrayKey<16>, usize>::new();
for i in 0..60 {
let key = format!("key_{i:03}");
vtree.insert(key, i);
}
let snapshot = vtree.snapshot();
for i in 60..80 {
let key = format!("key_{i:03}");
vtree.insert(key, i);
}
let tree = vtree.into_unversioned();
for i in 0..80 {
let key = format!("key_{i:03}");
assert_eq!(tree.get(&key), Some(&i), "Missing key {key}");
}
for i in 0..60 {
let key = format!("key_{i:03}");
assert_eq!(snapshot.get(&key), Some(&i));
}
for i in 60..80 {
let key = format!("key_{i:03}");
assert_eq!(snapshot.get(&key), None);
}
}
#[test]
fn test_insert_returns_bool() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
assert!(!tree.insert("key1", 100));
assert_eq!(tree.get("key1"), Some(&100));
assert!(tree.insert("key1", 200));
assert_eq!(tree.get("key1"), Some(&200));
assert!(tree.insert("key1", 300));
assert_eq!(tree.get("key1"), Some(&300));
assert!(!tree.insert("key2", 400));
assert_eq!(tree.get("key2"), Some(&400));
assert_eq!(tree.get("key1"), Some(&300));
assert!(tree.insert("key2", 500));
assert_eq!(tree.get("key2"), Some(&500));
}
#[test]
fn test_insert_and_replace_returns_old_value() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
assert_eq!(tree.insert_and_replace("key1", 100), None);
assert_eq!(tree.get("key1"), Some(&100));
assert_eq!(tree.insert_and_replace("key1", 200), Some(100));
assert_eq!(tree.get("key1"), Some(&200));
assert_eq!(tree.insert_and_replace("key1", 300), Some(200));
assert_eq!(tree.get("key1"), Some(&300));
assert_eq!(tree.insert_and_replace("key2", 400), None);
assert_eq!(tree.get("key2"), Some(&400));
assert_eq!(tree.insert_and_replace("key2", 500), Some(400));
assert_eq!(tree.get("key2"), Some(&500));
}
#[test]
fn test_raw_prefix_keys_are_supported() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<8>, i32>::new();
tree.insert_k(&ArrayKey::new_from_slice(b"d"), 1);
tree.insert_k(&ArrayKey::new_from_slice(b"da"), 2);
assert_eq!(tree.get_k(&ArrayKey::new_from_slice(b"d")), Some(&1));
assert_eq!(tree.get_k(&ArrayKey::new_from_slice(b"da")), Some(&2));
}
#[test]
fn test_remove_longer_key_does_not_delete_leaf_root_prefix() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<8>, i32>::new();
tree.insert_k(&ArrayKey::new_from_slice(b"d"), 1);
assert_eq!(tree.remove_k(&ArrayKey::new_from_slice(b"da")), None);
assert_eq!(tree.get_k(&ArrayKey::new_from_slice(b"d")), Some(&1));
}
#[test]
fn test_insert_with_snapshots() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
assert!(!tree.insert("key1", 100));
assert!(!tree.insert("key2", 200));
let snapshot = tree.snapshot();
assert!(tree.insert("key1", 300));
assert_eq!(tree.get("key1"), Some(&300));
assert_eq!(snapshot.get("key1"), Some(&100));
assert!(!tree.insert("key3", 400));
assert_eq!(tree.get("key3"), Some(&400));
assert_eq!(snapshot.get("key3"), None);
assert_eq!(tree.insert_and_replace("key1", 500), Some(300));
assert_eq!(tree.get("key1"), Some(&500));
assert_eq!(snapshot.get("key1"), Some(&100));
}
#[test]
fn test_prefix_keys_with_snapshots_preserve_isolation() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<8>, i32>::new();
tree.insert_k(&ArrayKey::new_from_slice(b"d"), 1);
let snapshot = tree.snapshot();
tree.insert_k(&ArrayKey::new_from_slice(b"da"), 2);
assert_eq!(tree.get_k(&ArrayKey::new_from_slice(b"d")), Some(&1));
assert_eq!(tree.get_k(&ArrayKey::new_from_slice(b"da")), Some(&2));
assert_eq!(snapshot.get_k(&ArrayKey::new_from_slice(b"d")), Some(&1));
assert_eq!(snapshot.get_k(&ArrayKey::new_from_slice(b"da")), None);
}
#[test]
fn traversal_empty_tree_visits_nothing() {
let tree = VersionedAdaptiveRadixTree::<ArrayKey<8>, i32>::new();
let mut count = 0;
tree.for_each_view(|_, _| count += 1);
assert_eq!(count, 0);
assert_eq!(tree.iter().count(), 0);
assert_eq!(tree.values_iter().count(), 0);
}
#[test]
fn traversal_single_root_leaf() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<8>, i32>::new();
let key = ArrayKey::new_from_slice(b"only");
tree.insert_k(&key, 7);
let iter_items: Vec<_> = tree.iter().map(|(key, value)| (key, *value)).collect();
assert_eq!(iter_items, vec![(key, 7)]);
let mut views = Vec::new();
tree.for_each_view(|key_view, value| {
views.push((key_view.to_vec(), *value));
});
assert_eq!(views, vec![(b"only".to_vec(), 7)]);
assert_eq!(tree.values_iter().copied().collect::<Vec<_>>(), vec![7]);
}
#[test]
fn traversal_handles_key_that_is_prefix_of_another() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<8>, i32>::new();
tree.insert_k(&ArrayKey::new_from_slice(b"d"), 1);
tree.insert_k(&ArrayKey::new_from_slice(b"da"), 2);
let prefix = ArrayKey::new_from_slice(b"d");
let values: Vec<_> = tree
.prefix_iter_k(&prefix)
.map(|(key, value)| (key.as_ref().to_vec(), *value))
.collect();
assert_eq!(values, vec![(b"d".to_vec(), 1), (b"da".to_vec(), 2)]);
let mut views = Vec::new();
tree.prefix_for_each_view_k(&prefix, |key_view, value| {
views.push((key_view.to_vec(), *value));
});
assert_eq!(views, vec![(b"d".to_vec(), 1), (b"da".to_vec(), 2)]);
}
#[test]
fn prefix_match_iter_returns_saved_prefixes() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
for (idx, key) in [
b"".as_slice(),
b"a".as_slice(),
b"alpha".as_slice(),
b"alphabet".as_slice(),
b"alphabetical".as_slice(),
b"apple".as_slice(),
]
.iter()
.enumerate()
{
tree.insert_k(&ArrayKey::new_from_slice(key), idx as i32);
}
let got: Vec<Vec<u8>> = tree
.prefix_match_iter(ArrayKey::new_from_slice(b"alphabet"))
.map(|(key, _)| key.as_ref().to_vec())
.collect();
assert_eq!(
got,
vec![
b"".to_vec(),
b"a".to_vec(),
b"alpha".to_vec(),
b"alphabet".to_vec(),
]
);
}
#[test]
fn prefix_match_for_each_returns_saved_prefixes() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
for (idx, key) in [
b"".as_slice(),
b"a".as_slice(),
b"alpha".as_slice(),
b"alphabet".as_slice(),
b"alphabetical".as_slice(),
b"apple".as_slice(),
]
.iter()
.enumerate()
{
tree.insert_k(&ArrayKey::new_from_slice(key), idx as i32);
}
let probe = ArrayKey::new_from_slice(b"alphabet");
let mut got = Vec::new();
tree.prefix_match_for_each_k(&probe, |key, value| {
got.push((key.to_vec(), *value));
});
assert_eq!(
got,
vec![
(b"".to_vec(), 0),
(b"a".to_vec(), 1),
(b"alpha".to_vec(), 2),
(b"alphabet".to_vec(), 3),
]
);
}
#[test]
fn prefix_traversal_matches_inside_compressed_prefix() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
tree.insert_k(&ArrayKey::new_from_slice(b"abcdef"), 1);
tree.insert_k(&ArrayKey::new_from_slice(b"abcxyz"), 2);
let prefix = ArrayKey::new_from_slice(b"abc");
let values: Vec<_> = tree
.prefix_iter_k(&prefix)
.map(|(key, value)| (key.as_ref().to_vec(), *value))
.collect();
assert_eq!(
values,
vec![(b"abcdef".to_vec(), 1), (b"abcxyz".to_vec(), 2)]
);
}
#[test]
fn prefix_traversal_no_match_visits_nothing() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
tree.insert_k(&ArrayKey::new_from_slice(b"abcdef"), 1);
let prefix = ArrayKey::new_from_slice(b"abz");
let mut count = 0;
tree.prefix_for_each_view_k(&prefix, |_, _| count += 1);
assert_eq!(count, 0);
assert_eq!(tree.prefix_iter_k(&prefix).count(), 0);
}
#[test]
fn versioned_traversal_order_matches_unversioned_tree() {
let mut versioned = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
let mut unversioned = crate::tree::AdaptiveRadixTree::<ArrayKey<16>, i32>::new();
let keys = [
b"banana".as_slice(),
b"app".as_slice(),
b"apple".as_slice(),
b"band".as_slice(),
b"can".as_slice(),
b"bandana".as_slice(),
];
for (idx, key) in keys.iter().enumerate() {
let key = ArrayKey::new_from_slice(key);
versioned.insert_k(&key, idx as i32);
unversioned.insert_k(&key, idx as i32);
}
let versioned_items: Vec<_> = versioned
.iter()
.map(|(key, value)| (key.as_ref().to_vec(), *value))
.collect();
let unversioned_items: Vec<_> = unversioned
.iter()
.map(|(key, value)| (key.as_ref().to_vec(), *value))
.collect();
assert_eq!(versioned_items, unversioned_items);
let start = ArrayKey::new_from_slice(b"banana");
let end = ArrayKey::new_from_slice(b"can");
let versioned_range: Vec<_> = versioned
.range(start..end)
.map(|(key, value)| (key.as_ref().to_vec(), *value))
.collect();
assert_eq!(
versioned_range,
vec![
(b"banana".to_vec(), 0),
(b"band".to_vec(), 3),
(b"bandana".to_vec(), 5)
]
);
}
#[test]
fn versioned_traversal_preserves_snapshot_isolation() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
for (key, value) in [
(b"a".as_slice(), 1),
(b"b".as_slice(), 2),
(b"c".as_slice(), 3),
] {
tree.insert_k(&ArrayKey::new_from_slice(key), value);
}
let snapshot = tree.snapshot();
tree.insert_k(&ArrayKey::new_from_slice(b"b"), 20);
tree.insert_k(&ArrayKey::new_from_slice(b"d"), 4);
assert_eq!(tree.remove_k(&ArrayKey::new_from_slice(b"a")), Some(1));
let snapshot_items: Vec<_> = snapshot
.iter()
.map(|(key, value)| (key.as_ref().to_vec(), *value))
.collect();
assert_eq!(
snapshot_items,
vec![(b"a".to_vec(), 1), (b"b".to_vec(), 2), (b"c".to_vec(), 3)]
);
let tree_items: Vec<_> = tree
.iter()
.map(|(key, value)| (key.as_ref().to_vec(), *value))
.collect();
assert_eq!(
tree_items,
vec![(b"b".to_vec(), 20), (b"c".to_vec(), 3), (b"d".to_vec(), 4)]
);
}
#[test]
fn traversal_covers_wide_node_shapes() {
for size in [4usize, 16, 48, 256] {
let mut tree = VersionedAdaptiveRadixTree::<DenseSequentialKey, usize>::new();
for idx in 0..size {
tree.insert_k(&dense_sequential_key(idx as u32), idx);
}
let values: Vec<_> = tree.values_iter().copied().collect();
assert_eq!(values, (0..size).collect::<Vec<_>>());
let mut viewed = Vec::new();
tree.for_each_view(|key_view, value| {
viewed.push((key_view.to_vec(), *value));
});
assert_eq!(viewed.len(), size);
assert_eq!(viewed[0].1, 0);
assert_eq!(viewed[size - 1].1, size - 1);
}
}
fn cache_key(obj: u64, symbol: u32) -> ArrayKey<16> {
let mut bytes = [0; 16];
bytes[..8].copy_from_slice(&obj.to_be_bytes());
bytes[8..12].copy_from_slice(&symbol.to_be_bytes());
ArrayKey::new_from_slice(&bytes[..12])
}
fn obj_prefix(obj: u64) -> ArrayKey<16> {
let mut bytes = [0; 16];
bytes[..8].copy_from_slice(&obj.to_be_bytes());
ArrayKey::new_from_slice(&bytes[..8])
}
#[test]
fn prefix_for_each_view_supports_moor_shaped_cache_keys() {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, usize>::new();
for symbol in 0..5 {
tree.insert_k(&cache_key(10, symbol), symbol as usize);
}
for symbol in 0..3 {
tree.insert_k(&cache_key(11, symbol), 100 + symbol as usize);
}
let prefix = obj_prefix(10);
let mut values = Vec::new();
tree.prefix_for_each_view_k(&prefix, |key_view, value| {
assert!(key_view.to_vec().starts_with(prefix.as_ref()));
values.push(*value);
});
assert_eq!(values, vec![0, 1, 2, 3, 4]);
}
#[test]
fn intersect_with_returns_common_keys_and_values() {
let mut left = VersionedAdaptiveRadixTree::<ArrayKey<32>, i32>::new();
let mut right = VersionedAdaptiveRadixTree::<ArrayKey<32>, i32>::new();
for (key, value) in [
(b"a".as_slice(), 1),
(b"ab".as_slice(), 2),
(b"abc".as_slice(), 3),
(b"abd".as_slice(), 4),
(b"bzz".as_slice(), 5),
(b"cat".as_slice(), 6),
] {
left.insert_k(&ArrayKey::new_from_slice(key), value);
}
for (key, value) in [
(b"ab".as_slice(), 20),
(b"abc".as_slice(), 30),
(b"bzz".as_slice(), 50),
(b"dog".as_slice(), 70),
] {
right.insert_k(&ArrayKey::new_from_slice(key), value);
}
let mut seen = Vec::new();
left.intersect_with(&right, |key, left_value, right_value| {
seen.push((key.as_ref().to_vec(), *left_value, *right_value));
});
seen.sort();
assert_eq!(
seen,
vec![
(b"ab".to_vec(), 2, 20),
(b"abc".to_vec(), 3, 30),
(b"bzz".to_vec(), 5, 50),
]
);
}
#[test]
fn intersect_lending_with_returns_common_keys_and_values() {
let mut left = VersionedAdaptiveRadixTree::<ArrayKey<32>, i32>::new();
let mut right = VersionedAdaptiveRadixTree::<ArrayKey<32>, i32>::new();
for (key, value) in [
(b"a".as_slice(), 1),
(b"ab".as_slice(), 2),
(b"abc".as_slice(), 3),
(b"abd".as_slice(), 4),
(b"bzz".as_slice(), 5),
] {
left.insert_k(&ArrayKey::new_from_slice(key), value);
}
for (key, value) in [
(b"ab".as_slice(), 20),
(b"abc".as_slice(), 30),
(b"bzz".as_slice(), 50),
(b"dog".as_slice(), 70),
] {
right.insert_k(&ArrayKey::new_from_slice(key), value);
}
let mut seen = Vec::new();
left.intersect_lending_with(&right, |key, left_value, right_value| {
seen.push((key.to_vec(), *left_value, *right_value));
});
seen.sort();
assert_eq!(
seen,
vec![
(b"ab".to_vec(), 2, 20),
(b"abc".to_vec(), 3, 30),
(b"bzz".to_vec(), 5, 50),
]
);
}
#[test]
fn intersect_values_with_and_count() {
let mut left = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
let mut right = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
for (key, value) in [
(b"aa".as_slice(), 1),
(b"ab".as_slice(), 2),
(b"ac".as_slice(), 3),
] {
left.insert_k(&ArrayKey::new_from_slice(key), value);
}
for (key, value) in [
(b"ab".as_slice(), 20),
(b"ac".as_slice(), 30),
(b"zz".as_slice(), 40),
] {
right.insert_k(&ArrayKey::new_from_slice(key), value);
}
let mut pairs = Vec::new();
left.intersect_values_with(&right, |left_value, right_value| {
pairs.push((*left_value, *right_value));
});
pairs.sort_unstable();
assert_eq!(pairs, vec![(2, 20), (3, 30)]);
assert_eq!(left.intersect_count(&right), 2);
}
#[test]
fn intersect_with_empty_tree_visits_nothing() {
let mut left = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
let right = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
left.insert_k(&ArrayKey::new_from_slice(b"a"), 1);
left.insert_k(&ArrayKey::new_from_slice(b"b"), 2);
let mut count = 0usize;
left.intersect_with(&right, |_key, _left_value, _right_value| {
count += 1;
});
assert_eq!(count, 0);
assert_eq!(left.intersect_count(&right), 0);
}
#[test]
fn intersect_uses_snapshot_state_after_cow_mutations() {
let mut left = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
let mut right = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
for (key, value) in [
(b"aa".as_slice(), 1),
(b"ab".as_slice(), 2),
(b"ac".as_slice(), 3),
] {
left.insert_k(&ArrayKey::new_from_slice(key), value);
}
for (key, value) in [
(b"ab".as_slice(), 20),
(b"ac".as_slice(), 30),
(b"ad".as_slice(), 40),
] {
right.insert_k(&ArrayKey::new_from_slice(key), value);
}
let left_snapshot = left.snapshot();
let right_snapshot = right.snapshot();
left.insert_k(&ArrayKey::new_from_slice(b"ab"), 200);
left.remove_k(&ArrayKey::new_from_slice(b"ac"));
left.insert_k(&ArrayKey::new_from_slice(b"ad"), 4);
right.insert_k(&ArrayKey::new_from_slice(b"ad"), 400);
right.remove_k(&ArrayKey::new_from_slice(b"ab"));
let mut snapshot_pairs = Vec::new();
left_snapshot.intersect_with(&right_snapshot, |key, left_value, right_value| {
snapshot_pairs.push((key.as_ref().to_vec(), *left_value, *right_value));
});
snapshot_pairs.sort();
assert_eq!(
snapshot_pairs,
vec![(b"ab".to_vec(), 2, 20), (b"ac".to_vec(), 3, 30)]
);
let mut live_pairs = Vec::new();
left.intersect_with(&right, |key, left_value, right_value| {
live_pairs.push((key.as_ref().to_vec(), *left_value, *right_value));
});
live_pairs.sort();
assert_eq!(live_pairs, vec![(b"ad".to_vec(), 4, 400)]);
}
#[test]
fn dense_sequential_snapshot_isolation_survives_wide_node_cow() {
for size in [48usize, 256] {
let keys: Vec<_> = (0..size).map(|i| dense_sequential_key(i as u32)).collect();
let mut tree = VersionedAdaptiveRadixTree::<DenseSequentialKey, usize>::new();
for (i, key) in keys.iter().enumerate() {
assert!(!tree.insert_k(key, i));
}
let snapshot = tree.snapshot();
for (i, key) in keys.iter().enumerate() {
assert!(tree.insert_k(key, i + 10_000));
}
for (i, key) in keys.iter().enumerate() {
assert_eq!(snapshot.get_k(key), Some(&i));
assert_eq!(tree.get_k(key), Some(&(i + 10_000)));
}
let scratch_key = dense_sequential_key((size + 1024) as u32);
assert!(!tree.insert_k(&scratch_key, 77));
assert_eq!(tree.remove_k(&scratch_key), Some(77));
assert_eq!(snapshot.get_k(&scratch_key), None);
assert_eq!(tree.get_k(&scratch_key), None);
for (i, key) in keys.iter().enumerate() {
assert_eq!(tree.remove_k(key), Some(i + 10_000));
assert!(!tree.insert_k(key, i + 20_000));
}
for (i, key) in keys.iter().enumerate() {
assert_eq!(snapshot.get_k(key), Some(&i));
assert_eq!(tree.get_k(key), Some(&(i + 20_000)));
}
}
}
#[test]
fn dense_sequential_multiple_snapshots_mutate_independently() {
let size = 4096usize;
let keys: Vec<_> = (0..size).map(|i| dense_sequential_key(i as u32)).collect();
let mut tree = VersionedAdaptiveRadixTree::<DenseSequentialKey, usize>::new();
for (i, key) in keys.iter().enumerate() {
assert!(!tree.insert_k(key, i));
}
let mut snapshot_a = tree.snapshot();
let snapshot_b = tree.snapshot();
for (i, key) in keys.iter().enumerate() {
assert!(tree.insert_k(key, i + 10_000));
}
for (i, key) in keys.iter().enumerate().step_by(3) {
assert_eq!(snapshot_a.remove_k(key), Some(i));
assert!(!snapshot_a.insert_k(key, i + 20_000));
}
for i in size..(size + 128) {
let key = dense_sequential_key(i as u32);
assert!(!snapshot_a.insert_k(&key, i + 30_000));
assert_eq!(tree.get_k(&key), None);
assert_eq!(snapshot_b.get_k(&key), None);
}
for (i, key) in keys.iter().enumerate() {
assert_eq!(tree.get_k(key), Some(&(i + 10_000)));
assert_eq!(snapshot_b.get_k(key), Some(&i));
let expected_a = if i % 3 == 0 { i + 20_000 } else { i };
assert_eq!(snapshot_a.get_k(key), Some(&expected_a));
}
}
#[test]
fn test_into_unversioned_preserves_tree_structure() {
let mut vtree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
let mut regular_tree = crate::tree::AdaptiveRadixTree::<ArrayKey<16>, i32>::new();
let test_data = vec![
("apple", 1),
("application", 2),
("app", 3),
("banana", 4),
("band", 5),
("bandana", 6),
("can", 7),
("cannot", 8),
];
for (key, value) in &test_data {
vtree.insert(*key, *value);
regular_tree.insert(*key, *value);
}
let converted_tree = vtree.into_unversioned();
for (key, expected_value) in &test_data {
assert_eq!(converted_tree.get(*key), Some(expected_value));
assert_eq!(regular_tree.get(*key), Some(expected_value));
assert_eq!(converted_tree.get(*key), regular_tree.get(*key));
}
let non_existent = ["xyz", "apple_pie", "ban", "candidate"];
for key in &non_existent {
assert_eq!(converted_tree.get(*key), None);
assert_eq!(regular_tree.get(*key), None);
assert_eq!(converted_tree.get(*key), regular_tree.get(*key));
}
}
#[test]
fn test_into_unversioned_memory_efficiency() {
let mut vtree = VersionedAdaptiveRadixTree::<ArrayKey<16>, Box<i32>>::new();
vtree.insert("key1", Box::new(42));
vtree.insert("key2", Box::new(84));
let tree = vtree.into_unversioned();
assert_eq!(**tree.get("key1").unwrap(), 42);
assert_eq!(**tree.get("key2").unwrap(), 84);
}
}
#[cfg(test)]
mod shuttle_tests {
use super::*;
use crate::keys::array_key::ArrayKey;
use shuttle::{Config, Runner, sync::Arc as ShuttleArc, thread};
#[test]
fn shuttle_concurrent_snapshots() {
let runner = Runner::new(
shuttle::scheduler::DfsScheduler::new(Some(1000), false),
Config::new(),
);
runner.run(|| {
let tree = ShuttleArc::new(std::sync::Mutex::new(VersionedAdaptiveRadixTree::<
ArrayKey<16>,
i32,
>::new()));
{
let mut t = tree.lock().unwrap();
for i in 0..10 {
t.insert(i, i * 10);
}
}
let snapshot1 = {
let t = tree.lock().unwrap();
t.snapshot()
};
let snapshot2 = {
let t = tree.lock().unwrap();
t.snapshot()
};
let tree1 = ShuttleArc::clone(&tree);
let handle1 = thread::spawn(move || {
for i in 0..10 {
assert_eq!(snapshot1.get(i), Some(&(i * 10)));
}
snapshot1
});
let handle2 = thread::spawn(move || {
for i in 0..10 {
assert_eq!(snapshot2.get(i), Some(&(i * 10)));
}
snapshot2
});
let handle3 = thread::spawn(move || {
let mut t = tree1.lock().unwrap();
t.insert(100, 1000);
assert_eq!(t.get(100), Some(&1000));
});
let snapshot1 = handle1.join().unwrap();
let snapshot2 = handle2.join().unwrap();
handle3.join().unwrap();
assert_eq!(snapshot1.get(100), None);
assert_eq!(snapshot2.get(100), None);
let t = tree.lock().unwrap();
assert_eq!(t.get(100), Some(&1000));
});
}
#[test]
fn shuttle_snapshot_sharing_across_threads() {
let runner = Runner::new(
shuttle::scheduler::DfsScheduler::new(Some(1000), false),
Config::new(),
);
runner.run(|| {
let mut tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
for i in 0..5 {
tree.insert(i, i * 2);
}
let snapshot = ShuttleArc::new(tree.snapshot());
let snapshot1 = ShuttleArc::clone(&snapshot);
let snapshot2 = ShuttleArc::clone(&snapshot);
let snapshot3 = ShuttleArc::clone(&snapshot);
let handle1 = thread::spawn(move || {
let mut results = Vec::new();
for i in 0..5 {
if let Some(val) = snapshot1.get(i) {
results.push(*val);
}
}
results
});
let handle2 = thread::spawn(move || {
let mut results = Vec::new();
for i in 0..5 {
if let Some(val) = snapshot2.get(i) {
results.push(*val);
}
}
results
});
let handle3 = thread::spawn(move || {
let mut results = Vec::new();
for i in 0..5 {
if let Some(val) = snapshot3.get(i) {
results.push(*val);
}
}
results
});
let results1 = handle1.join().unwrap();
let results2 = handle2.join().unwrap();
let results3 = handle3.join().unwrap();
let expected: Vec<i32> = (0..5).map(|i| i * 2).collect();
assert_eq!(results1, expected);
assert_eq!(results2, expected);
assert_eq!(results3, expected);
});
}
#[test]
fn shuttle_concurrent_snapshot_mutations() {
let runner = Runner::new(
shuttle::scheduler::DfsScheduler::new(Some(1000), false),
Config::new(),
);
runner.run(|| {
let mut base_tree = VersionedAdaptiveRadixTree::<ArrayKey<16>, i32>::new();
for i in 0..3 {
base_tree.insert(i, i);
}
let snapshot1 = ShuttleArc::new(std::sync::Mutex::new(base_tree.snapshot()));
let snapshot2 = ShuttleArc::new(std::sync::Mutex::new(base_tree.snapshot()));
let s1 = ShuttleArc::clone(&snapshot1);
let s2 = ShuttleArc::clone(&snapshot2);
let handle1 = thread::spawn(move || {
let mut snap = s1.lock().unwrap();
snap.insert(10, 100);
snap.insert(11, 110);
assert_eq!(snap.get(10), Some(&100));
assert_eq!(snap.get(11), Some(&110));
for i in 0..3 {
assert_eq!(snap.get(i), Some(&i));
}
});
let handle2 = thread::spawn(move || {
let mut snap = s2.lock().unwrap();
snap.insert(20, 200);
snap.insert(21, 210);
assert_eq!(snap.get(20), Some(&200));
assert_eq!(snap.get(21), Some(&210));
for i in 0..3 {
assert_eq!(snap.get(i), Some(&i));
}
});
handle1.join().unwrap();
handle2.join().unwrap();
{
let snap1 = snapshot1.lock().unwrap();
assert_eq!(snap1.get(10), Some(&100));
assert_eq!(snap1.get(11), Some(&110));
assert_eq!(snap1.get(20), None); assert_eq!(snap1.get(21), None);
}
{
let snap2 = snapshot2.lock().unwrap();
assert_eq!(snap2.get(20), Some(&200));
assert_eq!(snap2.get(21), Some(&210));
assert_eq!(snap2.get(10), None); assert_eq!(snap2.get(11), None);
}
});
}
#[test]
fn shuttle_many_readers_one_writer() {
let runner = Runner::new(
shuttle::scheduler::DfsScheduler::new(Some(1000), false),
Config::new(),
);
runner.run(|| {
let tree = ShuttleArc::new(std::sync::Mutex::new(VersionedAdaptiveRadixTree::<
ArrayKey<16>,
i32,
>::new()));
{
let mut t = tree.lock().unwrap();
for i in 0..10 {
t.insert(i, i * 3);
}
}
let snapshot = {
let t = tree.lock().unwrap();
ShuttleArc::new(t.snapshot())
};
let tree_for_writer = ShuttleArc::clone(&tree);
let mut reader_handles = Vec::new();
for reader_id in 0..3 {
let snap = ShuttleArc::clone(&snapshot);
let handle = thread::spawn(move || {
let mut sum = 0;
for i in 0..10 {
if let Some(val) = snap.get(i) {
sum += val;
}
}
(reader_id, sum)
});
reader_handles.push(handle);
}
let writer_handle = thread::spawn(move || {
let mut tree = tree_for_writer.lock().unwrap();
for i in 100..105 {
tree.insert(i, i * 5);
}
let mut writer_sum = 0;
for i in 100..105 {
if let Some(val) = tree.get(i) {
writer_sum += val;
}
}
writer_sum
});
let expected_reader_sum = (0..10).map(|i| i * 3).sum::<i32>();
for handle in reader_handles {
let (reader_id, sum) = handle.join().unwrap();
assert_eq!(sum, expected_reader_sum, "Reader {reader_id} got wrong sum");
}
let writer_sum = writer_handle.join().unwrap();
let expected_writer_sum = (100..105).map(|i| i * 5).sum::<i32>();
assert_eq!(writer_sum, expected_writer_sum);
for i in 100..105 {
assert_eq!(snapshot.get(i), None);
}
});
}
#[test]
fn shuttle_snapshot_drop_safety() {
let runner = Runner::new(
shuttle::scheduler::DfsScheduler::new(Some(1000), false),
Config::new(),
);
runner.run(|| {
let tree = ShuttleArc::new(std::sync::Mutex::new(VersionedAdaptiveRadixTree::<
ArrayKey<16>,
i32,
>::new()));
{
let mut t = tree.lock().unwrap();
for i in 0..5 {
t.insert(i, i * 7);
}
}
let tree1 = ShuttleArc::clone(&tree);
let tree2 = ShuttleArc::clone(&tree);
let handle1 = thread::spawn(move || {
let snapshot = {
let t = tree1.lock().unwrap();
t.snapshot()
};
let mut sum = 0;
for i in 0..5 {
if let Some(val) = snapshot.get(i) {
sum += val;
}
}
sum
});
let handle2 = thread::spawn(move || {
let snapshot = {
let t = tree2.lock().unwrap();
t.snapshot()
};
let mut count = 0;
for i in 0..5 {
if snapshot.get(i).is_some() {
count += 1;
}
}
count
});
let sum = handle1.join().unwrap();
let count = handle2.join().unwrap();
assert_eq!(sum, (0..5).map(|i| i * 7).sum::<i32>());
assert_eq!(count, 5);
let t = tree.lock().unwrap();
for i in 0..5 {
assert_eq!(t.get(i), Some(&(i * 7)));
}
});
}
}