Skip to content

Commit

Permalink
Fix Stacked Borrow violations in BptreeMap (#125)
Browse files Browse the repository at this point in the history
  • Loading branch information
preston-evans98 authored Oct 2, 2024
1 parent 3a12b6d commit b32cb4c
Show file tree
Hide file tree
Showing 5 changed files with 267 additions and 198 deletions.
36 changes: 18 additions & 18 deletions src/internals/bptree/cursor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,8 +66,8 @@ impl<K: Clone + Ord + Debug, V: Clone> LinCowCellCapable<CursorRead<K, V>, Curso
// Now when the lock is dropped, both sides see the correct info and garbage for drops.

// We are done, time to seal everything.
new.first_seen.iter().for_each(|n| unsafe {
(**n).make_ro();
new.first_seen.iter().for_each(|n| {
Node::make_ro_raw(*n);
});
// Clear first seen, we won't be dropping them from here.
new.first_seen.clear();
Expand Down Expand Up @@ -105,19 +105,19 @@ impl<K: Clone + Ord + Debug, V: Clone> SuperBlock<K, V> {
// let last_seen: Vec<*mut Node<K, V>> = Vec::with_capacity(16);
let mut first_seen = Vec::with_capacity(16);
// Do a pre-verify to be sure it's sane.
assert!(unsafe { (*root).verify() });
assert!(Node::verify_raw(root));
// Collect anythinng from root into this txid if needed.
// Set txid to txid on all tree nodes from the root.
first_seen.push(root);
unsafe { (*root).sblock_collect(&mut first_seen) };
Node::sblock_collect_raw(root, &mut first_seen);

// Lock them all
first_seen.iter().for_each(|n| unsafe {
(**n).make_ro();
first_seen.iter().for_each(|n| {
Node::make_ro_raw(*n);
});

// Determine our count internally.
let (length, _) = unsafe { (*root).tree_density() };
let (length, _) = Node::tree_density_raw(root);

// Good to go!
SuperBlock {
Expand Down Expand Up @@ -184,8 +184,8 @@ pub(crate) trait CursorReadOps<K: Clone + Ord + Debug, V: Clone> {
#[cfg(test)]
fn get_tree_density(&self) -> (usize, usize) {
// Walk the tree and calculate the packing efficiency.
let rref = self.get_root_ref();
rref.tree_density()
let rref = self.get_root();
Node::tree_density_raw(rref)
}

fn search<Q>(&self, k: &Q) -> Option<&V>
Expand Down Expand Up @@ -250,7 +250,7 @@ pub(crate) trait CursorReadOps<K: Clone + Ord + Debug, V: Clone> {
panic!("Tree depth exceeded max limit (65536). This may indicate memory corruption.");
}

fn range<'n, R, T>(&'n self, range: R) -> RangeIter<'n, '_, K, V>
fn range<'n, R, T>(&'n self, range: R) -> RangeIter<'n, 'n, K, V>
where
K: Borrow<T>,
T: Ord + ?Sized,
Expand All @@ -259,21 +259,21 @@ pub(crate) trait CursorReadOps<K: Clone + Ord + Debug, V: Clone> {
RangeIter::new(self.get_root(), range, self.len())
}

fn kv_iter<'n>(&'n self) -> Iter<'n, '_, K, V> {
fn kv_iter<'n>(&'n self) -> Iter<'n, 'n, K, V> {
Iter::new(self.get_root(), self.len())
}

fn k_iter<'n>(&'n self) -> KeyIter<'n, '_, K, V> {
fn k_iter<'n>(&'n self) -> KeyIter<'n, 'n, K, V> {
KeyIter::new(self.get_root(), self.len())
}

fn v_iter<'n>(&'n self) -> ValueIter<'n, '_, K, V> {
fn v_iter<'n>(&'n self) -> ValueIter<'n, 'n, K, V> {
ValueIter::new(self.get_root(), self.len())
}

#[cfg(test)]
fn verify(&self) -> bool {
self.get_root_ref().no_cycles() && self.get_root_ref().verify() && {
Node::no_cycles_raw(self.get_root()) && Node::verify_raw(self.get_root()) && {
let (l, _) = self.get_tree_density();
l == self.len()
}
Expand Down Expand Up @@ -554,10 +554,10 @@ impl<K: Clone + Ord + Debug, V: Clone> CursorWrite<K, V> {

#[cfg(test)]
pub(crate) fn tree_density(&self) -> (usize, usize) {
self.get_root_ref().tree_density()
Node::<K, V>::tree_density_raw(self.get_root())
}

pub(crate) fn range_mut<'n, R, T>(&'n mut self, range: R) -> RangeMutIter<'n, '_, K, V>
pub(crate) fn range_mut<'n, R, T>(&'n mut self, range: R) -> RangeMutIter<'n, 'n, K, V>
where
K: Borrow<T>,
T: Ord + ?Sized,
Expand Down Expand Up @@ -604,7 +604,7 @@ impl<K: Clone + Ord + Debug, V: Clone> Drop for SuperBlock<K, V> {
let mut first_seen = Vec::with_capacity(16);
// eprintln!("{:?}", self.root);
first_seen.push(self.root);
unsafe { (*self.root).sblock_collect(&mut first_seen) };
Node::sblock_collect_raw(self.root, &mut first_seen);
first_seen.iter().for_each(|n| Node::free(*n));
}
}
Expand Down Expand Up @@ -1096,7 +1096,7 @@ where
K: Clone + Ord + Debug + 'a,
V: Clone,
{
if self_meta!(node).is_leaf() {
if unsafe {&* node}.meta.is_leaf() {
leaf_ref!(node, K, V).get_mut_ref(k)
} else {
// This nmref binds the life of the reference ...
Expand Down
13 changes: 7 additions & 6 deletions src/internals/bptree/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ impl<'a, K: Clone + Ord + Debug, V: Clone> Iterator for LeafIter<'a, K, V> {

// Return the leaf as we found at the start, regardless of the
// stack operations.
Some(leaf_ref!(leafref, K, V))
Some(leaf_ref_shared!(leafref, K, V))
}

fn size_hint(&self) -> (usize, Option<usize>) {
Expand Down Expand Up @@ -181,7 +181,8 @@ where
}
break;
} else {
let bref = branch_ref!(work_node, K, V);
let bref = branch_ref_shared!(work_node, K, V);
let bref_count = bref.count();
match bound {
Bound::Excluded(q) | Bound::Included(q) => {
let idx = bref.locate_node(q);
Expand All @@ -192,8 +193,8 @@ where
}
Bound::Unbounded => {
// count shows the most right node.
stack.push_back((work_node, bref.count()));
work_node = branch_ref!(work_node, K, V).get_idx_unchecked(bref.count());
stack.push_back((work_node, bref_count));
work_node = branch_ref!(work_node, K, V).get_idx_unchecked(bref_count);
}
}
}
Expand Down Expand Up @@ -297,7 +298,7 @@ impl<'a, K: Clone + Ord + Debug, V: Clone> Iterator for RevLeafIter<'a, K, V> {

// Return the leaf as we found at the start, regardless of the
// stack operations.
Some(leaf_ref!(leafref, K, V))
Some(leaf_ref_shared!(leafref, K, V))
}

fn size_hint(&self) -> (usize, Option<usize>) {
Expand Down Expand Up @@ -671,7 +672,7 @@ impl<K: Clone + Ord + Debug, V: Clone> DoubleEndedIterator for RangeIter<'_, '_,
fn next_back(&mut self) -> Option<Self::Item> {
loop {
if let Some((node, idx)) = self.right_iter.get_mut() {
let leaf = leaf_ref!(*node, K, V);
let leaf = leaf_ref_shared!(*node, K, V);
// Get idx checked.
if let Some(r) = leaf.get_kv_idx_checked(*idx) {
if let Some((lnode, lidx)) = self.left_iter.get_mut() {
Expand Down
18 changes: 18 additions & 0 deletions src/internals/bptree/macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,24 @@ macro_rules! branch_ref {
}};
}

/// Like [`branch_ref`], but yields &Leaf without coercing from &mut Leaf. This is useful
/// to avoid triggering Miri's analysis.
macro_rules! branch_ref_shared {
($x:expr, $k:ty, $v:ty) => {{
debug_assert!(unsafe { (*$x).meta.is_branch() });
unsafe { &*($x as *const Branch<$k, $v>) }
}};
}

/// Like [`leaf_ref`], but yields &Leaf without coercing from &mut Leaf. This is useful
/// to avoid triggering Miri's analysis.
macro_rules! leaf_ref_shared {
($x:expr, $k:ty, $v:ty) => {{
debug_assert!(unsafe { (*$x).meta.is_leaf() });
unsafe { &*($x as *const Leaf<$k, $v>) }
}};
}

macro_rules! leaf_ref {
($x:expr, $k:ty, $v:ty) => {{
debug_assert!(unsafe { (*$x).meta.is_leaf() });
Expand Down
Loading

0 comments on commit b32cb4c

Please sign in to comment.