use ::core::borrow::Borrow;
use ::core::fmt::Debug;
use ::core::iter::FusedIterator;
use ::core::marker::PhantomData;
use ::core::ops::{Bound, RangeBounds};
use ::core::sync::atomic::{AtomicPtr, AtomicU64, Ordering};
use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use alloc::sync::Arc;
use alloc::vec;
use alloc::vec::Vec;
use parking_lot::{
ArcRwLockReadGuard, ArcRwLockWriteGuard, Mutex, MutexGuard, RawRwLock, RwLock, RwLockReadGuard, RwLockWriteGuard,
};
use crate::cdc::change::ChangeEvent;
use crate::concurrent::operation::*;
use crate::core::constants::DEFAULT_INNER_SIZE;
use crate::core::node::*;
use super::r#ref::Ref;
#[inline]
fn yield_now() {
#[cfg(feature = "std")]
std::thread::yield_now();
#[cfg(not(feature = "std"))]
::core::hint::spin_loop();
}
const ROOT_PUBLICATION_SPIN_LIMIT: usize = 16;
const STABLE_READ_BLOCKING_FALLBACK_AFTER: usize = 2;
const PUBLICATION_BACKLOG_DRAIN_THRESHOLD: usize = 64;
type NodeIndex<T, Node> = BTreeMap<T, Arc<RwLock<Node>>>;
const PUBLISHED_ROUTES_PER_CHUNK: usize = 128;
const PUBLISHED_REBUILD_ROUTES_PER_CHUNK: usize = PUBLISHED_ROUTES_PER_CHUNK * 2 / 3;
const PUBLISHED_ROUTE_MERGE_THRESHOLD: usize = PUBLISHED_ROUTES_PER_CHUNK * 3 / 4;
struct PublishedChunk<T, Node> {
entries: Vec<(T, Arc<RwLock<Node>>)>,
}
impl<T: Clone, Node> Clone for PublishedChunk<T, Node> {
fn clone(&self) -> Self {
Self {
entries: self.entries.clone(),
}
}
}
struct PublishedNodeIndex<T, Node> {
chunks: Vec<Arc<PublishedChunk<T, Node>>>,
len: usize,
}
impl<T, Node> Clone for PublishedNodeIndex<T, Node> {
fn clone(&self) -> Self {
Self {
chunks: self.chunks.clone(),
len: self.len,
}
}
}
impl<T, Node> PublishedNodeIndex<T, Node>
where
T: Ord + Clone,
{
fn from_canonical(index: &NodeIndex<T, Node>) -> Self {
let mut chunks = Vec::with_capacity(index.len().div_ceil(PUBLISHED_REBUILD_ROUTES_PER_CHUNK));
let mut entries = Vec::with_capacity(PUBLISHED_REBUILD_ROUTES_PER_CHUNK);
for (key, node) in index {
entries.push((key.clone(), node.clone()));
if entries.len() == PUBLISHED_REBUILD_ROUTES_PER_CHUNK {
chunks.push(Arc::new(PublishedChunk { entries }));
entries = Vec::with_capacity(PUBLISHED_REBUILD_ROUTES_PER_CHUNK);
}
}
if !entries.is_empty() {
chunks.push(Arc::new(PublishedChunk { entries }));
}
Self {
chunks,
len: index.len(),
}
}
fn iter(&self) -> impl Iterator<Item = (&T, &Arc<RwLock<Node>>)> {
self.chunks
.iter()
.flat_map(|chunk| chunk.entries.iter().map(|(key, node)| (key, node)))
}
fn first_key_value(&self) -> Option<(&T, &Arc<RwLock<Node>>)> {
self.chunks.first()?.entries.first().map(|(key, node)| (key, node))
}
fn last_key_value(&self) -> Option<(&T, &Arc<RwLock<Node>>)> {
self.chunks.last()?.entries.last().map(|(key, node)| (key, node))
}
fn chunk_for<Q>(&self, key: &Q) -> usize
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
self.chunks.partition_point(|chunk| {
let max = &chunk.entries.last().expect("published chunks are non-empty").0;
<T as Borrow<Q>>::borrow(max) < key
})
}
fn first_for_bound<Q>(&self, bound: Bound<&Q>) -> Option<(&T, &Arc<RwLock<Node>>)>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
let key = match bound {
Bound::Included(key) | Bound::Excluded(key) => key,
Bound::Unbounded => return self.first_key_value(),
};
let mut chunk_index = self.chunk_for(key);
while let Some(chunk) = self.chunks.get(chunk_index) {
let entry_index = chunk.entries.partition_point(|(candidate, _)| match bound {
Bound::Included(_) => <T as Borrow<Q>>::borrow(candidate) < key,
Bound::Excluded(_) => <T as Borrow<Q>>::borrow(candidate) <= key,
Bound::Unbounded => false,
});
if let Some((found, node)) = chunk.entries.get(entry_index) {
return Some((found, node));
}
chunk_index += 1;
}
None
}
fn insert(&mut self, key: T, node: Arc<RwLock<Node>>) -> Option<Arc<RwLock<Node>>> {
if self.chunks.is_empty() {
self.chunks.push(Arc::new(PublishedChunk {
entries: vec![(key, node)],
}));
self.len = 1;
return None;
}
let mut chunk_index = self.chunk_for(&key);
if chunk_index == self.chunks.len() {
chunk_index -= 1;
}
let chunk = Arc::make_mut(&mut self.chunks[chunk_index]);
match chunk.entries.binary_search_by(|(candidate, _)| candidate.cmp(&key)) {
Ok(index) => Some(::core::mem::replace(&mut chunk.entries[index].1, node)),
Err(index) => {
chunk.entries.insert(index, (key, node));
self.len += 1;
if chunk.entries.len() > PUBLISHED_ROUTES_PER_CHUNK {
let right = chunk.entries.split_off(chunk.entries.len() / 2);
self.chunks
.insert(chunk_index + 1, Arc::new(PublishedChunk { entries: right }));
}
None
}
}
}
fn remove<Q>(&mut self, key: &Q) -> Option<Arc<RwLock<Node>>>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
let chunk_index = self.chunk_for(key);
let entry_index = self
.chunks
.get(chunk_index)?
.entries
.binary_search_by(|(candidate, _)| <T as Borrow<Q>>::borrow(candidate).cmp(key))
.ok()?;
let chunk = Arc::make_mut(&mut self.chunks[chunk_index]);
let (_, removed) = chunk.entries.remove(entry_index);
self.len -= 1;
if chunk.entries.is_empty() {
self.chunks.remove(chunk_index);
} else if chunk_index > 0
&& self.chunks[chunk_index - 1].entries.len() + self.chunks[chunk_index].entries.len()
<= PUBLISHED_ROUTE_MERGE_THRESHOLD
{
let right = self.chunks.remove(chunk_index);
Arc::make_mut(&mut self.chunks[chunk_index - 1])
.entries
.extend(right.entries.iter().cloned());
} else if chunk_index + 1 < self.chunks.len()
&& self.chunks[chunk_index].entries.len() + self.chunks[chunk_index + 1].entries.len()
<= PUBLISHED_ROUTE_MERGE_THRESHOLD
{
let right = self.chunks.remove(chunk_index + 1);
Arc::make_mut(&mut self.chunks[chunk_index])
.entries
.extend(right.entries.iter().cloned());
}
Some(removed)
}
}
#[inline]
fn node_identity<Node>(node: &Arc<RwLock<Node>>) -> usize {
Arc::as_ptr(node) as usize
}
struct RetiredIndex<T, Node>(*mut PublishedNodeIndex<T, Node>);
unsafe impl<T: Send, Node: Send> Send for RetiredIndex<T, Node> {}
impl<T, Node> Drop for RetiredIndex<T, Node> {
fn drop(&mut self) {
unsafe { drop(Box::from_raw(self.0)) }
}
}
struct PublishedIndex<T, Node> {
current: AtomicPtr<PublishedNodeIndex<T, Node>>,
domain: ps_reclaim::Domain,
}
impl<T, Node> PublishedIndex<T, Node> {
fn new() -> Self {
Self {
current: AtomicPtr::new(Box::into_raw(Box::new(PublishedNodeIndex {
chunks: Vec::new(),
len: 0,
}))),
domain: ps_reclaim::Domain::new(),
}
}
}
impl<T, Node> PublishedIndex<T, Node>
where
T: Ord + Clone + Send + 'static,
Node: Send + 'static,
{
fn snapshot(&self) -> PublishedNodeIndex<T, Node> {
let current = self.current.load(Ordering::Acquire);
unsafe { (&*current).clone() }
}
fn replace(&self, replacement: PublishedNodeIndex<T, Node>) -> RetiredIndex<T, Node> {
let replacement = Box::into_raw(Box::new(replacement));
let retired = self.current.swap(replacement, Ordering::AcqRel);
RetiredIndex(retired)
}
fn retire(&self, retired: RetiredIndex<T, Node>) {
self.domain.retire(move || drop(retired));
}
fn advance(&self) {
if self.domain.pending() >= PUBLICATION_BACKLOG_DRAIN_THRESHOLD {
self.domain.advance();
}
}
}
impl<T, Node> Drop for PublishedIndex<T, Node> {
fn drop(&mut self) {
let current = *self.current.get_mut();
unsafe { drop(Box::from_raw(current)) }
}
}
pub(crate) struct Topology<T, Node> {
index: RwLock<NodeIndex<T, Node>>,
published_keys: Mutex<BTreeMap<usize, T>>,
published: PublishedIndex<T, Node>,
generation: AtomicU64,
}
impl<T, Node> Debug for Topology<T, Node> {
fn fmt(&self, formatter: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
formatter
.debug_struct("Topology")
.field("nodes", &self.index.read().len())
.field("generation", &self.generation.load(Ordering::Relaxed))
.finish()
}
}
impl<T, Node> Topology<T, Node> {
fn new() -> Self {
Self {
index: RwLock::new(BTreeMap::new()),
published_keys: Mutex::new(BTreeMap::new()),
published: PublishedIndex::new(),
generation: AtomicU64::new(0),
}
}
#[inline]
pub(crate) fn read(&self) -> RwLockReadGuard<'_, NodeIndex<T, Node>> {
self.index.read()
}
}
impl<T, Node> Topology<T, Node>
where
T: Ord + Clone + Send + 'static,
Node: Send + 'static,
{
#[inline]
fn write(&self) -> TopologyWriteGuard<'_, T, Node> {
let index = self.index.write();
self.generation.fetch_add(1, Ordering::AcqRel);
TopologyWriteGuard {
topology: self,
index: Some(index),
published: None,
published_keys: None,
publish: true,
dirty: false,
}
}
#[inline]
fn write_rekey(&self) -> TopologyWriteGuard<'_, T, Node> {
let index = self.index.write();
TopologyWriteGuard {
topology: self,
index: Some(index),
published: None,
published_keys: None,
publish: false,
dirty: false,
}
}
#[inline]
fn try_write(&self) -> Option<TopologyWriteGuard<'_, T, Node>> {
let index = self.index.try_write()?;
self.generation.fetch_add(1, Ordering::AcqRel);
Some(TopologyWriteGuard {
topology: self,
index: Some(index),
published: None,
published_keys: None,
publish: true,
dirty: false,
})
}
}
pub(crate) struct TopologyWriteGuard<'a, T, Node>
where
T: Ord + Clone + Send + 'static,
Node: Send + 'static,
{
topology: &'a Topology<T, Node>,
index: Option<RwLockWriteGuard<'a, NodeIndex<T, Node>>>,
published: Option<PublishedNodeIndex<T, Node>>,
published_keys: Option<MutexGuard<'a, BTreeMap<usize, T>>>,
publish: bool,
dirty: bool,
}
impl<T, Node> ::core::ops::Deref for TopologyWriteGuard<'_, T, Node>
where
T: Ord + Clone + Send + 'static,
Node: Send + 'static,
{
type Target = NodeIndex<T, Node>;
fn deref(&self) -> &Self::Target {
self.index.as_deref().expect("topology guard already released")
}
}
impl<'a, T, Node> TopologyWriteGuard<'a, T, Node>
where
T: Ord + Clone + Send + 'static,
Node: Send + 'static,
{
pub(crate) fn enable_publication(&mut self) {
if self.publish {
return;
}
debug_assert!(
!self.dirty,
"publication must be enabled before mutating an opt-out topology guard"
);
self.topology.generation.fetch_add(1, Ordering::AcqRel);
self.publish = true;
}
fn ensure_publication_snapshot(&mut self) {
if self.published.is_none() {
self.published_keys = Some(self.topology.published_keys.lock());
self.published = Some(self.topology.published.snapshot());
}
}
fn rebuild_publication(&mut self) {
if self.published_keys.is_none() {
self.published_keys = Some(self.topology.published_keys.lock());
}
let canonical = self.index.as_deref().expect("topology guard already released");
let rebuilt = PublishedNodeIndex::from_canonical(canonical);
let rebuilt_keys = canonical
.iter()
.map(|(key, node)| (node_identity(node), key.clone()))
.collect();
**self
.published_keys
.as_mut()
.expect("publication identity lock was initialized") = rebuilt_keys;
self.published = Some(rebuilt);
self.dirty = true;
}
pub(crate) fn is_last_node(&self, node: &Arc<RwLock<Node>>) -> bool {
self.index
.as_deref()
.and_then(BTreeMap::last_key_value)
.is_some_and(|(_, candidate)| Arc::ptr_eq(candidate, node))
}
pub(crate) fn rekey_last_node(&mut self, old_key: &T, new_key: T, node: Arc<RwLock<Node>>) {
debug_assert!(!self.publish, "last-node rekey must use the opt-out guard");
debug_assert!(
!self.dirty,
"an opt-out guard may perform only one explicit last-node rekey"
);
debug_assert!(
self.is_last_node(&node),
"only the canonical last node may skip publication"
);
let index = self.index.as_deref_mut().expect("topology guard already released");
let removed = index.remove(old_key);
debug_assert!(
removed.as_ref().is_some_and(|removed| Arc::ptr_eq(removed, &node)),
"last-node rekey must remove its expected canonical route"
);
let replaced = index.insert(new_key, node);
debug_assert!(
replaced.is_none(),
"last-node rekey must not collide with another canonical route"
);
self.dirty = true;
}
fn repair_last_route_before_attach(&mut self) {
debug_assert!(self.publish, "attachment repair requires publication");
let Some((canonical_key, last_node)) = self
.index
.as_deref()
.expect("topology guard already released")
.last_key_value()
.map(|(key, node)| (key.clone(), node.clone()))
else {
return;
};
self.ensure_publication_snapshot();
let published_key = self
.published_keys
.as_ref()
.expect("publication identity map initialized")
.get(&node_identity(&last_node))
.cloned();
let Some(published_key) = published_key else {
self.rebuild_publication();
return;
};
if published_key == canonical_key {
return;
}
let repaired_consistently = {
let published = self.published.as_mut().expect("publication snapshot initialized");
let published_keys = self
.published_keys
.as_mut()
.expect("publication identity map initialized");
let removed = published.remove(&published_key);
let displaced = published.insert(canonical_key.clone(), last_node.clone());
published_keys.insert(node_identity(&last_node), canonical_key);
removed.is_some_and(|old_node| Arc::ptr_eq(&old_node, &last_node)) && displaced.is_none()
};
if !repaired_consistently {
self.rebuild_publication();
} else {
self.dirty = true;
}
}
pub(crate) fn insert(&mut self, key: T, node: Arc<RwLock<Node>>) -> Option<Arc<RwLock<Node>>> {
debug_assert!(self.publish, "generic topology insertion requires publication");
let replaced = self
.index
.as_deref_mut()
.expect("topology guard already released")
.insert(key.clone(), node.clone());
self.dirty = true;
if self.publish {
self.ensure_publication_snapshot();
if let Some(replaced) = &replaced {
let removed_consistently = {
let published = self.published.as_mut().expect("publication snapshot initialized");
let published_keys = self
.published_keys
.as_mut()
.expect("publication identity map initialized");
published_keys
.remove(&node_identity(replaced))
.and_then(|old_key| published.remove(&old_key))
.is_some_and(|old_node| Arc::ptr_eq(&old_node, replaced))
};
if !removed_consistently {
self.rebuild_publication();
return Some(replaced.clone());
}
}
let displaced = self
.published
.as_mut()
.expect("publication snapshot initialized")
.insert(key.clone(), node.clone());
self.published_keys
.as_mut()
.expect("publication identity map initialized")
.insert(node_identity(&node), key);
if displaced.is_some_and(|old_node| !Arc::ptr_eq(&old_node, &node)) {
self.rebuild_publication();
}
}
replaced
}
pub(crate) fn remove<Q>(&mut self, key: &Q) -> Option<Arc<RwLock<Node>>>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
debug_assert!(self.publish, "generic topology removal requires publication");
let removed = self
.index
.as_deref_mut()
.expect("topology guard already released")
.remove(key)?;
self.dirty = true;
if self.publish {
self.ensure_publication_snapshot();
let removed_consistently = {
let published = self.published.as_mut().expect("publication snapshot initialized");
let published_keys = self
.published_keys
.as_mut()
.expect("publication identity map initialized");
published_keys
.remove(&node_identity(&removed))
.and_then(|route_key| published.remove::<T>(&route_key))
.is_some_and(|route_node| Arc::ptr_eq(&route_node, &removed))
};
if !removed_consistently {
self.rebuild_publication();
}
}
Some(removed)
}
}
impl<T, Node> Drop for TopologyWriteGuard<'_, T, Node>
where
T: Ord + Clone + Send + 'static,
Node: Send + 'static,
{
fn drop(&mut self) {
#[cfg(debug_assertions)]
if self.publish && self.dirty {
let canonical = self.index.as_deref().expect("topology guard exists while validating");
let published = self.published.as_ref().expect("publishing guard carries snapshot");
let published_keys = self
.published_keys
.as_deref()
.expect("publishing guard carries identity map");
debug_assert_eq!(
canonical.len(),
published.len,
"canonical and published node counts diverged"
);
debug_assert_eq!(
canonical.len(),
published_keys.len(),
"canonical and published identity counts diverged"
);
for ((_, canonical_node), (route_key, route_node)) in canonical.iter().zip(published.iter()) {
debug_assert!(
Arc::ptr_eq(route_node, canonical_node),
"canonical and published node order diverged"
);
debug_assert!(
published_keys.get(&node_identity(canonical_node)) == Some(route_key),
"published identity key does not match the route index"
);
}
}
if !self.publish {
drop(self.published_keys.take());
drop(self.index.take());
return;
}
if !self.dirty {
self.topology.generation.fetch_add(1, Ordering::Release);
drop(self.published_keys.take());
drop(self.index.take());
return;
}
let retired = self.topology.published.replace(
self.published
.take()
.expect("publishing guard must carry a read snapshot"),
);
self.topology.generation.fetch_add(1, Ordering::Release);
self.topology.published.retire(retired);
drop(self.published_keys.take());
drop(self.index.take());
self.topology.published.advance();
}
}
pub(crate) fn no_identity_adoption<T>(_stored: &T, _incoming: &mut T) {}
fn first_for_borrowed_bound<'a, T, Q, V>(
index: &'a BTreeMap<T, V>,
bound: Bound<&Q>,
borrow_order_matches: bool,
) -> Option<(&'a T, &'a V)>
where
T: Ord + Borrow<Q>,
Q: Ord + ?Sized,
{
if borrow_order_matches {
return index.range::<Q, _>((bound, Bound::Unbounded)).next();
}
index.iter().find(|(key, _)| match bound {
Bound::Included(value) => <T as Borrow<Q>>::borrow(key) >= value,
Bound::Excluded(value) => <T as Borrow<Q>>::borrow(key) > value,
Bound::Unbounded => true,
})
}
fn first_published_for_borrowed_bound<'a, T, Q, V>(
index: &'a PublishedNodeIndex<T, V>,
bound: Bound<&Q>,
borrow_order_matches: bool,
) -> Option<(&'a T, &'a Arc<RwLock<V>>)>
where
T: Ord + Clone + Borrow<Q>,
Q: Ord + ?Sized,
{
if borrow_order_matches {
return index.first_for_bound(bound);
}
index.iter().find(|(key, _)| match bound {
Bound::Included(value) => <T as Borrow<Q>>::borrow(key) >= value,
Bound::Excluded(value) => <T as Borrow<Q>>::borrow(key) > value,
Bound::Unbounded => true,
})
}
fn node_for_borrowed_end<'a, T, Q, V>(
index: &'a BTreeMap<T, V>,
end: &Q,
borrow_order_matches: bool,
) -> Option<(&'a T, &'a V)>
where
T: Ord + Borrow<Q>,
Q: Ord + ?Sized,
{
if borrow_order_matches {
return index
.range::<Q, _>((Bound::Included(end), Bound::Unbounded))
.next()
.or_else(|| index.last_key_value());
}
let mut last_equal = None;
for entry @ (key, _) in index {
match <T as Borrow<Q>>::borrow(key).cmp(end) {
::core::cmp::Ordering::Less => {}
::core::cmp::Ordering::Equal => last_equal = Some(entry),
::core::cmp::Ordering::Greater => return Some(entry),
}
}
last_equal.or_else(|| index.last_key_value())
}
#[derive(Debug)]
pub struct BTreeSet<T, Node = Vec<T>>
where
T: Ord + Clone + 'static,
Node: NodeLike<T>,
{
pub(crate) index: Topology<T, Node>,
node_capacity: usize,
borrow_order_matches: bool,
#[cfg(feature = "cdc")]
event_id: AtomicU64,
}
impl<T: Ord + Clone + 'static, Node: NodeLike<T>> Default for BTreeSet<T, Node> {
fn default() -> Self {
Self {
index: Topology::new(),
node_capacity: DEFAULT_INNER_SIZE,
borrow_order_matches: true,
#[cfg(feature = "cdc")]
event_id: AtomicU64::new(0),
}
}
}
impl<T, Node> BTreeSet<T, Node>
where
T: Debug + Ord + Clone + Send,
Node: NodeLike<T> + Send + 'static,
{
pub fn new() -> Self {
Self::default()
}
pub fn with_maximum_node_size(node_capacity: usize) -> Self {
Self {
index: Topology::new(),
node_capacity,
borrow_order_matches: true,
#[cfg(feature = "cdc")]
event_id: AtomicU64::new(0),
}
}
pub(crate) fn with_grouped_borrow_routing(mut self) -> Self {
self.borrow_order_matches = false;
self
}
pub fn attach_node(&self, node: Node) {
self.attach_nodes(::core::iter::once(node));
}
pub fn attach_nodes(&self, nodes: impl IntoIterator<Item = Node>) {
let mut nodes = nodes.into_iter().peekable();
if nodes.peek().is_none() {
return;
}
let mut index = self.index.write();
index.repair_last_route_before_attach();
for node in nodes {
let node_id = node
.max()
.cloned()
.expect("node should contain at least one value to be correct node");
index.insert(node_id, Arc::new(RwLock::new(node)));
}
}
#[cfg(feature = "cdc")]
pub(crate) fn export_topology(&self) -> (usize, Vec<Vec<T>>) {
let index = self.index.read();
let nodes = index
.values()
.map(|node| node.read().iter().cloned().collect())
.collect();
(self.node_capacity, nodes)
}
#[allow(clippy::type_complexity)]
fn put_checked_inner<const EMIT_CDC: bool>(
&self,
value: T,
adopt: fn(&T, &mut T),
) -> Result<(Option<T>, Vec<ChangeEvent<T>>), (ArcRwLockWriteGuard<RawRwLock, Node>, usize, T)> {
loop {
let mut cdc = vec![];
let index = self.index.read();
let target_node_entry = match index.range(value.clone()..).next() {
Some(entry) => entry,
None => {
if let Some(last) = index.last_key_value() {
last
} else {
drop(index);
let mut spins = 0;
let mut index = loop {
if let Some(guard) = self.index.try_write() {
break guard;
}
if spins >= ROOT_PUBLICATION_SPIN_LIMIT {
break self.index.write();
}
spins += 1;
::core::hint::spin_loop();
};
if !index.is_empty() {
continue;
}
let mut first_node = Node::with_capacity(self.node_capacity);
first_node.insert(value.clone());
#[cfg(feature = "cdc")]
if EMIT_CDC {
let node_insertion = ChangeEvent::CreateNode {
event_id: self.event_id.fetch_add(1, Ordering::Relaxed).into(),
max_value: value.clone(),
};
cdc.push(node_insertion);
}
index.insert(value, Arc::new(RwLock::new(first_node)));
return Ok((None, cdc));
}
}
};
let mut node_guard = target_node_entry.1.clone().write_arc();
#[allow(unused_assignments)]
let mut operation = None;
if !node_guard.need_to_split(self.node_capacity, &value) {
let old_max = node_guard.max().cloned();
let (inserted, idx) = NodeLike::insert(&mut *node_guard, value.clone());
if inserted {
#[cfg(feature = "cdc")]
if EMIT_CDC {
let node_element_insertion = ChangeEvent::InsertAt {
event_id: self.event_id.fetch_add(1, Ordering::Relaxed).into(),
max_value: old_max.clone().unwrap_or(value.clone()),
value: value.clone(),
index: idx,
};
cdc.push(node_element_insertion);
}
if node_guard.max().cloned() == old_max {
return Ok((None, cdc));
}
operation = Some(Operation::UpdateMax(
target_node_entry.1.clone(),
target_node_entry.0.clone(),
));
} else {
return Err((node_guard, idx, old_max.unwrap()));
}
} else {
operation = Some(Operation::Split(
target_node_entry.1.clone(),
target_node_entry.0.clone(),
value.clone(),
));
}
drop(node_guard);
drop(index);
let op = operation.unwrap();
let mut index = match &op {
Operation::UpdateMax(_, _) => self.index.write_rekey(),
Operation::Split(_, _, _) | Operation::MakeUnreachable(_, _) => self.index.write(),
};
match &op {
Operation::Split(_, _, _) => {
if let Ok((value, value_cdc)) = op.commit::<EMIT_CDC>(&mut index, adopt) {
#[cfg(feature = "cdc")]
if EMIT_CDC {
for unassigned_event in value_cdc {
let event_id = self.event_id.fetch_add(1, Ordering::Relaxed).into();
cdc.push(unassigned_event.assign_id(event_id));
}
}
return Ok((value, cdc));
} else {
continue;
}
}
Operation::UpdateMax(_, _) => {
return if let Ok((value, value_cdc)) = op.commit::<EMIT_CDC>(&mut index, adopt) {
#[cfg(feature = "cdc")]
if EMIT_CDC {
for unassigned_event in value_cdc {
let event_id = self.event_id.fetch_add(1, Ordering::Relaxed).into();
cdc.push(unassigned_event.assign_id(event_id));
}
}
Ok((value, cdc))
} else {
Ok((None, cdc))
}
}
Operation::MakeUnreachable(_, _) => unreachable!(),
}
}
}
fn put_inner<const EMIT_CDC: bool>(&self, value: T, adopt: fn(&T, &mut T)) -> (Option<T>, Vec<ChangeEvent<T>>) {
match self.put_checked_inner::<EMIT_CDC>(value.clone(), adopt) {
Ok(res) => res,
Err((mut node_guard, idx, max)) => {
let mut value = value;
if let Some(stored) = node_guard.get_ith(idx) {
adopt(stored, &mut value);
}
let mut cdc = vec![];
#[cfg(feature = "cdc")]
if EMIT_CDC {
if node_guard.len() == 1 {
let node_removal = ChangeEvent::RemoveNode {
event_id: self.event_id.fetch_add(1, Ordering::Relaxed).into(),
max_value: max.clone(),
};
let node_insertion = ChangeEvent::CreateNode {
event_id: self.event_id.fetch_add(1, Ordering::Relaxed).into(),
max_value: value.clone(),
};
cdc.push(node_removal);
cdc.push(node_insertion);
} else if idx == node_guard.len() - 1 {
let new_max = if node_guard.len() <= 1 {
None
} else {
node_guard.get_ith(node_guard.len() - 2)
};
let node_element_removal = ChangeEvent::RemoveAt {
event_id: self.event_id.fetch_add(1, Ordering::Relaxed).into(),
max_value: max.clone(),
value: value.clone(),
index: idx,
};
let node_element_insertion = ChangeEvent::InsertAt {
event_id: self.event_id.fetch_add(1, Ordering::Relaxed).into(),
max_value: new_max.expect("length was checked so should be ok").clone(),
value: value.clone(),
index: idx,
};
cdc.push(node_element_removal);
cdc.push(node_element_insertion);
} else {
let node_element_removal = ChangeEvent::RemoveAt {
event_id: self.event_id.fetch_add(1, Ordering::Relaxed).into(),
max_value: max.clone(),
value: value.clone(),
index: idx,
};
let node_element_insertion = ChangeEvent::InsertAt {
event_id: self.event_id.fetch_add(1, Ordering::Relaxed).into(),
max_value: max.clone(),
value: value.clone(),
index: idx,
};
cdc.push(node_element_removal);
cdc.push(node_element_insertion);
}
}
(NodeLike::replace(&mut *node_guard, idx, value.clone()), cdc)
}
}
}
pub(crate) fn put(&self, value: T) -> Option<T> {
self.put_inner::<false>(value, no_identity_adoption).0
}
#[cfg(feature = "multimap")]
pub(crate) fn put_with(&self, value: T, adopt: fn(&T, &mut T)) -> Option<T> {
self.put_inner::<false>(value, adopt).0
}
#[allow(clippy::type_complexity)]
pub(crate) fn put_checked(
&self,
value: T,
) -> Result<(Option<T>, Vec<ChangeEvent<T>>), (ArcRwLockWriteGuard<RawRwLock, Node>, usize, T)> {
self.put_checked_inner::<false>(value, no_identity_adoption)
}
pub(crate) fn put_cdc(&self, value: T) -> (Option<T>, Vec<ChangeEvent<T>>) {
self.put_inner::<true>(value, no_identity_adoption)
}
#[cfg(all(feature = "multimap", feature = "cdc"))]
pub(crate) fn put_cdc_with(&self, value: T, adopt: fn(&T, &mut T)) -> (Option<T>, Vec<ChangeEvent<T>>) {
self.put_inner::<true>(value, adopt)
}
#[allow(clippy::type_complexity)]
pub(crate) fn put_cdc_checked(
&self,
value: T,
) -> Result<(Option<T>, Vec<ChangeEvent<T>>), (ArcRwLockWriteGuard<RawRwLock, Node>, usize, T)> {
self.put_checked_inner::<true>(value, no_identity_adoption)
}
pub fn insert(&self, value: T) -> bool {
self.put(value).is_none()
}
fn remove_inner<const EMIT_CDC: bool, Q>(&self, value: &Q) -> (Option<T>, Vec<ChangeEvent<T>>)
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
let mut cdc = vec![];
let index = self.index.read();
if let Some(target_node_entry) =
first_for_borrowed_bound(&index, Bound::Included(value), self.borrow_order_matches)
.or_else(|| index.last_key_value())
{
let mut node_guard = target_node_entry.1.clone().write_arc();
let old_max = node_guard.max().cloned();
let deleted = NodeLike::delete(&mut *node_guard, value);
if deleted.is_none() {
return (None, cdc);
}
let (deleted, idx) = deleted.expect("should be ok as checked before");
let operation = if node_guard.len() > 0 {
#[cfg(feature = "cdc")]
if EMIT_CDC {
let node_element_removal = ChangeEvent::RemoveAt {
event_id: self.event_id.fetch_add(1, Ordering::Relaxed).into(),
max_value: old_max.clone().expect("Max value should exist as Node is not empty"),
value: deleted.clone(),
index: idx,
};
cdc.push(node_element_removal);
}
if old_max.as_ref() == node_guard.max() {
return (Some(deleted), cdc);
}
Some(Operation::UpdateMax(
target_node_entry.1.clone(),
target_node_entry.0.clone(),
))
} else {
Some(Operation::MakeUnreachable(
target_node_entry.1.clone(),
target_node_entry.0.clone(),
))
};
drop(node_guard);
drop(index);
let operation = operation.unwrap();
let mut index = match &operation {
Operation::UpdateMax(_, _) => self.index.write_rekey(),
Operation::Split(_, _, _) | Operation::MakeUnreachable(_, _) => self.index.write(),
};
return if let Ok((_, value_cdc)) = operation.commit::<EMIT_CDC>(&mut index, no_identity_adoption) {
#[cfg(feature = "cdc")]
if EMIT_CDC {
for unassigned_event in value_cdc {
let event_id = self.event_id.fetch_add(1, Ordering::Relaxed).into();
cdc.push(unassigned_event.assign_id(event_id));
}
}
(Some(deleted), cdc)
} else {
(Some(deleted), cdc)
};
}
(None, vec![])
}
pub fn remove_cdc<Q>(&self, value: &Q) -> (Option<T>, Vec<ChangeEvent<T>>)
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
self.remove_inner::<true, Q>(value)
}
pub fn remove<Q>(&self, value: &Q) -> Option<T>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
self.remove_inner::<false, Q>(value).0
}
#[inline(always)]
fn lock_node_for_value_optimistic<Q>(&self, value: &Q) -> Option<ArcRwLockReadGuard<RawRwLock, Node>>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
let node = {
let index = self.index.read();
match first_for_borrowed_bound(&index, Bound::Included(value), self.borrow_order_matches) {
Some((_, node)) => Some(node.clone()),
None => index
.last_key_value()
.map(|(_, node)| node.clone())
.or_else(|| index.first_key_value().map(|(_, node)| node.clone())),
}
}?;
Some(node.read_arc())
}
#[inline(always)]
fn lock_node_for_value<Q>(&self, value: &Q) -> Option<ArcRwLockReadGuard<RawRwLock, Node>>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
let mut retries = 0;
let mut writer_spins = 0;
loop {
let generation = self.index.generation.load(Ordering::Acquire);
if !generation.is_multiple_of(2) {
if writer_spins < ROOT_PUBLICATION_SPIN_LIMIT {
writer_spins += 1;
::core::hint::spin_loop();
} else {
yield_now();
}
continue;
}
writer_spins = 0;
if retries >= STABLE_READ_BLOCKING_FALLBACK_AFTER {
let index = self.index.read();
let node = match first_for_borrowed_bound(&index, Bound::Included(value), self.borrow_order_matches) {
Some((_, node)) => Some(node.clone()),
None => index
.last_key_value()
.map(|(_, node)| node.clone())
.or_else(|| index.first_key_value().map(|(_, node)| node.clone())),
}?;
let node_guard = node.read_arc();
drop(index);
return Some(node_guard);
}
let pin = self.index.published.domain.pin();
let snapshot = self.index.published.current.load(Ordering::Acquire);
let index = unsafe { &*snapshot };
let node =
match first_published_for_borrowed_bound(index, Bound::Included(value), self.borrow_order_matches) {
Some((_, node)) => Some(node.clone()),
None => index
.last_key_value()
.map(|(_, node)| node.clone())
.or_else(|| index.first_key_value().map(|(_, node)| node.clone())),
};
let Some(node) = node else {
if self.index.generation.load(Ordering::Acquire) == generation {
return None;
}
retries += 1;
continue;
};
drop(pin);
let node_guard = node.read_arc();
if self.index.generation.load(Ordering::Acquire) == generation {
return Some(node_guard);
}
retries += 1;
}
}
#[inline(always)]
fn get_with_guard<Q, R>(
node_guard: ArcRwLockReadGuard<RawRwLock, Node>,
value: &Q,
read: impl FnOnce(&T) -> R,
) -> Option<R>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
let position = node_guard.try_select(value)?;
node_guard.get_ith(position).map(read)
}
#[inline(always)]
pub(crate) fn get_with<Q, R>(&self, value: &Q, read: impl FnOnce(&T) -> R) -> Option<R>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
let mut retries = 0;
let mut writer_spins = 0;
let mut read = Some(read);
loop {
let generation = self.index.generation.load(Ordering::Acquire);
if !generation.is_multiple_of(2) {
if writer_spins < ROOT_PUBLICATION_SPIN_LIMIT {
writer_spins += 1;
::core::hint::spin_loop();
} else {
yield_now();
}
continue;
}
writer_spins = 0;
if retries >= STABLE_READ_BLOCKING_FALLBACK_AFTER {
let index = self.index.read();
let node = first_for_borrowed_bound(&index, Bound::Included(value), self.borrow_order_matches)
.or_else(|| index.last_key_value())
.or_else(|| index.first_key_value())
.map(|(_, node)| node.clone())?;
let node_guard = node.read_arc();
drop(index);
let position = node_guard.try_select(value)?;
return node_guard
.get_ith(position)
.map(read.take().expect("read closure is consumed only on return"));
}
let pin = self.index.published.domain.pin();
let snapshot = self.index.published.current.load(Ordering::Acquire);
let index = unsafe { &*snapshot };
let node = first_published_for_borrowed_bound(index, Bound::Included(value), self.borrow_order_matches)
.or_else(|| index.last_key_value())
.or_else(|| index.first_key_value())
.map(|(_, node)| node);
let Some(node) = node else {
if self.index.generation.load(Ordering::Acquire) == generation {
return None;
}
retries += 1;
continue;
};
if let Some(node_guard) = node.try_read() {
if self.index.generation.load(Ordering::Acquire) != generation {
retries += 1;
continue;
}
let position = node_guard.try_select(value)?;
return node_guard
.get_ith(position)
.map(read.take().expect("read closure is consumed only on return"));
}
let node = node.clone();
drop(pin);
let node_guard = node.read_arc();
if self.index.generation.load(Ordering::Acquire) != generation {
retries += 1;
continue;
}
let position = node_guard.try_select(value)?;
return node_guard
.get_ith(position)
.map(read.take().expect("read closure is consumed only on return"));
}
}
#[inline(always)]
pub(crate) fn get_with_optimistic<Q, R>(&self, value: &Q, read: impl FnOnce(&T) -> R) -> Option<R>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
Self::get_with_guard(self.lock_node_for_value_optimistic(value)?, value, read)
}
pub fn contains<Q>(&self, value: &Q) -> bool
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
self.get_with(value, |_| ()).is_some()
}
pub fn get<'a, Q>(&'a self, value: &'a Q) -> Option<Ref<T, Node>>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
{
if let Some(node_guard) = self.lock_node_for_value(value) {
let potential_position = node_guard.try_select(value);
if let Some(position) = potential_position {
return Some(Ref {
node_guard,
position,
phantom_data: PhantomData,
});
}
}
None
}
pub fn len(&self) -> usize {
self.index.read().values().map(|node| node.read().len()).sum()
}
pub fn is_empty(&self) -> bool {
self.index.read().values().all(|node| node.read().is_empty())
}
pub fn capacity(&self) -> usize {
self.index
.read()
.values()
.map(|node| {
let guard = node.read();
guard.capacity()
})
.sum()
}
pub fn node_count(&self) -> usize {
self.index.read().len()
}
}
impl<T> FromIterator<T> for BTreeSet<T>
where
T: Debug + Ord + Clone + Send,
{
fn from_iter<K: IntoIterator<Item = T>>(iter: K) -> Self {
let btree = BTreeSet::new();
iter.into_iter().for_each(|item| {
btree.insert(item);
});
btree
}
}
impl<T, const N: usize> From<[T; N]> for BTreeSet<T>
where
T: Debug + Ord + Clone + Send,
{
fn from(value: [T; N]) -> Self {
let btree: BTreeSet<T> = Default::default();
value.into_iter().for_each(|item| {
btree.insert(item);
});
btree
}
}
pub struct Iter<'a, T, Node>
where
T: Debug + Ord + Clone + Send + 'static,
Node: NodeLike<T> + Send + 'static,
{
tree: &'a BTreeSet<T, Node>,
current_front_batch: Option<alloc::vec::IntoIter<T>>,
current_back_batch: Option<alloc::vec::IntoIter<T>>,
exhausted_front_node: Option<Arc<RwLock<Node>>>,
exhausted_back_node: Option<Arc<RwLock<Node>>>,
front_partial: Option<(Arc<RwLock<Node>>, usize)>,
back_partial: Option<(Arc<RwLock<Node>>, usize)>,
front_batch_limit: usize,
back_batch_limit: usize,
current_front_value: Option<T>,
current_back_value: Option<T>,
met: bool,
}
const INITIAL_BATCH: usize = 4;
const MAX_BATCH: usize = 4096;
impl<'a, T, Node> Iter<'a, T, Node>
where
T: Debug + Ord + Clone + Send + 'static,
Node: NodeLike<T> + Send + 'static,
{
pub fn new(btree: &'a BTreeSet<T, Node>) -> Self {
Self {
tree: btree,
current_front_batch: None,
current_back_batch: None,
exhausted_front_node: None,
exhausted_back_node: None,
front_partial: None,
back_partial: None,
front_batch_limit: INITIAL_BATCH,
back_batch_limit: INITIAL_BATCH,
current_front_value: None,
current_back_value: None,
met: false,
}
}
fn install_front_batch(&mut self) -> bool {
let index = self.tree.index.read();
let candidate = match self.current_front_value.as_ref() {
Some(last_yielded) => index.range((Bound::Excluded(last_yielded), Bound::Unbounded)).next(),
None => index.first_key_value(),
};
let entry = match (candidate, self.exhausted_front_node.as_ref()) {
(Some((key, node)), Some(exhausted)) if Arc::ptr_eq(node, exhausted) => {
index.range((Bound::Excluded(key), Bound::Unbounded)).next()
}
(candidate, _) => candidate,
};
let Some((_, entry)) = entry else {
return false;
};
let node = entry.clone();
let guard = node.read_arc();
drop(index);
let rank_skip = self
.current_front_value
.as_ref()
.and_then(|value| guard.rank(Bound::Excluded(value), true))
.map_or(0, |rank| rank + 1);
let partial_skip = match self.front_partial.as_ref() {
Some((partial, taken)) if Arc::ptr_eq(partial, &node) => *taken,
_ => 0,
};
let skip = if self.current_front_value.is_some() {
rank_skip
} else {
partial_skip
};
let available = guard.len().saturating_sub(skip);
let take = available.min(self.front_batch_limit);
let batch = guard.iter().skip(skip).take(take).cloned().collect::<Vec<_>>();
drop(guard);
if take == available {
self.exhausted_front_node = Some(node);
self.front_partial = None;
} else {
self.exhausted_front_node = None;
self.front_partial = Some((node, skip + take));
}
self.front_batch_limit = self.front_batch_limit.saturating_mul(2).min(MAX_BATCH);
self.current_front_batch = Some(batch.into_iter());
true
}
fn install_back_batch(&mut self) -> bool {
let index = self.tree.index.read();
let candidate = match self.current_back_value.as_ref() {
Some(last_yielded) => index
.range((Bound::Included(last_yielded), Bound::Unbounded))
.next()
.or_else(|| index.last_key_value()),
None => index.last_key_value(),
};
let entry = match (candidate, self.exhausted_back_node.as_ref()) {
(Some((key, node)), Some(exhausted)) if Arc::ptr_eq(node, exhausted) => index.range(..key).next_back(),
(candidate, _) => candidate,
};
let Some((_, entry)) = entry else {
return false;
};
let node = entry.clone();
let guard = node.read_arc();
drop(index);
let truncate = self
.current_back_value
.as_ref()
.and_then(|value| guard.rank(Bound::Excluded(value), false))
.map_or(0, |rank| rank + 1);
let partial_truncate = match self.back_partial.as_ref() {
Some((partial, taken)) if Arc::ptr_eq(partial, &node) => *taken,
_ => 0,
};
let truncate = if self.current_back_value.is_some() {
truncate
} else {
partial_truncate
};
let available = guard.len().saturating_sub(truncate);
let take = available.min(self.back_batch_limit);
let skip = available - take;
let batch = guard.iter().skip(skip).take(take).cloned().collect::<Vec<_>>();
drop(guard);
if take == available {
self.exhausted_back_node = Some(node);
self.back_partial = None;
} else {
self.exhausted_back_node = None;
self.back_partial = Some((node, truncate + take));
}
self.back_batch_limit = self.back_batch_limit.saturating_mul(2).min(MAX_BATCH);
self.current_back_batch = Some(batch.into_iter());
true
}
}
impl<'a, T, Node> Iterator for Iter<'a, T, Node>
where
T: Debug + Ord + Clone + Send + 'static,
Node: NodeLike<T> + Send + 'static,
{
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
loop {
if self.met {
return None;
}
if self.current_front_batch.is_none() && !self.install_front_batch() {
return None;
}
let batch = self.current_front_batch.as_mut().expect("installed above");
if let Some(value) = batch.next() {
if let Some(current_front_value) = self.current_front_value.as_ref() {
if value.le(current_front_value) {
continue;
}
}
if let Some(current_back_value) = self.current_back_value.as_ref() {
if value.ge(current_back_value) {
self.met = true;
return None;
}
}
self.current_front_value = Some(value.clone());
return Some(value);
} else {
self.current_front_batch = None;
}
}
}
}
impl<'a, T, Node> DoubleEndedIterator for Iter<'a, T, Node>
where
T: Debug + Ord + Clone + Send + 'static,
Node: NodeLike<T> + Send + 'static,
{
fn next_back(&mut self) -> Option<Self::Item> {
loop {
if self.met {
return None;
}
if self.current_back_batch.is_none() && !self.install_back_batch() {
return None;
}
let batch = self.current_back_batch.as_mut().expect("installed above");
if let Some(value) = batch.next_back() {
if let Some(current_back_value) = self.current_back_value.as_ref() {
if value.ge(current_back_value) {
continue;
}
}
if let Some(current_front_value) = self.current_front_value.as_ref() {
if value.le(current_front_value) {
self.met = true;
return None;
}
}
self.current_back_value = Some(value.clone());
return Some(value);
} else {
self.current_back_batch = None;
}
}
}
}
impl<'a, T: Debug + Ord + Clone + Send, Node: NodeLike<T> + Send + 'static> FusedIterator for Iter<'a, T, Node> {}
impl<'a, T, Node> IntoIterator for &'a BTreeSet<T, Node>
where
T: Debug + Ord + Send + Clone,
Node: NodeLike<T> + Send + 'static,
{
type Item = T;
type IntoIter = Iter<'a, T, Node>;
fn into_iter(self) -> Self::IntoIter {
Iter::new(self)
}
}
pub struct Range<'a, T, Node>
where
T: Debug + Ord + Clone + Send + 'static,
Node: NodeLike<T> + Send + 'static,
{
iter: Iter<'a, T, Node>,
}
impl<'a, T, Node> Range<'a, T, Node>
where
T: Debug + Ord + Clone + Send + 'static,
Node: NodeLike<T> + Send + 'static,
{
pub fn new<Q, R>(btree: &'a BTreeSet<T, Node>, range: R) -> Self
where
T: Borrow<Q>,
Q: Ord + ?Sized,
R: RangeBounds<Q>,
{
let index = btree.index.read();
let start_bound = range.start_bound();
let end_bound = range.end_bound();
let mut met = match (start_bound, end_bound) {
(Bound::Included(start), Bound::Included(end)) => start > end,
(Bound::Included(start), Bound::Excluded(end))
| (Bound::Excluded(start), Bound::Included(end))
| (Bound::Excluded(start), Bound::Excluded(end)) => start >= end,
_ => false,
};
let current_front_entry = first_for_borrowed_bound(&index, start_bound, btree.borrow_order_matches);
let front_value = if let Some((front_key, front_node)) = current_front_entry {
let front_guard = front_node.clone().read_arc();
let rank = match start_bound {
Bound::Included(v) => front_guard.rank(Bound::Included(v), true),
Bound::Excluded(v) => front_guard.rank(Bound::Excluded(v), true),
Bound::Unbounded => None,
};
if let Some(rank) = rank {
let value = front_guard.iter().nth(rank).cloned();
drop(front_guard);
value
} else {
drop(front_guard);
if let Some((_, pre_front_node)) = index.range::<T, _>(..front_key).next_back() {
let pre_front_guard = pre_front_node.clone().read_arc();
pre_front_guard.iter().last().cloned()
} else {
None
}
}
} else {
None
};
let current_back_entry = match end_bound {
Bound::Included(end) | Bound::Excluded(end) => {
node_for_borrowed_end(&index, end, btree.borrow_order_matches)
}
Bound::Unbounded => index.last_key_value(),
};
let back_value = if let Some((back_key, back_node)) = current_back_entry {
let back_guard = back_node.clone().read_arc();
let rank = match end_bound {
Bound::Included(v) => back_guard.rank(Bound::Included(v), false),
Bound::Excluded(v) => back_guard.rank(Bound::Excluded(v), false),
Bound::Unbounded => None,
};
if let Some(rank) = rank {
let value = back_guard.iter().nth_back(rank).cloned();
drop(back_guard);
value
} else {
drop(back_guard);
if let Some((_, next_back_node)) = index
.range::<T, _>((Bound::Excluded(back_key), Bound::Unbounded))
.next()
{
let next_back_guard = next_back_node.clone().read_arc();
next_back_guard.iter().next().cloned()
} else {
None
}
}
} else {
None
};
if front_value.is_none() && back_value.is_none() {
if start_bound != Bound::Unbounded || end_bound != Bound::Unbounded {
if let Some(max) = index
.last_key_value()
.and_then(|(_, node)| node.clone().read_arc().max().cloned())
{
if let Bound::Included(v) = start_bound {
if v > max.borrow() {
met = true;
}
} else if let Bound::Excluded(v) = start_bound {
if v >= max.borrow() {
met = true;
}
}
}
if let Some(min) = index
.first_key_value()
.and_then(|(_, node)| node.clone().read_arc().min().cloned())
{
if let Bound::Included(v) = end_bound {
if v < min.borrow() {
met = true;
}
} else if let Bound::Excluded(v) = end_bound {
if v <= min.borrow() {
met = true;
}
}
}
}
}
Self {
iter: Iter {
tree: btree,
current_front_batch: None,
current_back_batch: None,
exhausted_front_node: None,
exhausted_back_node: None,
front_partial: None,
back_partial: None,
front_batch_limit: INITIAL_BATCH,
back_batch_limit: INITIAL_BATCH,
current_front_value: front_value,
current_back_value: back_value,
met,
},
}
}
}
impl<'a, T, Node> Iterator for Range<'a, T, Node>
where
T: Debug + Ord + Clone + Send + 'static,
Node: NodeLike<T> + Send + 'static,
{
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
self.iter.next()
}
}
impl<'a, T, Node> DoubleEndedIterator for Range<'a, T, Node>
where
T: Debug + Ord + Clone + Send + 'static,
Node: NodeLike<T> + Send + 'static,
{
fn next_back(&mut self) -> Option<Self::Item> {
self.iter.next_back()
}
}
impl<'a, T, Node> FusedIterator for Range<'a, T, Node>
where
T: Debug + Ord + Clone + Send + 'static,
Node: NodeLike<T> + Send + 'static,
{
}
impl<'a, T, Node> BTreeSet<T, Node>
where
T: Debug + Ord + Clone + Send + 'static,
Node: NodeLike<T> + Send + 'static,
{
pub fn iter(&'a self) -> Iter<'a, T, Node> {
Iter::new(self)
}
pub fn range<Q, R>(&'a self, range: R) -> Range<'a, T, Node>
where
T: Borrow<Q>,
Q: Ord + ?Sized,
R: RangeBounds<Q>,
{
Range::new(self, range)
}
}
impl<T> BTreeSet<T>
where
T: Debug + Ord + Clone + Send + 'static,
{
pub fn remove_range<R, Q>(&self, range: R)
where
Q: Ord + ?Sized,
T: Borrow<Q>,
R: RangeBounds<Q>,
{
let mut detached_nodes = Vec::new();
let mut index = self.index.write();
let start_bound = range.start_bound();
let end_bound = range.end_bound();
let Some((front_key, front_node)) = first_for_borrowed_bound(&index, start_bound, self.borrow_order_matches)
else {
return;
};
let front_key = front_key.clone();
let front_node = front_node.clone();
let back_entry = match end_bound {
Bound::Included(end) | Bound::Excluded(end) => {
node_for_borrowed_end(&index, end, self.borrow_order_matches)
}
Bound::Unbounded => index.last_key_value(),
};
let Some((back_key, back_node)) = back_entry else {
return;
};
let back_key = back_key.clone();
let back_node = back_node.clone();
if back_key < front_key {
return;
}
let removed_prefix_len = |guard: &Vec<T>| -> usize {
match end_bound {
Bound::Included(end) => guard.rank(Bound::Excluded(end), true).map_or(0, |last| last + 1),
Bound::Excluded(end) => guard.rank(Bound::Included(end), true).map_or(0, |last| last + 1),
Bound::Unbounded => guard.len(),
}
};
if Arc::ptr_eq(&front_node, &back_node) {
let mut guard = front_node.clone().write_arc();
let front_position = guard.rank(start_bound, true).map_or(0, |last| last + 1);
let back_position = removed_prefix_len(&guard);
if back_position <= front_position {
return;
}
let original_len = guard.len();
guard.drain(front_position..back_position);
if back_position == original_len {
index.remove::<T>(&front_key);
if let Some(new_max) = guard.last().cloned() {
index.insert(new_max, front_node);
}
}
return;
}
let mut front_guard = front_node.clone().write_arc();
let mut back_guard = back_node.clone().write_arc();
let front_position = front_guard.rank(start_bound, true).map_or(0, |last| last + 1);
let back_position = removed_prefix_len(&back_guard);
let middle_keys = index
.range::<T, _>((Bound::Excluded(&front_key), Bound::Excluded(&back_key)))
.map(|(key, _)| key.clone())
.collect::<Vec<_>>();
for key in middle_keys {
let node = index
.remove::<T>(&key)
.expect("middle key was collected under the write lock");
let mut removed_node = node.write_arc();
detached_nodes.push(::core::mem::take(&mut *removed_node));
}
index.remove::<T>(&front_key);
front_guard.drain(front_position..);
if !front_guard.is_empty() {
let new_front_max = front_guard.last().unwrap().clone();
index.insert(new_front_max, front_node);
}
if back_position >= back_guard.len() {
index.remove::<T>(&back_key);
back_guard.drain(..);
} else if back_position > 0 {
back_guard.drain(..back_position);
}
}
}
#[cfg(test)]
mod tests {
use crate::concurrent::operation::Operation;
use crate::concurrent::set::{BTreeSet, Iter, DEFAULT_INNER_SIZE, INITIAL_BATCH};
use crate::core::node::NodeLike;
use rand::Rng;
use std::collections::HashSet;
use std::ops::Bound::{self, Included};
use std::sync::mpsc;
use std::sync::{Arc, Barrier, Mutex};
use std::thread;
use std::time::Duration;
#[test]
fn test_node_size_two_preserves_all_u64_values() {
let set = BTreeSet::<u64>::with_maximum_node_size(2);
for value in 0..10_u64 {
set.insert(value);
}
assert_eq!(set.iter().collect::<Vec<_>>(), (0..10).collect::<Vec<_>>());
}
#[test]
fn test_node_size_three_preserves_all_u8_values() {
let set = BTreeSet::<u8>::with_maximum_node_size(3);
for value in 0..20_u8 {
set.insert(value);
}
assert_eq!(set.iter().collect::<Vec<_>>(), (0..20).collect::<Vec<_>>());
}
#[test]
fn concurrent_first_writers_preserve_disjoint_ranges() {
const WRITERS: u64 = 8;
const VALUES_PER_WRITER: u64 = 1_000;
let set = Arc::new(BTreeSet::<u64>::new());
let start = Arc::new(Barrier::new(WRITERS as usize));
let handles = (0..WRITERS)
.map(|writer| {
let set = Arc::clone(&set);
let start = Arc::clone(&start);
thread::spawn(move || {
start.wait();
let first = writer * VALUES_PER_WRITER;
for value in first..first + VALUES_PER_WRITER {
assert!(set.insert(value));
}
})
})
.collect::<Vec<_>>();
for handle in handles {
handle.join().unwrap();
}
let expected = (0..WRITERS * VALUES_PER_WRITER).collect::<Vec<_>>();
assert_eq!(set.len(), expected.len());
assert_eq!(set.iter().collect::<Vec<_>>(), expected);
}
#[test]
fn published_point_reads_remain_definitive_across_splits() {
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
const STABLE_KEYS: usize = 256;
const FINAL_KEYS: usize = 4_096;
const READERS: usize = 4;
let set = Arc::new(BTreeSet::<usize>::with_maximum_node_size(8));
for key in 0..STABLE_KEYS {
set.insert(key);
}
let start = Arc::new(Barrier::new(READERS + 1));
let done = Arc::new(AtomicBool::new(false));
let published_up_to = Arc::new(AtomicUsize::new(STABLE_KEYS - 1));
let readers = (0..READERS)
.map(|reader| {
let set = Arc::clone(&set);
let start = Arc::clone(&start);
let done = Arc::clone(&done);
let published_up_to = Arc::clone(&published_up_to);
thread::spawn(move || {
start.wait();
let mut probe = reader;
while !done.load(Ordering::Acquire) {
let key = probe % STABLE_KEYS;
assert_eq!(set.get_with(&key, |value| *value), Some(key));
let newest_acknowledged = published_up_to.load(Ordering::Acquire);
assert_eq!(
set.get_with(&newest_acknowledged, |value| *value),
Some(newest_acknowledged),
"an acknowledged insert disappeared from the published route"
);
probe += READERS;
}
})
})
.collect::<Vec<_>>();
start.wait();
for key in STABLE_KEYS..FINAL_KEYS {
set.insert(key);
published_up_to.store(key, Ordering::Release);
}
done.store(true, Ordering::Release);
for reader in readers {
reader.join().unwrap();
}
assert_eq!(set.len(), FINAL_KEYS);
for key in 0..FINAL_KEYS {
assert_eq!(set.get_with(&key, |value| *value), Some(key));
}
}
#[test]
fn published_pointer_survives_reclamation_interleavings() {
use std::sync::atomic::{AtomicBool, Ordering};
let set = Arc::new(BTreeSet::<usize>::with_maximum_node_size(2));
for key in 0..8 {
set.insert(key);
}
let start = Arc::new(Barrier::new(2));
let done = Arc::new(AtomicBool::new(false));
let reader_set = Arc::clone(&set);
let reader_start = Arc::clone(&start);
let reader_done = Arc::clone(&done);
let reader = thread::spawn(move || {
reader_start.wait();
let mut probe = 0;
while !reader_done.load(Ordering::Acquire) {
let key = probe % 8;
assert_eq!(reader_set.get_with(&key, |value| *value), Some(key));
probe += 1;
}
});
start.wait();
for key in 8..80 {
set.insert(key);
}
for key in 8..80 {
assert_eq!(set.remove(&key), Some(key));
}
done.store(true, Ordering::Release);
reader.join().unwrap();
for key in 0..8 {
assert_eq!(set.get_with(&key, |value| *value), Some(key));
}
}
#[test]
fn published_route_chunks_split_and_merge_without_losing_keys() {
let set = BTreeSet::<usize>::with_maximum_node_size(2);
for key in 0..600 {
assert!(set.insert(key));
}
for key in (0..600).step_by(2) {
assert_eq!(set.remove(&key), Some(key));
}
for key in 0..600 {
assert_eq!(set.contains(&key), key % 2 == 1, "probe {key}");
}
for key in (1..600).step_by(2) {
assert_eq!(set.remove(&key), Some(key));
}
assert!(set.is_empty());
}
#[test]
fn mixed_structural_and_node_lock_paths_complete_without_deadlock() {
const THREADS: usize = 8;
const OPERATIONS: usize = 1_000;
let set = Arc::new(BTreeSet::<usize>::with_maximum_node_size(8));
for value in 0..256 {
set.insert(value);
}
let start = Arc::new(Barrier::new(THREADS));
let (done_tx, done_rx) = mpsc::channel();
let handles = (0..THREADS)
.map(|worker| {
let set = Arc::clone(&set);
let start = Arc::clone(&start);
let done_tx = done_tx.clone();
thread::spawn(move || {
start.wait();
for operation in 0..OPERATIONS {
let value = (operation * 17 + worker * 31) % 512;
match (operation + worker) % 5 {
0 => {
set.insert(value);
}
1 => {
set.remove(&value);
}
2 => {
let _ = set.contains(&value);
}
3 => {
let _ = set.get_with(&value, Clone::clone);
}
_ => {
set.remove_range(value..=value);
set.insert(value);
}
}
}
done_tx.send(()).unwrap();
})
})
.collect::<Vec<_>>();
drop(done_tx);
for _ in 0..THREADS {
done_rx
.recv_timeout(Duration::from_secs(10))
.expect("mixed structural/node-lock workload did not complete");
}
for handle in handles {
handle.join().unwrap();
}
}
#[test]
fn test_concurrent_insert() {
let set = Arc::new(BTreeSet::<i32>::new());
let num_threads = 128;
let operations_per_thread = 10000;
let mut handles = vec![];
let test_data: Vec<Vec<(i32, i32)>> = (0..num_threads)
.map(|_| {
let mut rng = rand::rng();
(0..operations_per_thread)
.map(|_| {
let value = rng.random_range(0..100000);
let operation = rng.random_range(0..2);
(operation, value)
})
.collect()
})
.collect();
let expected_values = Arc::new(Mutex::new(HashSet::new()));
for thread_idx in 0..num_threads {
let set_clone = Arc::clone(&set);
let expected_values = Arc::clone(&expected_values);
let thread_data = test_data[thread_idx].clone();
let handle = thread::spawn(move || {
for (operation, value) in thread_data {
if operation == 0 {
let _a = set_clone.insert(value);
expected_values.lock().unwrap().insert(value);
}
}
});
handles.push(handle);
}
for handle in handles {
handle.join().unwrap();
}
let expected_values = expected_values.lock().unwrap();
assert_eq!(set.len(), expected_values.len());
for value in expected_values.iter() {
assert!(set.contains(value));
}
}
#[test]
fn test_insert_desc() {
let set = Arc::new(BTreeSet::<i32>::new());
assert!(set.insert(2));
assert!(set.insert(1));
}
#[test]
fn test_insert_st() {
let set = Arc::new(BTreeSet::<i32>::new());
let mut rng = rand::rng();
let n = 2048 * 100;
let range = 0..n;
let mut inserted_values = HashSet::new();
for _ in range {
let value = rng.random_range(0..n);
if inserted_values.insert(value) {
set.insert(value);
}
}
assert_eq!(
set.len(),
inserted_values.len(),
"Length did not match, missing: {:?}",
set.index
.read()
.values()
.flat_map(|node| node.read().iter().cloned().collect::<Vec<_>>())
.collect::<HashSet<_>>()
.symmetric_difference(&inserted_values)
.collect::<Vec<_>>()
);
for i in inserted_values {
assert!(
set.contains(&i),
"Did not find: {} with index: {:?}",
i,
set.index.read().keys().cloned().collect::<Vec<_>>(),
);
}
}
#[test]
fn test_single_element() {
let set = BTreeSet::<i32>::new();
set.insert(1);
let mut iter = set.into_iter();
assert_eq!(iter.next(), Some(1));
assert_eq!(iter.next(), None);
assert_eq!(iter.next_back(), None);
}
#[test]
fn test_multiple_elements() {
let set = BTreeSet::<i32>::new();
set.insert(1);
set.insert(2);
set.insert(3);
let mut iter = set.into_iter();
assert_eq!(iter.next(), Some(1));
assert_eq!(iter.next_back(), Some(3));
assert_eq!(iter.next(), Some(2));
assert_eq!(iter.next(), None);
assert_eq!(iter.next_back(), None);
}
#[test]
fn test_bidirectional_iteration() {
let set = BTreeSet::<i32>::with_maximum_node_size(3);
for i in 1..=20 {
set.insert(i);
}
let mut iter = set.into_iter();
for i in 0..10 {
let tree = set.index.read().keys().cloned().collect::<Vec<_>>();
let expected_next = i + 1;
let actual_next = iter.next();
assert_eq!(actual_next, Some(expected_next), "Tree: {:?}", tree);
let expected_next_back = 20 - i;
let actual_next_back = iter.next_back();
assert_eq!(actual_next_back, Some(expected_next_back), "Tree: {:?}", tree);
}
assert_eq!(iter.next(), None);
assert_eq!(iter.next_back(), None);
}
#[test]
fn test_fused_iterator() {
let set = BTreeSet::<i32>::new();
set.insert(1);
let mut iter = set.into_iter();
assert_eq!(iter.next(), Some(1));
assert_eq!(iter.next(), None);
assert_eq!(iter.next(), None);
}
#[test]
fn test_fused_iterator_back() {
let set = BTreeSet::<i32>::new();
set.insert(1);
let mut iter = set.into_iter();
assert_eq!(iter.next_back(), Some(1));
assert_eq!(iter.next_back(), None);
assert_eq!(iter.next_back(), None);
}
#[test]
fn test_out_of_bounds_range() {
let btree: BTreeSet<usize> = BTreeSet::from_iter(0..10);
assert_eq!(btree.range((Included(5), Included(10))).count(), 5);
assert_eq!(btree.range((Included(5), Included(11))).count(), 5);
assert_eq!(btree.range((Included(5), Included(10 + DEFAULT_INNER_SIZE))).count(), 5);
assert_eq!(btree.range((Included(0), Included(11))).count(), 10);
}
#[test]
fn test_iterating_over_blocks() {
let btree = BTreeSet::from_iter((0..(DEFAULT_INNER_SIZE + 10)).into_iter());
assert_eq!(btree.iter().count(), (0..(DEFAULT_INNER_SIZE + 10)).count());
let start = btree.range(0..DEFAULT_INNER_SIZE).into_iter().collect::<Vec<_>>();
assert_eq!(start, (0..DEFAULT_INNER_SIZE).collect::<Vec<_>>());
assert_eq!(
btree.range(0..=DEFAULT_INNER_SIZE).into_iter().collect::<Vec<_>>(),
(0..=DEFAULT_INNER_SIZE).collect::<Vec<_>>()
);
assert_eq!(
btree.range(0..=DEFAULT_INNER_SIZE + 1).count(),
(0..=DEFAULT_INNER_SIZE + 1).count()
);
assert_eq!(btree.iter().rev().count(), (0..(DEFAULT_INNER_SIZE + 10)).count());
assert_eq!(
btree.range(0..DEFAULT_INNER_SIZE).rev().count(),
(0..DEFAULT_INNER_SIZE).count()
);
assert_eq!(
btree.range(0..=DEFAULT_INNER_SIZE).rev().count(),
(0..=DEFAULT_INNER_SIZE).count()
);
assert_eq!(
btree.range(0..=DEFAULT_INNER_SIZE + 1).rev().count(),
(0..=DEFAULT_INNER_SIZE + 1).count()
);
}
#[test]
fn test_empty_set() {
let btree: BTreeSet<usize> = BTreeSet::new();
assert_eq!(btree.iter().count(), 0);
assert_eq!(btree.range(0..0).count(), 0);
assert_eq!(btree.range(0..).count(), 0);
assert_eq!(btree.range(..0).count(), 0);
assert_eq!(btree.range(..).count(), 0);
assert_eq!(btree.range(0..=0).count(), 0);
assert_eq!(btree.range(..1).count(), 0);
assert_eq!(btree.iter().rev().count(), 0);
assert_eq!(btree.range(0..0).rev().count(), 0);
assert_eq!(btree.range(..).rev().count(), 0);
assert_eq!(btree.range(..1).rev().count(), 0);
assert_eq!(btree.range(..DEFAULT_INNER_SIZE).count(), 0);
assert_eq!(btree.range(DEFAULT_INNER_SIZE..DEFAULT_INNER_SIZE * 2).count(), 0);
}
#[test]
fn test_remove_range() {
let btree = BTreeSet::from_iter(0..(DEFAULT_INNER_SIZE * 2));
let expected_len = DEFAULT_INNER_SIZE * 2;
let actual_len = btree.len();
assert_eq!(expected_len, actual_len);
btree.remove_range(5..15);
let expected_len = expected_len - 10;
let actual_len = btree.len();
assert_eq!(expected_len, actual_len);
btree.remove_range(DEFAULT_INNER_SIZE - 5..DEFAULT_INNER_SIZE + 5);
let expected_len = expected_len - 10;
let actual_len = btree.len();
assert_eq!(expected_len, actual_len);
btree.remove_range(..DEFAULT_INNER_SIZE / 2);
let expected_len = expected_len - (DEFAULT_INNER_SIZE / 2) + 10;
let actual_len = btree.len();
assert_eq!(expected_len, actual_len);
let from = (DEFAULT_INNER_SIZE * 3) / 2;
btree.remove_range(from..);
let expected_len = expected_len - DEFAULT_INNER_SIZE / 2;
let actual_len = btree.len();
assert_eq!(expected_len, actual_len);
btree.remove_range(..);
assert_eq!(btree.len(), 0);
for i in 0..(DEFAULT_INNER_SIZE * 2) {
btree.insert(i);
}
let expected_len = DEFAULT_INNER_SIZE * 2;
let actual_len = btree.len();
assert_eq!(expected_len, actual_len);
btree.remove_range((std::ops::Bound::Excluded(5), std::ops::Bound::Excluded(15)));
let expected_len = expected_len - 9;
let actual_len = btree.len();
assert_eq!(expected_len, actual_len);
btree.remove_range((
std::ops::Bound::Included(DEFAULT_INNER_SIZE),
std::ops::Bound::Excluded(DEFAULT_INNER_SIZE + 10),
));
let expected_len = expected_len - 10;
let actual_len = btree.len();
assert_eq!(expected_len, actual_len);
btree.remove_range(DEFAULT_INNER_SIZE * 3..DEFAULT_INNER_SIZE * 4);
let expected_len = expected_len;
let actual_len = btree.len();
assert_eq!(expected_len, actual_len);
btree.remove_range(DEFAULT_INNER_SIZE * 2 - 5..DEFAULT_INNER_SIZE * 3);
let expected_len = expected_len - 5;
let actual_len = btree.len();
assert_eq!(expected_len, actual_len);
}
#[test]
fn remove_range_end_bound_regressions() {
let set = BTreeSet::<u64>::with_maximum_node_size(4);
for value in 0..10 {
set.insert(value);
}
set.remove_range(7..);
assert_eq!(set.iter().collect::<Vec<_>>(), (0..7).collect::<Vec<_>>());
let set = BTreeSet::<u64>::with_maximum_node_size(4);
for value in 0..10 {
set.insert(value);
}
set.remove_range(..);
assert_eq!(set.len(), 0);
assert!(set.is_empty());
let set = BTreeSet::<u64>::with_maximum_node_size(4);
for value in 0..10 {
set.insert(value);
}
set.remove_range(3..=5);
assert_eq!(set.iter().collect::<Vec<_>>(), vec![0, 1, 2, 6, 7, 8, 9]);
let set = BTreeSet::<u64>::with_maximum_node_size(4);
for value in 0..10 {
set.insert(value);
}
set.remove_range(2..=2);
assert_eq!(set.iter().collect::<Vec<_>>(), vec![0, 1, 3, 4, 5, 6, 7, 8, 9]);
let set = BTreeSet::<u64>::with_maximum_node_size(3);
for value in 0..9 {
set.insert(value);
}
let boundary = *set.index.read().first_key_value().expect("node must exist").0;
set.remove_range(0..boundary);
let expected = (0..9).filter(|value| *value >= boundary).collect::<Vec<_>>();
assert_eq!(set.iter().collect::<Vec<_>>(), expected);
}
#[test]
fn remove_range_matches_btreeset_oracle() {
use std::ops::Bound;
fn oracle_case(node_size: usize, values: &[u64], start: Bound<u64>, end: Bound<u64>) {
let set = BTreeSet::<u64>::with_maximum_node_size(node_size);
for &value in values {
set.insert(value);
}
let mut oracle = values.iter().copied().collect::<std::collections::BTreeSet<_>>();
let range = (start, end);
oracle.retain(|value| !std::ops::RangeBounds::contains(&range, value));
set.remove_range(range);
assert_eq!(
set.iter().collect::<Vec<_>>(),
oracle.iter().copied().collect::<Vec<_>>(),
"node_size={node_size}, start={start:?}, end={end:?}"
);
assert_eq!(
set.len(),
oracle.len(),
"node_size={node_size}, start={start:?}, end={end:?}"
);
}
let values = (0..15u64).map(|value| value * 2).collect::<Vec<_>>();
let mut bounds = vec![Bound::Unbounded];
for probe in 0..=30u64 {
bounds.push(Bound::Included(probe));
bounds.push(Bound::Excluded(probe));
}
for node_size in [4usize, 7, 64] {
for &start in &bounds {
for &end in &bounds {
oracle_case(node_size, &values, start, end);
}
}
}
}
#[test]
fn remove_range_clears_detached_nodes() {
let set = BTreeSet::<u64>::with_maximum_node_size(4);
for value in 0..32 {
set.insert(value);
}
let detached = set
.index
.read()
.range((Included(&2), Bound::Unbounded))
.nth(1)
.unwrap()
.1
.clone();
let detached_values = detached.read().iter().copied().collect::<Vec<_>>();
assert!(detached_values.iter().all(|value| (2..30).contains(value)));
set.remove_range(2..30);
assert!(detached.read().is_empty());
assert!(detached_values.iter().all(|value| !set.contains(value)));
}
#[test]
fn remove_reaches_value_above_every_index_key() {
let set = BTreeSet::<u64>::new();
for value in [1u64, 2, 3] {
set.insert(value);
}
{
let node = set.index.read().last_key_value().expect("node must exist").1.clone();
let mut guard = node.write();
NodeLike::insert(&mut *guard, 5u64);
}
assert!(set.contains(&5));
assert_eq!(set.remove(&5), Some(5), "value above every index key must be removable");
assert!(!set.contains(&5));
assert_eq!(set.iter().collect::<Vec<_>>(), vec![1, 2, 3]);
}
fn drain_node_with_pending_unlink(set: &BTreeSet<u64>, values: &[u64], stale_key: u64) -> Operation<u64, Vec<u64>> {
let node = set.index.read().last_key_value().expect("node must exist").1.clone();
{
let mut guard = node.write();
for value in values {
NodeLike::delete(&mut *guard, value).expect("seeded value must be present");
}
}
Operation::MakeUnreachable(node, stale_key)
}
#[test]
fn split_commit_against_drained_node_fails_instead_of_dropping_insert() {
let set = BTreeSet::<u64>::new();
for seeded in [10u64, 20, 30] {
set.insert(seeded);
}
let node = set.index.read().last_key_value().expect("node must exist").1.clone();
let pending_split = Operation::Split(node.clone(), 30u64, 15u64);
{
let mut guard = node.write();
for seeded in [10u64, 20, 30] {
NodeLike::delete(&mut *guard, &seeded).expect("seeded value must be present");
}
}
assert!(pending_split
.commit::<false>(&mut set.index.write(), super::no_identity_adoption)
.is_err());
assert!(
set.index.read().get(&30).is_some(),
"drained node must stay linked for the retry"
);
assert!(set.insert(15));
assert!(set.contains(&15));
assert_eq!(set.remove(&15), Some(15));
assert!(set.is_empty());
}
#[cfg(feature = "cdc")]
#[test]
fn split_commit_against_drained_node_does_not_panic_in_cdc_build() {
let set = BTreeSet::<u64>::new();
for seeded in [10u64, 20, 30] {
set.insert(seeded);
}
let node = set.index.read().last_key_value().expect("node must exist").1.clone();
let pending_split = Operation::Split(node.clone(), 30u64, 15u64);
{
let mut guard = node.write();
for seeded in [10u64, 20, 30] {
NodeLike::delete(&mut *guard, &seeded).expect("seeded value must be present");
}
}
assert!(pending_split
.commit::<true>(&mut set.index.write(), super::no_identity_adoption)
.is_err());
assert!(set.index.read().get(&30).is_some());
let (old, _events) = set.put_cdc(15);
assert!(old.is_none());
assert!(set.contains(&15));
}
#[test]
fn insert_into_emptied_node_survives_stale_make_unreachable() {
for value in [5u64, 40u64] {
let set = BTreeSet::<u64>::new();
for seeded in [10u64, 20, 30] {
set.insert(seeded);
}
let pending_unlink = drain_node_with_pending_unlink(&set, &[10, 20, 30], 30);
assert!(set.insert(value));
let _ = pending_unlink.commit::<false>(&mut set.index.write(), super::no_identity_adoption);
assert!(set.contains(&value), "value {value} lost after stale unlink");
assert_eq!(set.iter().collect::<Vec<_>>(), vec![value]);
assert_eq!(set.remove(&value), Some(value));
assert!(!set.contains(&value));
assert_eq!(set.len(), 0);
}
}
#[test]
fn stale_make_unreachable_rekeys_refilled_node_instead_of_unlinking() {
for value in [5u64, 40u64] {
let set = BTreeSet::<u64>::new();
for seeded in [10u64, 20, 30] {
set.insert(seeded);
}
let node = set.index.read().last_key_value().expect("node must exist").1.clone();
let pending_unlink = drain_node_with_pending_unlink(&set, &[10, 20, 30], 30);
{
let mut guard = node.write();
NodeLike::insert(&mut *guard, value);
}
let pending_repair = Operation::UpdateMax(node.clone(), 30u64);
assert!(pending_unlink
.commit::<false>(&mut set.index.write(), super::no_identity_adoption)
.is_ok());
let _ = pending_repair.commit::<false>(&mut set.index.write(), super::no_identity_adoption);
assert!(set.contains(&value), "value {value} lost to stale unlink");
assert_eq!(set.remove(&value), Some(value));
assert!(set.is_empty());
}
}
#[test]
fn published_route_updates_when_a_non_last_boundary_shrinks() {
let set = BTreeSet::<u64>::with_maximum_node_size(8);
set.attach_nodes([vec![1, 10], vec![20, 30]]);
assert_eq!(set.remove(&10), Some(10));
assert!(set.insert(5));
assert!(set.contains(&5));
assert_eq!(set.get(&5).map(|value| *value.get()), Some(5));
}
#[test]
fn attach_repairs_a_stale_last_node_boundary() {
let set = BTreeSet::<u64>::with_maximum_node_size(8);
set.attach_node(vec![1, 10]);
assert_eq!(set.remove(&10), Some(10));
set.attach_node(vec![5, 20]);
assert_eq!(set.get(&5).map(|value| *value.get()), Some(5));
assert!(set.contains(&5));
}
#[test]
fn attach_repairs_a_stale_low_last_node_boundary() {
let set = BTreeSet::<u64>::with_maximum_node_size(8);
set.attach_node(vec![1, 10]);
assert!(set.insert(20));
set.attach_node(vec![25, 30]);
assert_eq!(set.get(&20).map(|value| *value.get()), Some(20));
assert_eq!(set.get(&25).map(|value| *value.get()), Some(25));
}
#[test]
fn attach_recovers_from_a_missing_published_identity() {
let set = BTreeSet::<u64>::with_maximum_node_size(8);
set.attach_node(vec![1, 10]);
let last_identity = {
let index = set.index.read();
super::node_identity(index.last_key_value().unwrap().1)
};
assert!(set.index.published_keys.lock().remove(&last_identity).is_some());
set.attach_node(vec![20, 30]);
for value in [1, 10, 20, 30] {
assert_eq!(set.get(&value).map(|found| *found.get()), Some(value));
}
}
#[test]
fn insert_recovers_from_a_missing_replaced_node_identity() {
let set = BTreeSet::<u64>::with_maximum_node_size(8);
set.attach_node(vec![1, 10]);
let old_node = set.index.read().last_key_value().unwrap().1.clone();
assert!(set
.index
.published_keys
.lock()
.remove(&super::node_identity(&old_node))
.is_some());
{
let mut index = set.index.write();
let replaced = index.insert(10, Arc::new(parking_lot::RwLock::new(vec![5, 10])));
assert!(replaced.is_some_and(|node| Arc::ptr_eq(&node, &old_node)));
}
assert!(!set.contains(&1));
assert!(set.contains(&5));
assert!(set.contains(&10));
}
#[test]
fn remove_recovers_from_a_missing_node_identity() {
let set = BTreeSet::<u64>::with_maximum_node_size(8);
set.attach_node(vec![1, 10]);
let old_node = set.index.read().last_key_value().unwrap().1.clone();
assert!(set
.index
.published_keys
.lock()
.remove(&super::node_identity(&old_node))
.is_some());
{
let mut index = set.index.write();
let removed = index.remove(&10).expect("canonical route exists");
assert!(Arc::ptr_eq(&removed, &old_node));
}
assert!(set.is_empty());
assert!(!set.contains(&1));
}
#[test]
fn missing_published_remove_does_not_clone_a_shared_chunk() {
let mut published = super::PublishedNodeIndex::<u64, Vec<u64>> {
chunks: Vec::new(),
len: 0,
};
for key in 0..16 {
published.insert(key, Arc::new(parking_lot::RwLock::new(vec![key])));
}
let snapshot = published.clone();
assert!(Arc::ptr_eq(&published.chunks[0], &snapshot.chunks[0]));
assert!(published.remove(&100).is_none());
assert!(Arc::ptr_eq(&published.chunks[0], &snapshot.chunks[0]));
}
#[test]
fn published_chunks_have_split_merge_hysteresis() {
let mut published = super::PublishedNodeIndex::<u64, Vec<u64>> {
chunks: Vec::new(),
len: 0,
};
for key in 0..=128 {
published.insert(key, Arc::new(parking_lot::RwLock::new(vec![key])));
}
assert_eq!(published.chunks.len(), 2);
for key in 0..33 {
assert!(published.remove(&key).is_some());
}
assert_eq!(published.chunks.len(), 1);
published.insert(0, Arc::new(parking_lot::RwLock::new(vec![0])));
assert_eq!(
published.chunks.len(),
1,
"one insert after a merge must not split again"
);
assert!(published.remove(&0).is_some());
assert_eq!(
published.chunks.len(),
1,
"one remove after a merge must not change chunking"
);
}
#[test]
fn attach_boundary_repair_is_a_complete_publication_by_itself() {
let set = BTreeSet::<u64>::with_maximum_node_size(8);
set.attach_node(vec![1, 10]);
assert_eq!(set.remove(&10), Some(10));
{
let mut index = set.index.write();
index.repair_last_route_before_attach();
}
set.attach_node(vec![5, 20]);
assert_eq!(set.get(&5).map(|found| *found.get()), Some(5));
}
#[cfg(debug_assertions)]
#[test]
#[should_panic(expected = "generic topology removal requires publication")]
fn publication_must_be_enabled_before_an_opt_out_guard_mutates() {
let set = BTreeSet::<u64>::with_maximum_node_size(8);
set.attach_node(vec![1, 10]);
let mut index = set.index.write_rekey();
index.remove(&10);
}
#[test]
fn published_point_routes_match_a_sequential_oracle_under_churn() {
let set = BTreeSet::<u64>::with_maximum_node_size(4);
let mut oracle = std::collections::BTreeSet::new();
let mut state = 0x8f4d_2a71_c390_6be5u64;
for step in 0..2_000 {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
let key = state % 64;
if state & 1 == 0 {
assert_eq!(set.insert(key), oracle.insert(key), "insert step {step}, key {key}");
} else {
assert_eq!(
set.remove(&key).is_some(),
oracle.remove(&key),
"remove step {step}, key {key}"
);
}
for probe in 0..64 {
assert_eq!(
set.contains(&probe),
oracle.contains(&probe),
"point route diverged at step {step}, probe {probe}"
);
}
assert_eq!(
set.iter().collect::<Vec<_>>(),
oracle.iter().copied().collect::<Vec<_>>()
);
}
}
#[test]
fn concurrent_remove_reinsert_over_emptying_nodes_preserves_all_keys() {
const THREADS: u64 = 4;
const ITERATIONS: u64 = 1_000;
let set = Arc::new(BTreeSet::<u64>::with_maximum_node_size(2));
for key in 0..THREADS {
set.insert(key);
}
let start = Arc::new(Barrier::new(THREADS as usize));
let (done_tx, done_rx) = mpsc::channel();
let handles = (0..THREADS)
.map(|key| {
let set = Arc::clone(&set);
let start = Arc::clone(&start);
let done_tx = done_tx.clone();
thread::spawn(move || {
start.wait();
for _ in 0..ITERATIONS {
let mut attempts = 0;
while set.remove(&key).is_none() {
let index = set.index.read();
let present = index.values().any(|node| node.read().contains(&key));
drop(index);
assert!(present, "acknowledged insert of {key} was lost");
attempts += 1;
assert!(attempts < 10_000, "key {key} present but never became removable");
std::hint::spin_loop();
}
assert!(set.insert(key), "{key} still present after acknowledged remove");
}
done_tx.send(()).unwrap();
})
})
.collect::<Vec<_>>();
drop(done_tx);
for _ in 0..THREADS {
done_rx
.recv_timeout(Duration::from_secs(30))
.expect("remove/reinsert workload did not complete in time");
}
for handle in handles {
handle.join().unwrap();
}
for key in 0..THREADS {
assert!(set.contains(&key), "key {key} lost after churn");
assert_eq!(set.remove(&key), Some(key));
}
assert!(set.is_empty());
}
#[test]
fn test_remove_single_element() {
let set = BTreeSet::<i32>::new();
set.insert(5);
assert!(set.contains(&5));
assert!(set.remove(&5).is_some());
assert!(!set.contains(&5));
assert!(!set.remove(&5).is_some());
}
#[test]
fn test_remove_multiple_elements() {
let set = BTreeSet::<i32>::new();
for i in 0..2048 {
set.insert(i);
}
for i in 0..2048 {
assert!(set.remove(&i).is_some());
assert!(!set.contains(&i));
}
assert_eq!(set.len(), 0);
}
#[test]
fn test_remove_non_existent() {
let set = BTreeSet::<i32>::new();
set.insert(5);
assert!(!set.remove(&10).is_some());
assert!(set.contains(&5));
}
#[test]
fn test_remove_stress() {
let set = Arc::new(BTreeSet::<i32>::new());
const NUM_ELEMENTS: i32 = 10000;
for i in 0..NUM_ELEMENTS {
set.insert(i);
}
assert_eq!(set.len(), NUM_ELEMENTS as usize, "Incorrect size after insertion");
let num_threads = 8;
let elements_per_thread = NUM_ELEMENTS / num_threads;
let handles: Vec<_> = (0..num_threads)
.map(|t| {
let set = Arc::clone(&set);
thread::spawn(move || {
for i in (t * elements_per_thread)..((t + 1) * elements_per_thread) {
if i % 2 == 1 {
assert!(set.remove(&i).is_some(), "Failed to remove {}", i);
}
}
})
})
.collect();
for handle in handles {
handle.join().unwrap();
}
assert_eq!(set.len(), NUM_ELEMENTS as usize / 2, "Incorrect size after removal");
for i in 0..NUM_ELEMENTS {
if i % 2 == 0 {
assert!(set.contains(&i), "Even number {} should be in the set", i);
} else {
assert!(!set.contains(&i), "Odd number {} should not be in the set", i);
}
}
}
#[test]
fn test_remove_all_elements() {
let set = BTreeSet::<i32>::new();
let n = 2048;
for i in 0..n {
set.insert(i);
}
for i in 0..n {
assert!(set.remove(&i).is_some(), "Failed to remove {}", i);
}
assert_eq!(set.len(), 0, "Set should be empty");
for i in 0..n {
assert!(!set.contains(&i), "Element {} should not be in the set", i);
}
}
#[test]
fn test_range_edge_cases() {
let set = BTreeSet::<i32>::with_maximum_node_size(10);
for i in 0..20 {
set.insert(i);
}
assert_eq!(set.range(0..=0).collect::<Vec<_>>(), vec![0]);
assert_eq!(set.range(0..1).collect::<Vec<_>>(), vec![0]);
assert_eq!(set.range(5..=5).collect::<Vec<_>>(), vec![5]);
assert_eq!(set.range(5..6).collect::<Vec<_>>(), vec![5]);
assert_eq!(set.range(10..=10).collect::<Vec<_>>(), vec![10]);
assert_eq!(set.range(10..11).collect::<Vec<_>>(), vec![10]);
assert_eq!(set.range(0..=3).collect::<Vec<_>>(), vec![0, 1, 2, 3]);
assert_eq!(set.range(0..3).collect::<Vec<_>>(), vec![0, 1, 2]);
assert_eq!(set.range(5..=8).collect::<Vec<_>>(), vec![5, 6, 7, 8]);
assert_eq!(set.range(5..8).collect::<Vec<_>>(), vec![5, 6, 7]);
assert_eq!(set.range(10..=13).collect::<Vec<_>>(), vec![10, 11, 12, 13]);
assert_eq!(set.range(10..13).collect::<Vec<_>>(), vec![10, 11, 12]);
assert_eq!(set.range(4..=4).collect::<Vec<_>>(), vec![4]);
assert_eq!(set.range(4..5).collect::<Vec<_>>(), vec![4]);
assert_eq!(set.range(9..=9).collect::<Vec<_>>(), vec![9]);
assert_eq!(set.range(9..10).collect::<Vec<_>>(), vec![9]);
assert_eq!(set.range(19..=19).collect::<Vec<_>>(), vec![19]);
assert_eq!(set.range(19..20).collect::<Vec<_>>(), vec![19]);
assert_eq!(set.range(17..=19).collect::<Vec<_>>(), vec![17, 18, 19]);
assert_eq!(set.range(17..20).collect::<Vec<_>>(), vec![17, 18, 19]);
assert_eq!(set.range(7..=9).collect::<Vec<_>>(), vec![7, 8, 9]);
assert_eq!(set.range(7..10).collect::<Vec<_>>(), vec![7, 8, 9]);
assert_eq!(set.range(2..=4).collect::<Vec<_>>(), vec![2, 3, 4]);
assert_eq!(set.range(2..5).collect::<Vec<_>>(), vec![2, 3, 4]);
assert_eq!(set.range(0..=4).collect::<Vec<_>>(), vec![0, 1, 2, 3, 4]);
assert_eq!(set.range(0..5).collect::<Vec<_>>(), vec![0, 1, 2, 3, 4]);
assert_eq!(set.range(5..=9).collect::<Vec<_>>(), vec![5, 6, 7, 8, 9]);
assert_eq!(set.range(5..10).collect::<Vec<_>>(), vec![5, 6, 7, 8, 9]);
assert_eq!(
set.range(10..=19).collect::<Vec<_>>(),
vec![10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
);
assert_eq!(
set.range(10..20).collect::<Vec<_>>(),
vec![10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
);
assert_eq!(set.range(3..=6).collect::<Vec<_>>(), vec![3, 4, 5, 6]);
assert_eq!(set.range(3..7).collect::<Vec<_>>(), vec![3, 4, 5, 6]);
assert_eq!(set.range(8..=11).collect::<Vec<_>>(), vec![8, 9, 10, 11]);
assert_eq!(set.range(8..12).collect::<Vec<_>>(), vec![8, 9, 10, 11]);
assert_eq!(set.range(0..=0).rev().collect::<Vec<_>>(), vec![0]);
assert_eq!(set.range(0..1).rev().collect::<Vec<_>>(), vec![0]);
assert_eq!(set.range(5..=5).rev().collect::<Vec<_>>(), vec![5]);
assert_eq!(set.range(5..6).rev().collect::<Vec<_>>(), vec![5]);
assert_eq!(set.range(10..=10).rev().collect::<Vec<_>>(), vec![10]);
assert_eq!(set.range(10..11).rev().collect::<Vec<_>>(), vec![10]);
assert_eq!(set.range(0..=3).rev().collect::<Vec<_>>(), vec![3, 2, 1, 0]);
assert_eq!(set.range(0..3).rev().collect::<Vec<_>>(), vec![2, 1, 0]);
assert_eq!(set.range(5..=8).rev().collect::<Vec<_>>(), vec![8, 7, 6, 5]);
assert_eq!(set.range(5..8).rev().collect::<Vec<_>>(), vec![7, 6, 5]);
assert_eq!(set.range(10..=13).rev().collect::<Vec<_>>(), vec![13, 12, 11, 10]);
assert_eq!(set.range(10..13).rev().collect::<Vec<_>>(), vec![12, 11, 10]);
assert_eq!(set.range(4..=4).rev().collect::<Vec<_>>(), vec![4]);
assert_eq!(set.range(4..5).rev().collect::<Vec<_>>(), vec![4]);
assert_eq!(set.range(9..=9).rev().collect::<Vec<_>>(), vec![9]);
assert_eq!(set.range(9..10).rev().collect::<Vec<_>>(), vec![9]);
assert_eq!(set.range(19..=19).rev().collect::<Vec<_>>(), vec![19]);
assert_eq!(set.range(19..20).rev().collect::<Vec<_>>(), vec![19]);
assert_eq!(set.range(17..=19).rev().collect::<Vec<_>>(), vec![19, 18, 17]);
assert_eq!(set.range(17..20).rev().collect::<Vec<_>>(), vec![19, 18, 17]);
assert_eq!(set.range(7..=9).rev().collect::<Vec<_>>(), vec![9, 8, 7]);
assert_eq!(set.range(7..10).rev().collect::<Vec<_>>(), vec![9, 8, 7]);
assert_eq!(set.range(2..=4).rev().collect::<Vec<_>>(), vec![4, 3, 2]);
assert_eq!(set.range(2..5).rev().collect::<Vec<_>>(), vec![4, 3, 2]);
assert_eq!(set.range(0..=4).rev().collect::<Vec<_>>(), vec![4, 3, 2, 1, 0]);
assert_eq!(set.range(0..5).rev().collect::<Vec<_>>(), vec![4, 3, 2, 1, 0]);
assert_eq!(set.range(5..=9).rev().collect::<Vec<_>>(), vec![9, 8, 7, 6, 5]);
assert_eq!(set.range(5..10).rev().collect::<Vec<_>>(), vec![9, 8, 7, 6, 5]);
assert_eq!(
set.range(10..=19).rev().collect::<Vec<_>>(),
vec![19, 18, 17, 16, 15, 14, 13, 12, 11, 10]
);
assert_eq!(
set.range(10..20).rev().collect::<Vec<_>>(),
vec![19, 18, 17, 16, 15, 14, 13, 12, 11, 10]
);
assert_eq!(set.range(3..=6).rev().collect::<Vec<_>>(), vec![6, 5, 4, 3]);
assert_eq!(set.range(3..7).rev().collect::<Vec<_>>(), vec![6, 5, 4, 3]);
assert_eq!(set.range(8..=11).rev().collect::<Vec<_>>(), vec![11, 10, 9, 8]);
assert_eq!(set.range(8..12).rev().collect::<Vec<_>>(), vec![11, 10, 9, 8]);
assert!(set.range(20..).collect::<Vec<_>>().is_empty());
assert!(set.range(..0).collect::<Vec<_>>().is_empty());
assert!(set.range(20..).rev().collect::<Vec<_>>().is_empty());
assert!(set.range(..0).rev().collect::<Vec<_>>().is_empty());
}
#[test]
fn concurrent_range_constructions_at_node_boundaries_do_not_deadlock() {
const THREAD_ITERATIONS: usize = 20_000;
let set = Arc::new(BTreeSet::<u64>::with_maximum_node_size(4));
for value in 0..64 {
set.insert(value);
}
let (done_tx, done_rx) = mpsc::channel();
let forward = {
let set = Arc::clone(&set);
let done_tx = done_tx.clone();
thread::spawn(move || {
for iteration in 0..THREAD_ITERATIONS {
let start = (iteration % 64) as u64;
assert_eq!(set.range(start..).next(), Some(start));
}
done_tx.send(()).unwrap();
})
};
let backward = {
let set = Arc::clone(&set);
let done_tx = done_tx.clone();
thread::spawn(move || {
for iteration in 0..THREAD_ITERATIONS {
let end = (iteration % 64) as u64;
assert_eq!(set.range(..=end).next_back(), Some(end));
}
done_tx.send(()).unwrap();
})
};
drop(done_tx);
for _ in 0..2 {
done_rx
.recv_timeout(Duration::from_secs(30))
.expect("concurrent range constructions deadlocked");
}
forward.join().unwrap();
backward.join().unwrap();
}
fn three_node_set() -> BTreeSet<u64> {
let set = BTreeSet::<u64>::with_maximum_node_size(4);
for value in [0u64, 10, 20, 30, 40, 50, 60] {
set.insert(value);
}
assert_eq!(
set.index.read().keys().copied().collect::<Vec<_>>(),
vec![10, 30, 60],
"fixture geometry changed"
);
set
}
#[test]
fn forward_scan_repositions_when_current_node_vanishes() {
let set = three_node_set();
let mut iter = set.iter();
assert_eq!(iter.next(), Some(0));
assert_eq!(iter.next(), Some(10));
assert_eq!(iter.next(), Some(20));
set.index.write().remove(&30).expect("fixture entry");
assert_eq!(iter.next(), Some(30));
assert_eq!(iter.next(), Some(40));
assert_eq!(iter.next(), Some(50));
assert_eq!(iter.next(), Some(60));
assert_eq!(iter.next(), None);
}
#[test]
fn backward_scan_repositions_when_current_node_vanishes() {
let set = three_node_set();
let mut iter = set.iter();
assert_eq!(iter.next_back(), Some(60));
assert_eq!(iter.next_back(), Some(50));
set.index.write().remove(&60).expect("fixture entry");
assert_eq!(iter.next_back(), Some(40));
assert_eq!(iter.next_back(), Some(30));
assert_eq!(iter.next_back(), Some(20));
assert_eq!(iter.next_back(), Some(10));
assert_eq!(iter.next_back(), Some(0));
assert_eq!(iter.next_back(), None);
}
#[test]
fn forward_scan_does_not_re_yield_after_split_of_finished_node() {
let set = BTreeSet::<u64>::with_maximum_node_size(4);
for value in [0u64, 10, 20, 30, 40] {
set.insert(value);
}
let iter = Iter {
tree: &set,
current_front_batch: None,
current_back_batch: None,
exhausted_front_node: Some(set.index.read().first_key_value().expect("fixture node").1.clone()),
exhausted_back_node: None,
front_partial: None,
back_partial: None,
front_batch_limit: INITIAL_BATCH,
back_batch_limit: INITIAL_BATCH,
current_front_value: Some(20),
current_back_value: None,
met: false,
};
assert_eq!(iter.collect::<Vec<_>>(), vec![30, 40]);
}
#[test]
fn backward_scan_does_not_re_yield_values_from_scanned_range() {
let set = BTreeSet::<u64>::new();
set.attach_node(vec![0u64, 10]);
set.attach_node(vec![30u64, 40]);
{
let node = set.index.read().first_key_value().expect("fixture node").1.clone();
let mut guard = node.write();
NodeLike::insert(&mut *guard, 35u64);
}
let mut iter = Iter {
tree: &set,
current_front_batch: None,
current_back_batch: None,
exhausted_front_node: None,
exhausted_back_node: Some(set.index.read().last_key_value().expect("fixture node").1.clone()),
front_partial: None,
back_partial: None,
front_batch_limit: INITIAL_BATCH,
back_batch_limit: INITIAL_BATCH,
current_front_value: None,
current_back_value: Some(30),
met: false,
};
let mut collected = vec![];
while let Some(value) = iter.next_back() {
collected.push(value);
}
assert_eq!(collected, vec![10, 0]);
}
#[test]
fn backward_scan_does_not_skip_values_split_away_after_positioning() {
let set = BTreeSet::<u64>::with_maximum_node_size(4);
for value in [0u64, 10, 20, 30] {
set.insert(value);
}
let mut iter = set.iter();
set.insert(40);
assert!(set.node_count() > 1, "fixture must split");
let mut collected = vec![];
while let Some(value) = iter.next_back() {
collected.push(value);
}
for baseline in [30u64, 20, 10, 0] {
assert!(
collected.iter().any(|value| *value == baseline),
"baseline value {baseline} skipped by backward scan (yielded: {collected:?})"
);
}
assert!(
collected.windows(2).all(|pair| pair[0] > pair[1]),
"backward scan not strictly decreasing: {collected:?}"
);
}
#[test]
fn bidirectional_meet_into_opposite_held_node_does_not_self_deadlock() {
let set = Arc::new(BTreeSet::<u64>::with_maximum_node_size(4));
for value in [0u64, 10, 20, 30, 40] {
set.insert(value);
}
let (done_tx, done_rx) = mpsc::channel();
let handle = {
let set = Arc::clone(&set);
thread::spawn(move || {
let mut finished = set.iter();
assert_eq!(finished.next_back(), Some(40));
assert_eq!(finished.next(), Some(0));
assert_eq!(finished.next(), Some(10));
assert_eq!(finished.next(), Some(20));
assert_eq!(finished.next(), Some(30));
assert_eq!(finished.next(), None);
assert_eq!(finished.next_back(), None);
let mut iter = set.iter();
assert_eq!(iter.next(), Some(0));
assert_eq!(iter.next_back(), Some(40));
assert_eq!(iter.next_back(), Some(30));
assert_eq!(iter.next_back(), Some(20));
assert_eq!(iter.next_back(), Some(10));
assert_eq!(iter.next_back(), None);
assert_eq!(iter.next(), None);
drop(finished);
done_tx.send(()).unwrap();
})
};
done_rx
.recv_timeout(Duration::from_secs(10))
.expect("bidirectional meet-in-the-middle deadlocked");
handle.join().unwrap();
}
#[test]
fn structural_commits_complete_against_paused_scan() {
let set = Arc::new(BTreeSet::<u64>::with_maximum_node_size(2));
for value in 0..8 {
set.insert(value);
}
let scan_holds_guard = Arc::new(Barrier::new(3));
let (done_tx, done_rx) = mpsc::channel();
let scanner = {
let set = Arc::clone(&set);
let scan_holds_guard = Arc::clone(&scan_holds_guard);
let done_tx = done_tx.clone();
thread::spawn(move || {
let mut iter = set.iter();
assert_eq!(iter.next(), Some(0));
scan_holds_guard.wait();
thread::sleep(Duration::from_millis(100));
let mut collected = vec![];
while let Some(value) = iter.next_back() {
collected.push(value);
}
assert_eq!(collected, vec![7, 6, 5, 4, 3, 2, 1]);
done_tx.send(()).unwrap();
})
};
let writer = {
let set = Arc::clone(&set);
let scan_holds_guard = Arc::clone(&scan_holds_guard);
let done_tx = done_tx.clone();
thread::spawn(move || {
scan_holds_guard.wait();
set.remove_range(0..=0);
done_tx.send(()).unwrap();
})
};
drop(done_tx);
scan_holds_guard.wait();
for _ in 0..2 {
done_rx
.recv_timeout(Duration::from_secs(10))
.expect("scan or structural commit deadlocked");
}
scanner.join().unwrap();
writer.join().unwrap();
assert!(!set.contains(&0));
}
#[test]
fn full_scans_do_not_degrade_quadratically_with_node_count() {
use std::time::Instant;
const VALUES: u64 = 30_000;
let set = BTreeSet::<u64>::with_maximum_node_size(2);
for value in 0..VALUES {
set.insert(value);
}
assert!(
set.node_count() >= (VALUES / 4) as usize,
"fixture must be a many-node tree, got {} nodes",
set.node_count()
);
let started = Instant::now();
assert_eq!(set.iter().count(), VALUES as usize);
let forward = started.elapsed();
let started = Instant::now();
assert_eq!(set.iter().rev().count(), VALUES as usize);
let backward = started.elapsed();
let budget = Duration::from_secs(10);
assert!(
forward < budget,
"forward scan took {forward:?}, node advance is not logarithmic"
);
assert!(
backward < budget,
"backward scan took {backward:?}, node advance is not logarithmic"
);
}
#[test]
fn scans_stay_sorted_and_complete_under_monotonic_insert_churn() {
use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering};
const BASELINE: u64 = 400;
const EXTRA: u64 = 2_000;
const SCAN_BOUND: usize = 10_000;
let set = Arc::new(BTreeSet::<u64>::with_maximum_node_size(8));
for value in 0..BASELINE {
set.insert(value);
}
let done = Arc::new(AtomicBool::new(false));
let writer = {
let set = Arc::clone(&set);
let done = Arc::clone(&done);
thread::spawn(move || {
for value in BASELINE..BASELINE + EXTRA {
assert!(set.insert(value));
}
done.store(true, AtomicOrdering::Release);
})
};
let mut scans = 0usize;
loop {
let forward = set.iter().collect::<Vec<_>>();
assert!(
forward.windows(2).all(|pair| pair[0] < pair[1]),
"forward scan not strictly increasing (duplicate or unordered yield)"
);
assert_eq!(
forward.iter().filter(|value| **value < BASELINE).count() as u64,
BASELINE,
"forward scan truncated: baseline keys missing"
);
let backward = set.iter().rev().collect::<Vec<_>>();
assert!(
backward.windows(2).all(|pair| pair[0] > pair[1]),
"backward scan not strictly decreasing (duplicate or unordered yield)"
);
assert_eq!(
backward.iter().filter(|value| **value < BASELINE).count() as u64,
BASELINE,
"backward scan truncated: baseline keys missing"
);
scans += 1;
if done.load(AtomicOrdering::Acquire) || scans >= SCAN_BOUND {
break;
}
}
writer.join().unwrap();
let expected = (0..BASELINE + EXTRA).collect::<Vec<_>>();
assert_eq!(set.iter().collect::<Vec<_>>(), expected);
assert_eq!(set.iter().rev().collect::<Vec<_>>(), {
let mut reversed = expected;
reversed.reverse();
reversed
});
}
#[test]
fn collected_owned_values_survive_arbitrary_concurrent_mutation() {
use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering};
const BASELINE: u64 = 512;
const CHURN: u64 = 4_000;
let set = Arc::new(BTreeSet::<u64>::with_maximum_node_size(8));
for value in 0..BASELINE {
set.insert(value);
}
let done = Arc::new(AtomicBool::new(false));
let writer = {
let set = Arc::clone(&set);
let done = Arc::clone(&done);
thread::spawn(move || {
for value in BASELINE..BASELINE + CHURN {
assert!(set.insert(value));
assert_eq!(set.remove(&value), Some(value));
}
done.store(true, AtomicOrdering::Release);
})
};
let mut snapshots: Vec<Vec<u64>> = Vec::new();
loop {
let snapshot: Vec<u64> = set.iter().collect();
snapshots.push(snapshot);
if done.load(AtomicOrdering::Acquire) {
break;
}
}
writer.join().unwrap();
set.remove_range(..);
assert!(set.is_empty());
for snapshot in snapshots {
assert!(
snapshot.windows(2).all(|pair| pair[0] < pair[1]),
"snapshot not strictly increasing"
);
assert_eq!(
snapshot.iter().filter(|value| **value < BASELINE).count() as u64,
BASELINE,
"snapshot lost baseline keys"
);
}
}
#[test]
fn parallel_iter_and_mut() {
let set = Arc::new(BTreeSet::<i32>::new());
for i in 0..10_000 {
set.insert(i);
}
let set_clone = Arc::clone(&set);
let handle = thread::spawn(move || {
for _ in 0..1000 {
let mut _sum = 0;
for value in set_clone.iter() {
_sum += value;
}
}
});
for i in 10_000..20_000 {
set.insert(i);
}
handle.join().unwrap();
}
#[test]
fn a_scan_across_several_batch_installs_is_complete_and_ordered() {
let set: BTreeSet<u64> = BTreeSet::new();
let count = (INITIAL_BATCH * 20) as u64;
for i in 0..count {
set.insert(i);
}
let seen: Vec<u64> = set.iter().collect();
let expected: Vec<u64> = (0..count).collect();
assert_eq!(seen, expected, "a partial-batch scan lost or repeated elements");
}
#[test]
fn a_backward_scan_across_several_installs_is_complete_and_ordered() {
let set: BTreeSet<u64> = BTreeSet::new();
let count = (INITIAL_BATCH * 20) as u64;
for i in 0..count {
set.insert(i);
}
let seen: Vec<u64> = set.iter().rev().collect();
let expected: Vec<u64> = (0..count).rev().collect();
assert_eq!(
seen, expected,
"a partial-batch backward scan lost or repeated elements"
);
}
#[test]
fn a_scan_within_a_single_node_resumes_correctly() {
let set: BTreeSet<u64> = BTreeSet::with_maximum_node_size(DEFAULT_INNER_SIZE);
let count = 200u64;
for i in 0..count {
set.insert(i);
}
assert_eq!(set.node_count(), 1, "fixture wants one node");
let seen: Vec<u64> = set.iter().collect();
assert_eq!(seen, (0..count).collect::<Vec<_>>());
}
#[test]
fn a_single_element_range_yields_one_element() {
let set: BTreeSet<u64> = BTreeSet::new();
for i in 0..1_000u64 {
set.insert(i);
}
for probe in [0u64, 1, 499, 998, 999] {
let got: Vec<u64> = set.range(probe..=probe).collect();
assert_eq!(got, vec![probe], "range({probe}..={probe})");
}
assert!(set.range(1_000..=1_000).next().is_none(), "absent key");
}
#[test]
fn ranges_spanning_batch_boundaries_are_exact() {
let set: BTreeSet<u64> = BTreeSet::new();
for i in 0..500u64 {
set.insert(i);
}
for width in 1..=(INITIAL_BATCH * 8) as u64 {
let start = 100u64;
let got: Vec<u64> = set.range(start..start + width).collect();
let expected: Vec<u64> = (start..start + width).collect();
assert_eq!(got, expected, "range width {width}");
}
}
#[test]
fn a_double_ended_scan_meets_without_repeating() {
let set: BTreeSet<u64> = BTreeSet::new();
let count = (INITIAL_BATCH * 10) as u64;
for i in 0..count {
set.insert(i);
}
let mut iter = set.iter();
let mut front = Vec::new();
let mut back = Vec::new();
loop {
match iter.next() {
Some(v) => front.push(v),
None => break,
}
match iter.next_back() {
Some(v) => back.push(v),
None => break,
}
}
back.reverse();
front.extend(back);
front.sort_unstable();
assert_eq!(
front,
(0..count).collect::<Vec<_>>(),
"double-ended scan is not a partition"
);
}
#[test]
fn a_scan_under_concurrent_mutation_terminates() {
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
const SIZE: u64 = 4_000;
const CAP: usize = (SIZE * 20) as usize;
for _ in 0..8 {
let set: Arc<BTreeSet<u64>> = Arc::new(BTreeSet::new());
for i in 0..SIZE {
set.insert(i);
}
let stop = Arc::new(AtomicBool::new(false));
let writers: Vec<_> = (0..3)
.map(|w| {
let (set, stop) = (Arc::clone(&set), Arc::clone(&stop));
std::thread::spawn(move || {
let mut i = SIZE + w * 100_000;
while !stop.load(Ordering::Relaxed) {
set.insert(i);
set.remove(&i);
i += 1;
}
})
})
.collect();
let mut yielded = 0usize;
for _ in set.iter() {
yielded += 1;
if yielded >= CAP {
break;
}
}
stop.store(true, Ordering::Relaxed);
for w in writers {
w.join().expect("writer did not panic");
}
assert!(
yielded < CAP,
"scan did not make progress under concurrent mutation: {yielded} yields"
);
}
}
#[allow(clippy::manual_contains)]
fn yielded(seen: &[u64], value: u64) -> bool {
seen.iter().any(|item| *item == value)
}
#[test]
fn deleting_a_yielded_element_does_not_skip_a_live_one() {
for prefix in 1..12usize {
let set: BTreeSet<u64> = BTreeSet::new();
for i in 0..256u64 {
set.insert(i);
}
let mut seen = Vec::new();
for value in set.iter() {
seen.push(value);
if seen.len() == prefix {
set.remove(&0);
}
}
for expected in 1..256u64 {
assert!(
yielded(&seen, expected),
"prefix {prefix}: {expected} was present for the whole scan but was never yielded"
);
}
}
}
#[test]
fn deleting_a_yielded_element_backwards_does_not_skip_a_live_one() {
for prefix in 1..12usize {
let set: BTreeSet<u64> = BTreeSet::new();
for i in 0..256u64 {
set.insert(i);
}
let mut seen = Vec::new();
for value in set.iter().rev() {
seen.push(value);
if seen.len() == prefix {
set.remove(&255);
}
}
for expected in 0..255u64 {
assert!(
yielded(&seen, expected),
"prefix {prefix}: {expected} was present for the whole scan but was never yielded"
);
}
}
}
}