use crate::AccessGuard;
use crate::sync::Mutex;
use crate::tree_store::btree_base::{
BRANCH, BranchAccessor, LEAF, LeafAccessor, MAX_BTREE_DEPTH, OwnedEntryBuffer,
leaf_below_merge_threshold, leaf_fits_one_page, retained_after_removals,
};
use crate::tree_store::btree_iters::EntryGuard;
use crate::tree_store::btree_mutator::MutateHelper;
use crate::tree_store::page_store::{Page, PageHint, PageImpl};
use crate::tree_store::{BtreeHeader, PageAllocator, PageNumber, PageResolver, PageTracker};
use crate::types::{Key, Value};
use crate::{Result, StorageError};
#[cfg(feature = "experimental_cursor")]
use alloc::boxed::Box;
use alloc::string::ToString;
use alloc::sync::Arc;
use alloc::vec;
use alloc::vec::Vec;
use core::cmp::Ordering;
use core::marker::PhantomData;
use core::ops::Bound;
use core::ops::Bound::{Excluded, Included, Unbounded};
use core::ops::Range;
#[derive(Clone)]
struct Branch {
page: PageImpl,
child_index: usize,
}
impl Branch {
fn new(page: PageImpl, child_index: usize) -> Self {
Self { page, child_index }
}
fn into_parts(self) -> (PageImpl, usize) {
(self.page, self.child_index)
}
fn adjacent_child(
&self,
direction: Direction,
fixed_key_width: Option<usize>,
) -> Option<(usize, PageNumber)> {
let accessor = BranchAccessor::new(&self.page, fixed_key_width);
let child_index = match direction {
Direction::Next => {
let next = self.child_index + 1;
(next < accessor.count_children()).then_some(next)?
}
Direction::Previous => self.child_index.checked_sub(1)?,
};
Some((child_index, accessor.child_page(child_index).unwrap()))
}
}
#[derive(Copy, Clone)]
pub(super) enum Position<'a> {
Start,
End,
Before(&'a [u8]),
After(&'a [u8]),
}
impl<'a> Position<'a> {
pub(super) fn from_lower_bound(bound: Bound<&'a [u8]>) -> Self {
match bound {
Included(key) => Self::Before(key),
Excluded(key) => Self::After(key),
Unbounded => Self::Start,
}
}
pub(super) fn from_upper_bound(bound: Bound<&'a [u8]>) -> Self {
match bound {
Included(key) => Self::After(key),
Excluded(key) => Self::Before(key),
Unbounded => Self::End,
}
}
}
#[derive(Copy, Clone, PartialEq)]
enum Direction {
Next,
Previous,
}
impl Direction {
fn is_next(self) -> bool {
matches!(self, Self::Next)
}
fn opposite(self) -> Self {
match self {
Self::Next => Self::Previous,
Self::Previous => Self::Next,
}
}
}
fn lower_bound_entry<K: Key>(accessor: &LeafAccessor<'_>, position: Position<'_>) -> usize {
match position {
Position::Start => 0,
Position::End => accessor.num_pairs(),
Position::Before(query) | Position::After(query) => {
let (mut position_index, found) = accessor.position::<K>(query);
if matches!(position, Position::After(_)) && found {
position_index += 1;
}
position_index
}
}
}
fn child_to_visit<K: Key>(
accessor: &BranchAccessor<'_, '_, PageImpl>,
position: Position<'_>,
) -> usize {
match position {
Position::Start => 0,
Position::End => accessor.count_children() - 1,
Position::Before(query) | Position::After(query) => accessor.child_for_key::<K>(query).0,
}
}
fn descend_to_position<K: Key + 'static, V: Value + 'static, F>(
page: PageImpl,
position: Position<'_>,
path: &mut Vec<Branch>,
get_page: &mut F,
) -> Result<Leaf>
where
F: FnMut(PageNumber) -> Result<PageImpl>,
{
let mut page = page;
loop {
let (child_index, child_page) = match page.memory()[0] {
LEAF => {
let (leaf_position, len) = {
let accessor =
LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width());
(
lower_bound_entry::<K>(&accessor, position),
accessor.num_pairs(),
)
};
return Ok(Leaf {
page,
position: leaf_position,
len,
});
}
BRANCH => {
let accessor = BranchAccessor::new(&page, K::fixed_width());
let child_index = child_to_visit::<K>(&accessor, position);
(child_index, accessor.child_page(child_index).unwrap())
}
_ => unreachable!(),
};
if path.len() >= MAX_BTREE_DEPTH {
return Err(StorageError::Corrupted(
"Btree exceeded maximum depth".to_string(),
));
}
path.push(Branch::new(page, child_index));
page = get_page(child_page)?;
}
}
fn move_to_adjacent_leaf<K: Key + 'static, V: Value + 'static, F>(
path: &mut Vec<Branch>,
direction: Direction,
get_page: &mut F,
) -> Result<Option<Leaf>>
where
F: FnMut(PageNumber) -> Result<PageImpl>,
{
for index in (0..path.len()).rev() {
if let Some((child_index, child_page)) =
path[index].adjacent_child(direction, K::fixed_width())
{
path[index].child_index = child_index;
path.truncate(index + 1);
let page = get_page(child_page)?;
let edge = if direction.is_next() {
Position::Start
} else {
Position::End
};
return descend_to_position::<K, V, F>(page, edge, path, get_page).map(Some);
}
}
Ok(None)
}
fn prepare_leaf<K: Key + 'static, V: Value + 'static, F>(
leaf: &mut Option<Leaf>,
path: &mut Vec<Branch>,
direction: Direction,
get_page: &mut F,
) -> Result<bool>
where
F: FnMut(PageNumber) -> Result<PageImpl>,
{
while let Some(current) = leaf.as_ref() {
if (direction.is_next() && current.position < current.len)
|| (!direction.is_next() && current.position > 0)
{
return Ok(true);
}
let Some(next_leaf) = move_to_adjacent_leaf::<K, V, F>(path, direction, get_page)? else {
return Ok(false);
};
*leaf = Some(next_leaf);
}
Ok(false)
}
fn entry<K: Key + 'static, V: Value + 'static>(leaf: &Leaf, position: usize) -> EntryGuard<K, V> {
let (key, value) = LeafAccessor::new(leaf.page.memory(), K::fixed_width(), V::fixed_width())
.entry_ranges(position)
.expect("cursor entry must exist");
EntryGuard::new(leaf.page.clone(), key, value)
}
fn entry_ref<K: Key + 'static, V: Value + 'static>(
leaf: &Leaf,
position: usize,
) -> EntryRef<'_, K, V> {
let (key_range, value_range) =
LeafAccessor::new(leaf.page.memory(), K::fixed_width(), V::fixed_width())
.entry_ranges(position)
.expect("cursor entry must exist");
EntryRef {
page: &leaf.page,
key_range,
value_range,
_key_type: PhantomData,
_value_type: PhantomData,
}
}
fn key_data<K: Key + 'static, V: Value + 'static>(leaf: &Leaf, position: usize) -> Vec<u8> {
LeafAccessor::new(leaf.page.memory(), K::fixed_width(), V::fixed_width())
.entry(position)
.expect("cursor entry must exist")
.key()
.to_vec()
}
fn scan_boundary_key<K: Key + 'static, V: Value + 'static>(
leaf: &Leaf,
direction: Direction,
) -> Vec<u8> {
match direction {
Direction::Next => key_data::<K, V>(leaf, leaf.len - 1),
Direction::Previous => key_data::<K, V>(leaf, 0),
}
}
#[derive(Clone)]
struct Leaf {
page: PageImpl,
position: usize,
len: usize,
}
struct LeafRunRewrite {
parent_page: PageNumber,
direction: Direction,
replaced_children: Range<usize>,
entries: OwnedEntryBuffer,
removed_pairs: u64,
}
impl LeafRunRewrite {
fn new(parent_page: PageNumber, child_index: usize, direction: Direction) -> Self {
let origin = match direction {
Direction::Next => child_index,
Direction::Previous => child_index + 1,
};
let replaced_children = origin..origin;
Self {
parent_page,
direction,
replaced_children,
entries: OwnedEntryBuffer::default(),
removed_pairs: 0,
}
}
fn append_entries_from<K: Key, V: Value>(
&mut self,
page: PageImpl,
child_index: usize,
removed_indexes: &[usize],
) {
debug_assert!(removed_indexes.windows(2).all(|pair| pair[0] < pair[1]));
match self.direction {
Direction::Next => {
assert_eq!(child_index, self.replaced_children.end);
self.replaced_children.end += 1;
}
Direction::Previous => {
assert_eq!(child_index + 1, self.replaced_children.start);
self.replaced_children.start -= 1;
}
}
self.removed_pairs += removed_indexes.len() as u64;
let accessor = LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width());
self.entries
.extend_from_leaf(&accessor, removed_indexes, self.direction.is_next());
}
}
#[cfg(feature = "experimental_cursor")]
const INSERT_FLUSH_BYTES: usize = 1024 * 1024;
#[cfg(feature = "experimental_cursor")]
#[derive(Copy, Clone, PartialEq)]
enum RunDirection {
Ascending,
Descending,
}
#[cfg(feature = "experimental_cursor")]
struct InsertRun {
direction: RunDirection,
opening_next_key: Option<Vec<u8>>,
previous_key: Option<Vec<u8>>,
entries: OwnedEntryBuffer,
inserted_pairs: u64,
}
#[cfg(feature = "experimental_cursor")]
impl InsertRun {
fn buffered_previous(&self) -> Option<(&[u8], &[u8])> {
match self.direction {
RunDirection::Ascending => self.entries.back(),
RunDirection::Descending => None,
}
}
fn buffered_next(&self) -> Option<(&[u8], &[u8])> {
match self.direction {
RunDirection::Ascending => None,
RunDirection::Descending => self.entries.front(),
}
}
fn next_key(&self) -> Option<&[u8]> {
self.buffered_next()
.map(|(key, _)| key)
.or(self.opening_next_key.as_deref())
}
fn rejects<K: Key>(&self, key: &[u8]) -> bool {
if let Some(previous) = &self.previous_key
&& K::compare(key, previous).is_le()
{
return true;
}
if let Some(next) = self.next_key()
&& K::compare(key, next).is_ge()
{
return true;
}
false
}
}
pub(super) struct EntryRef<'a, K: Key + 'static, V: Value + 'static> {
page: &'a PageImpl,
key_range: Range<usize>,
value_range: Range<usize>,
_key_type: PhantomData<K>,
_value_type: PhantomData<V>,
}
impl<K: Key + 'static, V: Value + 'static> EntryRef<'_, K, V> {
pub(super) fn key_bytes(&self) -> &[u8] {
&self.page.memory()[self.key_range.clone()]
}
#[cfg(feature = "experimental_cursor")]
pub(super) fn to_guards<'g>(&self) -> (AccessGuard<'g, K>, AccessGuard<'g, V>) {
(
AccessGuard::with_page(self.page.clone(), self.key_range.clone()),
AccessGuard::with_page(self.page.clone(), self.value_range.clone()),
)
}
pub(super) fn key(&self) -> K::SelfType<'_> {
K::from_bytes(&self.page.memory()[self.key_range.clone()])
}
pub(super) fn value(&self) -> V::SelfType<'_> {
V::from_bytes(&self.page.memory()[self.value_range.clone()])
}
}
#[derive(Clone)]
pub(super) struct Cursor<K: Key + 'static, V: Value + 'static> {
root: PageNumber,
path: Vec<Branch>,
leaf: Option<Leaf>,
manager: PageResolver,
hint: PageHint,
_key_type: PhantomData<K>,
_value_type: PhantomData<V>,
}
impl<K: Key + 'static, V: Value + 'static> Cursor<K, V> {
pub(super) fn new(root: PageNumber, manager: PageResolver, hint: PageHint) -> Self {
Self {
root,
path: vec![],
leaf: None,
manager,
hint,
_key_type: PhantomData,
_value_type: PhantomData,
}
}
pub(super) fn seek_to(&mut self, position: Position<'_>) -> Result {
self.path.clear();
let root_page = self.manager.get_page(self.root, self.hint)?;
let Self {
manager,
hint,
path,
leaf,
..
} = self;
let mut get_page = |page| manager.get_page(page, *hint);
*leaf = Some(descend_to_position::<K, V, _>(
root_page,
position,
path,
&mut get_page,
)?);
Ok(())
}
fn ensure_has_entry(&mut self, direction: Direction) -> Result<bool> {
let Self {
manager,
hint,
path,
leaf,
..
} = self;
let mut get_page = |page| manager.get_page(page, *hint);
prepare_leaf::<K, V, _>(leaf, path, direction, &mut get_page)
}
pub(super) fn normalize_forward_gap(&mut self) -> Result {
if self
.leaf
.as_ref()
.is_none_or(|leaf| leaf.position != leaf.len)
{
return Ok(());
}
let Self {
manager,
hint,
path,
leaf,
..
} = self;
let mut get_page = |page| manager.get_page(page, *hint);
if let Some(next_leaf) =
move_to_adjacent_leaf::<K, V, _>(path, Direction::Next, &mut get_page)?
{
*leaf = Some(next_leaf);
}
Ok(())
}
pub(super) fn next(&mut self) -> Result<Option<EntryGuard<K, V>>> {
if !self.ensure_has_entry(Direction::Next)? {
return Ok(None);
}
let leaf = self.leaf.as_mut().expect("cursor must be positioned");
let position = leaf.position;
leaf.position += 1;
Ok(Some(entry(leaf, position)))
}
pub(super) fn prev(&mut self) -> Result<Option<EntryGuard<K, V>>> {
if !self.ensure_has_entry(Direction::Previous)? {
return Ok(None);
}
let leaf = self.leaf.as_mut().expect("cursor must be positioned");
leaf.position -= 1;
Ok(Some(entry(leaf, leaf.position)))
}
#[cfg(feature = "experimental_cursor")]
pub(super) fn peek_next(&mut self) -> Result<Option<EntryGuard<K, V>>> {
if !self.ensure_has_entry(Direction::Next)? {
return Ok(None);
}
let leaf = self.leaf.as_ref().expect("cursor must be positioned");
Ok(Some(entry(leaf, leaf.position)))
}
#[cfg(feature = "experimental_cursor")]
pub(super) fn peek_prev(&mut self) -> Result<Option<EntryGuard<K, V>>> {
if !self.ensure_has_entry(Direction::Previous)? {
return Ok(None);
}
let leaf = self.leaf.as_ref().expect("cursor must be positioned");
Ok(Some(entry(leaf, leaf.position - 1)))
}
fn page_number(&self) -> PageNumber {
self.leaf
.as_ref()
.expect("cursor must be positioned")
.page
.get_page_number()
}
fn position(&self) -> usize {
self.leaf
.as_ref()
.expect("cursor must be positioned")
.position
}
pub(super) fn compare_position(&self, other: &Self) -> Ordering {
let self_page = self.page_number();
let other_page = other.page_number();
if self_page == other_page {
return self.position().cmp(&other.position());
}
assert_eq!(self.path.len(), other.path.len());
for (self_frame, other_frame) in self.path.iter().zip(&other.path) {
match self_frame.child_index.cmp(&other_frame.child_index) {
Ordering::Equal => {}
ordering => return ordering,
}
}
unreachable!("distinct cursor pages must diverge in their branch path")
}
}
struct CursorPosition {
path: Vec<Branch>,
leaf: Leaf,
}
#[derive(Default)]
pub(super) struct CursorState {
position: Option<CursorPosition>,
removed_indexes: Vec<usize>,
detached_guards: bool,
poisoned: bool,
leaf_run_rewrite: Option<LeafRunRewrite>,
#[cfg(feature = "experimental_cursor")]
insert_run: Option<Box<InsertRun>>,
}
enum LeafCloseOutcome {
Unchanged,
Flushed { resume_key: Vec<u8> },
AbsorbedIntoRun,
}
pub(super) struct CursorMut<'a, 'b, K: Key + 'static, V: Value + 'static> {
root: &'b mut Option<BtreeHeader>,
page_allocator: &'b PageAllocator,
freed: &'b mut Vec<PageNumber>,
allocated: &'b Arc<PageTracker>,
state: CursorState,
_key_type: PhantomData<K>,
_value_type: PhantomData<V>,
_lifetime: PhantomData<&'a ()>,
}
impl CursorPosition {
fn into_parts(self) -> (Vec<(PageImpl, usize)>, Leaf) {
(
self.path.into_iter().map(Branch::into_parts).collect(),
self.leaf,
)
}
fn has_entry(&self, direction: Direction) -> bool {
match direction {
Direction::Next => self.leaf.position < self.leaf.len,
Direction::Previous => self.leaf.position > 0,
}
}
fn entry_index(&self, direction: Direction) -> usize {
match direction {
Direction::Next => self.leaf.position,
Direction::Previous => self.leaf.position - 1,
}
}
fn move_once(&mut self, direction: Direction) {
match direction {
Direction::Next => self.leaf.position += 1,
Direction::Previous => self.leaf.position -= 1,
}
}
}
impl<'a, 'b, K: Key + 'static, V: Value + 'static> CursorMut<'a, 'b, K, V> {
pub(super) fn new(
root: &'b mut Option<BtreeHeader>,
page_allocator: &'b PageAllocator,
freed: &'b mut Vec<PageNumber>,
allocated: &'b Arc<PageTracker>,
) -> Self {
Self::with_state(
root,
page_allocator,
freed,
allocated,
CursorState::default(),
)
}
pub(super) fn with_state(
root: &'b mut Option<BtreeHeader>,
page_allocator: &'b PageAllocator,
freed: &'b mut Vec<PageNumber>,
allocated: &'b Arc<PageTracker>,
state: CursorState,
) -> Self {
Self {
root,
page_allocator,
freed,
allocated,
state,
_key_type: PhantomData,
_value_type: PhantomData,
_lifetime: PhantomData,
}
}
pub(super) fn into_state(self) -> CursorState {
self.state
}
pub(super) fn seek_to(&mut self, target: Position<'_>) -> Result {
self.check_not_poisoned()?;
assert!(self.state.leaf_run_rewrite.is_none());
assert!(self.state.removed_indexes.is_empty());
#[cfg(feature = "experimental_cursor")]
assert!(self.state.insert_run.is_none());
self.state.position = None;
let Some(header) = *self.root else {
return Ok(());
};
let root_page = self.page_allocator.get_page(header.root, PageHint::None)?;
let page_allocator = self.page_allocator;
let mut get_page = |page| page_allocator.get_page(page, PageHint::None);
let mut path = vec![];
let leaf = descend_to_position::<K, V, _>(root_page, target, &mut path, &mut get_page)?;
self.state.position = Some(CursorPosition { path, leaf });
Ok(())
}
fn check_pending_removals(&self, direction: Direction) {
let Some(position) = self.state.position.as_ref() else {
return;
};
let gap = position.leaf.position;
let valid = match direction {
Direction::Next => self
.state
.removed_indexes
.last()
.is_none_or(|last| *last < gap),
Direction::Previous => self
.state
.removed_indexes
.last()
.is_none_or(|last| *last >= gap),
};
assert!(valid, "pending removals must match the scan direction");
}
fn ensure_has_entry(&mut self, direction: Direction) -> Result<bool> {
self.check_not_poisoned()?;
#[cfg(feature = "experimental_cursor")]
assert!(self.state.insert_run.is_none());
self.check_pending_removals(direction);
loop {
let Some(position) = self.state.position.as_ref() else {
return Ok(false);
};
if position.has_entry(direction) {
return Ok(true);
}
if !self.advance_past_closed_leaf(direction)? {
return Ok(false);
}
}
}
fn advance_past_closed_leaf(&mut self, direction: Direction) -> Result<bool> {
match self.close_current_leaf(direction)? {
LeafCloseOutcome::Unchanged => self.step_to_adjacent_leaf(direction),
LeafCloseOutcome::Flushed { resume_key } => {
self.resume_after_rewrite(direction, &resume_key)?;
Ok(self.state.position.is_some())
}
LeafCloseOutcome::AbsorbedIntoRun => {
let stepped = self
.step_to_adjacent_leaf(direction)
.inspect_err(|_| self.poison())?;
assert!(stepped);
Ok(stepped)
}
}
}
fn run_parent_frame(&self) -> &Branch {
let position = self
.state
.position
.as_ref()
.expect("cursor must be positioned");
position
.path
.last()
.expect("leaf runs require a parent branch")
}
fn run_parent_has_more_children(&self, direction: Direction) -> bool {
self.run_parent_frame()
.adjacent_child(direction, K::fixed_width())
.is_some()
}
fn step_to_adjacent_leaf(&mut self, direction: Direction) -> Result<bool> {
let Some(position) = self.state.position.as_mut() else {
return Ok(false);
};
let page_allocator = self.page_allocator;
let mut get_page = |page| page_allocator.get_page(page, PageHint::None);
if let Some(next_leaf) =
move_to_adjacent_leaf::<K, V, _>(&mut position.path, direction, &mut get_page)?
{
position.leaf = next_leaf;
Ok(true)
} else {
Ok(false)
}
}
pub(super) fn peek_next(&mut self) -> Result<Option<EntryRef<'_, K, V>>> {
if !self.ensure_has_entry(Direction::Next)? {
return Ok(None);
}
let position = self
.state
.position
.as_ref()
.expect("cursor must be positioned");
Ok(Some(entry_ref(
&position.leaf,
position.entry_index(Direction::Next),
)))
}
pub(super) fn peek_prev(&mut self) -> Result<Option<EntryRef<'_, K, V>>> {
if !self.ensure_has_entry(Direction::Previous)? {
return Ok(None);
}
let position = self
.state
.position
.as_ref()
.expect("cursor must be positioned");
Ok(Some(entry_ref(
&position.leaf,
position.entry_index(Direction::Previous),
)))
}
pub(super) fn move_next(&mut self) -> Result<bool> {
if !self.ensure_has_entry(Direction::Next)? {
return Ok(false);
}
self.state
.position
.as_mut()
.expect("cursor must be positioned")
.move_once(Direction::Next);
Ok(true)
}
#[cfg(feature = "experimental_cursor")]
pub(super) fn move_prev(&mut self) -> Result<bool> {
if !self.ensure_has_entry(Direction::Previous)? {
return Ok(false);
}
self.state
.position
.as_mut()
.expect("cursor must be positioned")
.move_once(Direction::Previous);
Ok(true)
}
#[cfg(feature = "experimental_cursor")]
pub(super) fn next(&mut self) -> Result<Option<EntryRef<'_, K, V>>> {
if !self.move_next()? {
return Ok(None);
}
let position = self
.state
.position
.as_ref()
.expect("cursor must be positioned");
Ok(Some(entry_ref(
&position.leaf,
position.entry_index(Direction::Previous),
)))
}
#[cfg(feature = "experimental_cursor")]
pub(super) fn prev(&mut self) -> Result<Option<EntryRef<'_, K, V>>> {
if !self.move_prev()? {
return Ok(None);
}
let position = self
.state
.position
.as_ref()
.expect("cursor must be positioned");
Ok(Some(entry_ref(
&position.leaf,
position.entry_index(Direction::Next),
)))
}
pub(super) fn remove_next(
&mut self,
) -> Result<Option<(AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
assert!(self.state.removed_indexes.is_empty());
if !self.ensure_has_entry(Direction::Next)? {
return Ok(None);
}
let position = self
.state
.position
.take()
.expect("cursor must be positioned");
let index = position.leaf.position;
self.remove_leaf_entry(position.leaf.page, position.path, index)
}
pub(super) fn remove_next_discard(&mut self) -> Result<bool> {
if !self.ensure_has_entry(Direction::Next)? {
return Ok(false);
}
self.record_removal(Direction::Next);
Ok(true)
}
pub(super) fn remove_next_deferred(
&mut self,
) -> Result<Option<(AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
if !self.ensure_has_entry(Direction::Next)? {
return Ok(None);
}
Ok(Some(self.record_removal_deferred(Direction::Next)))
}
pub(super) fn remove_prev(
&mut self,
) -> Result<Option<(AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
assert!(self.state.removed_indexes.is_empty());
if !self.ensure_has_entry(Direction::Previous)? {
return Ok(None);
}
let position = self
.state
.position
.take()
.expect("cursor must be positioned");
let index = position.leaf.position - 1;
self.remove_leaf_entry(position.leaf.page, position.path, index)
}
pub(super) fn remove_prev_deferred(
&mut self,
) -> Result<Option<(AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
if !self.ensure_has_entry(Direction::Previous)? {
return Ok(None);
}
Ok(Some(self.record_removal_deferred(Direction::Previous)))
}
#[cfg(feature = "experimental_cursor")]
#[allow(clippy::type_complexity)]
pub(super) fn remove_next_taking_key(
&mut self,
) -> Result<Option<(Vec<u8>, AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
self.remove_taking_key(Direction::Next)
}
#[cfg(feature = "experimental_cursor")]
#[allow(clippy::type_complexity)]
pub(super) fn remove_prev_taking_key(
&mut self,
) -> Result<Option<(Vec<u8>, AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
self.remove_taking_key(Direction::Previous)
}
#[cfg(feature = "experimental_cursor")]
#[allow(clippy::type_complexity)]
fn remove_taking_key(
&mut self,
direction: Direction,
) -> Result<Option<(Vec<u8>, AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
assert!(self.state.removed_indexes.is_empty());
if !self.ensure_has_entry(direction)? {
return Ok(None);
}
let position = self
.state
.position
.take()
.expect("cursor must be positioned");
let index = position.entry_index(direction);
let key = key_data::<K, V>(&position.leaf, index);
let entry = self.remove_leaf_entry(position.leaf.page, position.path, index)?;
Ok(entry.map(|(key_guard, value_guard)| (key, key_guard, value_guard)))
}
fn record_removal(&mut self, direction: Direction) -> usize {
let position = self
.state
.position
.as_mut()
.expect("cursor must be positioned");
let index = position.entry_index(direction);
let monotonic = match direction {
Direction::Next => self
.state
.removed_indexes
.last()
.is_none_or(|last| *last < index),
Direction::Previous => self
.state
.removed_indexes
.last()
.is_none_or(|last| *last > index),
};
assert!(
monotonic,
"removed indexes must be recorded monotonically in the scan direction"
);
self.state.removed_indexes.push(index);
position.move_once(direction);
index
}
fn record_removal_deferred(
&mut self,
direction: Direction,
) -> (AccessGuard<'a, K>, AccessGuard<'a, V>) {
let index = self.record_removal(direction);
self.state.detached_guards = true;
let leaf = &self
.state
.position
.as_ref()
.expect("cursor must be positioned")
.leaf;
let (key_range, value_range) =
LeafAccessor::new(leaf.page.memory(), K::fixed_width(), V::fixed_width())
.entry_ranges(index)
.expect("removed cursor entry must exist");
let page = leaf.page.to_arc();
(
AccessGuard::with_arc_page(page.clone(), key_range),
AccessGuard::with_arc_page(page, value_range),
)
}
pub(super) fn finish_pending_removals(&mut self) -> Result {
self.check_not_poisoned()?;
let direction = self
.state
.leaf_run_rewrite
.as_ref()
.map_or(Direction::Next, |run| run.direction);
if self.state.position.is_some() && !self.state.removed_indexes.is_empty() {
self.close_current_leaf(direction)?;
}
self.splice_open_run()
}
pub(super) fn poisoned(&self) -> bool {
self.state.poisoned
}
fn poison(&mut self) {
self.state.poisoned = true;
self.state.position = None;
self.state.removed_indexes.clear();
self.state.leaf_run_rewrite = None;
#[cfg(feature = "experimental_cursor")]
{
self.state.insert_run = None;
}
}
fn check_not_poisoned(&self) -> Result {
if self.state.poisoned {
return Err(StorageError::PreviousIo);
}
Ok(())
}
fn flush_removed_entries(&mut self, direction: Direction) -> Result<Vec<u8>> {
assert!(!self.state.removed_indexes.is_empty());
let position = self
.state
.position
.take()
.expect("cursor must be positioned");
let resume_key = scan_boundary_key::<K, V>(&position.leaf, direction);
let (path, leaf) = position.into_parts();
let allow_in_place = !self.state.detached_guards;
let removed_indexes = self.take_removals_ascending();
let result = self.mutate_helper().delete_leaf_entries(
leaf.page,
path,
&removed_indexes,
allow_in_place,
);
if result.is_err() {
self.poison();
}
result?;
Ok(resume_key)
}
fn close_current_leaf(&mut self, direction: Direction) -> Result<LeafCloseOutcome> {
assert!(self.state.position.is_some(), "cursor must be positioned");
let run_open = self.state.leaf_run_rewrite.is_some();
if self.state.removed_indexes.is_empty() {
if !run_open {
return Ok(LeafCloseOutcome::Unchanged);
}
} else {
let (underfilling, packs, has_parent) = {
let position = self.state.position.as_ref().unwrap();
let accessor = LeafAccessor::new(
position.leaf.page.memory(),
K::fixed_width(),
V::fixed_width(),
);
let (retained_pairs, retained_bytes) =
retained_after_removals(&accessor, &self.state.removed_indexes);
let page_size = self.page_allocator.get_page_size();
let underfilling = retained_pairs == 0
|| leaf_below_merge_threshold(
retained_pairs,
retained_bytes,
K::fixed_width(),
V::fixed_width(),
page_size,
);
let packs = underfilling
|| leaf_fits_one_page(
retained_pairs,
retained_bytes,
K::fixed_width(),
V::fixed_width(),
page_size,
);
(underfilling, packs, !position.path.is_empty())
};
if !(underfilling || run_open) || !has_parent {
let resume_key = self.flush_removed_entries(direction)?;
return Ok(LeafCloseOutcome::Flushed { resume_key });
}
let keeps_run_open = packs && self.run_parent_has_more_children(direction);
let removed_indexes = self.take_removals_ascending();
self.append_leaf_to_run(direction, &removed_indexes);
if keeps_run_open {
return Ok(LeafCloseOutcome::AbsorbedIntoRun);
}
}
let position = self.state.position.as_ref().unwrap();
let resume_key = scan_boundary_key::<K, V>(&position.leaf, direction);
self.splice_open_run()?;
Ok(LeafCloseOutcome::Flushed { resume_key })
}
fn take_removals_ascending(&mut self) -> Vec<usize> {
let mut removed_indexes = core::mem::take(&mut self.state.removed_indexes);
if removed_indexes.first() > removed_indexes.last() {
removed_indexes.reverse();
}
self.state.detached_guards = false;
removed_indexes
}
fn resume_after_rewrite(&mut self, direction: Direction, key: &[u8]) -> Result {
match direction {
Direction::Next => self.seek_to(Position::After(key)),
Direction::Previous => self.seek_to(Position::Before(key)),
}
}
fn append_leaf_to_run(&mut self, direction: Direction, removed_indexes: &[usize]) {
let (page, parent_page, child_index) = {
let frame = self.run_parent_frame();
let position = self.state.position.as_ref().unwrap();
(
position.leaf.page.clone(),
frame.page.get_page_number(),
frame.child_index,
)
};
let run = self
.state
.leaf_run_rewrite
.get_or_insert_with(|| LeafRunRewrite::new(parent_page, child_index, direction));
assert_eq!(run.parent_page, parent_page);
assert!(run.direction == direction);
run.append_entries_from::<K, V>(page, child_index, removed_indexes);
}
pub(super) fn splice_open_run(&mut self) -> Result {
self.check_not_poisoned()?;
let Some(run) = self.state.leaf_run_rewrite.take() else {
return Ok(());
};
let position = self
.state
.position
.take()
.expect("open run requires a position");
let result = self.splice_run(run, position);
if result.is_err() {
self.poison();
}
result
}
fn splice_run(&mut self, run: LeafRunRewrite, position: CursorPosition) -> Result {
let CursorPosition { path, leaf } = position;
drop(leaf);
assert_eq!(
path.last()
.expect("leaf runs require a parent branch")
.page
.get_page_number(),
run.parent_page
);
self.mutate_helper().replace_leaf_children(
path.into_iter().map(Branch::into_parts).collect(),
run.replaced_children,
run.entries,
run.removed_pairs,
)
}
fn mutate_helper<'c>(&'c mut self) -> MutateHelper<'a, 'c, K, V> {
assert!(self.state.leaf_run_rewrite.is_none());
#[cfg(feature = "experimental_cursor")]
assert!(self.state.insert_run.is_none());
MutateHelper::new(
&mut *self.root,
self.page_allocator,
&mut *self.freed,
self.allocated,
)
}
fn remove_leaf_entry(
&mut self,
leaf: PageImpl,
path: Vec<Branch>,
index: usize,
) -> Result<Option<(AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
assert!(self.state.removed_indexes.is_empty());
let path = path.into_iter().map(Branch::into_parts).collect();
let entry = self.mutate_helper().pop_leaf_entry(leaf, path, index)?;
Ok(Some(entry))
}
}
#[cfg(feature = "experimental_cursor")]
impl<K: Key + 'static, V: Value + 'static> CursorMut<'_, '_, K, V> {
pub(super) fn insert_before(&mut self, key: &[u8], value: &[u8]) -> Result<bool> {
self.check_not_poisoned()?;
assert!(self.state.removed_indexes.is_empty());
assert!(self.state.leaf_run_rewrite.is_none());
self.ensure_insert_run(RunDirection::Ascending)?;
let run = self.state.insert_run.as_mut().unwrap();
if run.rejects::<K>(key) {
if run.inserted_pairs == 0 {
self.state.insert_run = None;
}
return Ok(false);
}
run.entries.push_back(key, value);
run.previous_key = Some(key.to_vec());
run.inserted_pairs += 1;
if run.entries.total_bytes() >= INSERT_FLUSH_BYTES {
self.flush_insert_run(true)?;
}
Ok(true)
}
pub(super) fn insert_after(&mut self, key: &[u8], value: &[u8]) -> Result<bool> {
self.check_not_poisoned()?;
assert!(self.state.removed_indexes.is_empty());
assert!(self.state.leaf_run_rewrite.is_none());
self.ensure_insert_run(RunDirection::Descending)?;
let run = self.state.insert_run.as_mut().unwrap();
if run.rejects::<K>(key) {
if run.inserted_pairs == 0 {
self.state.insert_run = None;
}
return Ok(false);
}
run.entries.push_front(key, value);
run.inserted_pairs += 1;
if run.entries.total_bytes() >= INSERT_FLUSH_BYTES {
self.flush_insert_run(true)?;
}
Ok(true)
}
fn ensure_insert_run(&mut self, direction: RunDirection) -> Result {
if let Some(run) = &self.state.insert_run {
if run.direction == direction {
return Ok(());
}
self.flush_insert_run(true)?;
}
if self.state.insert_run.is_none() {
self.open_insert_run(direction)?;
}
Ok(())
}
fn open_insert_run(&mut self, direction: RunDirection) -> Result {
assert!(self.state.insert_run.is_none());
let opening_next_key = self.peek_next()?.map(|entry| entry.key_bytes().to_vec());
let previous_key = self.peek_prev()?.map(|entry| entry.key_bytes().to_vec());
let mut entries = OwnedEntryBuffer::default();
if direction == RunDirection::Ascending
&& let Some(position) = &self.state.position
{
let accessor = LeafAccessor::new(
position.leaf.page.memory(),
K::fixed_width(),
V::fixed_width(),
);
entries.extend_from_leaf_range(&accessor, 0..position.leaf.position, true);
}
self.state.insert_run = Some(Box::new(InsertRun {
direction,
opening_next_key,
previous_key,
entries,
inserted_pairs: 0,
}));
Ok(())
}
pub(super) fn flush_insert_run(&mut self, reseek: bool) -> Result {
self.check_not_poisoned()?;
let Some(run) = self.state.insert_run.take() else {
return Ok(());
};
if run.inserted_pairs == 0 {
return Ok(());
}
let resume_before = run.buffered_next().map(|(key, _)| key.to_vec());
let resume_after = run.previous_key;
let descending = run.direction == RunDirection::Descending;
let mut entries = run.entries;
let replaced = if let Some(position) = self.state.position.take() {
let accessor = LeafAccessor::new(
position.leaf.page.memory(),
K::fixed_width(),
V::fixed_width(),
);
if descending {
entries.extend_from_leaf_range(&accessor, 0..position.leaf.position, false);
}
entries.extend_from_leaf_range(
&accessor,
position.leaf.position..position.leaf.len,
true,
);
let CursorPosition { path, leaf } = position;
let replaced_leaf = leaf.page.get_page_number();
drop(leaf);
Some((
path.into_iter().map(Branch::into_parts).collect(),
replaced_leaf,
))
} else {
None
};
let result = self
.mutate_helper()
.splice_insert_run(replaced, &entries, run.inserted_pairs);
if result.is_err() {
self.poison();
}
result?;
if reseek {
if let Some(key) = &resume_before {
self.seek_to(Position::Before(key))?;
} else {
let key = resume_after.expect("a run with inserts bounds its gap on some side");
self.seek_to(Position::After(&key))?;
}
}
Ok(())
}
}
pub(super) struct CursorTree<'a, K: Key + 'static, V: Value + 'static> {
root: &'a mut Option<BtreeHeader>,
page_allocator: PageAllocator,
allocated: Arc<PageTracker>,
master_free_list: Arc<Mutex<Vec<PageNumber>>>,
freed: Vec<PageNumber>,
_key_type: PhantomData<K>,
_value_type: PhantomData<V>,
}
impl<'a, K: Key + 'static, V: Value + 'static> CursorTree<'a, K, V> {
fn new(
root: &'a mut Option<BtreeHeader>,
page_allocator: PageAllocator,
master_free_list: Arc<Mutex<Vec<PageNumber>>>,
allocated: Arc<PageTracker>,
) -> Self {
Self {
root,
page_allocator,
allocated,
master_free_list,
freed: vec![],
_key_type: PhantomData,
_value_type: PhantomData,
}
}
fn cursor(&mut self, state: CursorState) -> CursorMut<'a, '_, K, V> {
CursorMut::with_state(
&mut *self.root,
&self.page_allocator,
&mut self.freed,
&self.allocated,
state,
)
}
fn drain_freed(&mut self) {
if self.freed.is_empty() {
return;
}
let mut master_free_list = self.master_free_list.lock().unwrap();
for page in self.freed.drain(..) {
if !self
.page_allocator
.free_if_uncommitted(page, &self.allocated)
{
master_free_list.push(page);
}
}
}
}
#[cfg(feature = "experimental-api-5")]
pub(crate) struct BtreeCursor<K: Key + 'static, V: Value + 'static> {
inner: Option<Cursor<K, V>>,
}
#[cfg(feature = "experimental-api-5")]
impl<K: Key + 'static, V: Value + 'static> BtreeCursor<K, V> {
pub(crate) fn new(root: Option<BtreeHeader>, resolver: PageResolver, hint: PageHint) -> Self {
Self {
inner: root.map(|header| Cursor::new(header.root, resolver, hint)),
}
}
pub(crate) fn seek_lower_bound(&mut self, bound: Bound<&[u8]>) -> Result {
if let Some(cursor) = &mut self.inner {
cursor.seek_to(Position::from_lower_bound(bound))?;
}
Ok(())
}
pub(crate) fn seek_upper_bound(&mut self, bound: Bound<&[u8]>) -> Result {
if let Some(cursor) = &mut self.inner {
cursor.seek_to(Position::from_upper_bound(bound))?;
}
Ok(())
}
}
#[cfg(feature = "experimental_cursor")]
impl<K: Key + 'static, V: Value + 'static> BtreeCursor<K, V> {
#[allow(clippy::type_complexity)]
pub(crate) fn peek_next(
&mut self,
) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
let Some(cursor) = &mut self.inner else {
return Ok(None);
};
Ok(cursor.peek_next()?.map(entry_guards))
}
#[allow(clippy::type_complexity)]
pub(crate) fn peek_prev(
&mut self,
) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
let Some(cursor) = &mut self.inner else {
return Ok(None);
};
Ok(cursor.peek_prev()?.map(entry_guards))
}
#[allow(clippy::type_complexity)]
pub(crate) fn next(
&mut self,
) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
let Some(cursor) = &mut self.inner else {
return Ok(None);
};
Ok(cursor.next()?.map(entry_guards))
}
#[allow(clippy::type_complexity)]
pub(crate) fn prev(
&mut self,
) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
let Some(cursor) = &mut self.inner else {
return Ok(None);
};
Ok(cursor.prev()?.map(entry_guards))
}
}
#[cfg(feature = "experimental_cursor")]
fn entry_guards<K: Key + 'static, V: Value + 'static>(
entry: EntryGuard<K, V>,
) -> (AccessGuard<'static, K>, AccessGuard<'static, V>) {
let (page, key_range, value_range) = entry.into_raw();
(
AccessGuard::with_page(page.clone(), key_range),
AccessGuard::with_page(page, value_range),
)
}
#[cfg(feature = "experimental_cursor")]
pub(crate) struct BtreeCursorMut<'a, K: Key + 'static, V: Value + 'static> {
tree: CursorTree<'a, K, V>,
state: CursorState,
pending_reseek: Option<Vec<u8>>,
}
#[cfg(feature = "experimental_cursor")]
impl<'a, K: Key + 'static, V: Value + 'static> BtreeCursorMut<'a, K, V> {
pub(crate) fn new(
root: &'a mut Option<BtreeHeader>,
page_allocator: PageAllocator,
master_free_list: Arc<Mutex<Vec<PageNumber>>>,
allocated: Arc<PageTracker>,
) -> Self {
Self {
tree: CursorTree::new(root, page_allocator, master_free_list, allocated),
state: CursorState::default(),
pending_reseek: None,
}
}
pub(crate) fn seek_lower_bound(&mut self, bound: Bound<&[u8]>) -> Result {
self.pending_reseek = None;
self.with_cursor(|cursor| cursor.seek_to(Position::from_lower_bound(bound)))
}
pub(crate) fn seek_upper_bound(&mut self, bound: Bound<&[u8]>) -> Result {
self.pending_reseek = None;
self.with_cursor(|cursor| cursor.seek_to(Position::from_upper_bound(bound)))
}
fn settle_reseek(&mut self) -> Result {
let Some(key) = self.pending_reseek.take() else {
return Ok(());
};
self.with_cursor(|cursor| cursor.seek_to(Position::Before(&key)))
}
#[allow(clippy::type_complexity)]
pub(crate) fn peek_next(&mut self) -> Result<Option<(AccessGuard<'_, K>, AccessGuard<'_, V>)>> {
self.settle_reseek()?;
if let Some(run) = &self.state.insert_run {
if let Some((key, value)) = run.buffered_next() {
return Ok(Some((
AccessGuard::with_owned_value(key.to_vec()),
AccessGuard::with_owned_value(value.to_vec()),
)));
}
let Some(next_key) = &run.opening_next_key else {
return Ok(None);
};
let header = self.tree.root.expect("a captured next key implies a root");
let mut cursor: Cursor<K, V> = Cursor::new(
header.root,
self.tree.page_allocator.resolver(),
PageHint::None,
);
cursor.seek_to(Position::Before(next_key))?;
let entry = cursor
.next()?
.expect("the captured next key is in the tree");
return Ok(Some(entry_guards(entry)));
}
self.with_cursor(|cursor| Ok(cursor.peek_next()?.map(|entry| entry.to_guards())))
}
#[allow(clippy::type_complexity)]
pub(crate) fn peek_prev(&mut self) -> Result<Option<(AccessGuard<'_, K>, AccessGuard<'_, V>)>> {
self.settle_reseek()?;
if let Some(run) = &self.state.insert_run {
if let Some((key, value)) = run.buffered_previous() {
return Ok(Some((
AccessGuard::with_owned_value(key.to_vec()),
AccessGuard::with_owned_value(value.to_vec()),
)));
}
if run.direction == RunDirection::Descending
&& let Some(position) = &self.state.position
&& position.leaf.position > 0
{
let entry = entry_ref::<K, V>(&position.leaf, position.leaf.position - 1);
return Ok(Some(entry.to_guards()));
}
return Ok(None);
}
self.with_cursor(|cursor| Ok(cursor.peek_prev()?.map(|entry| entry.to_guards())))
}
#[allow(clippy::type_complexity)]
pub(crate) fn next(&mut self) -> Result<Option<(AccessGuard<'_, K>, AccessGuard<'_, V>)>> {
self.settle_reseek()?;
self.with_cursor(|cursor| {
cursor.flush_insert_run(true)?;
Ok(cursor.next()?.map(|entry| entry.to_guards()))
})
}
#[allow(clippy::type_complexity)]
pub(crate) fn prev(&mut self) -> Result<Option<(AccessGuard<'_, K>, AccessGuard<'_, V>)>> {
self.settle_reseek()?;
self.with_cursor(|cursor| {
cursor.flush_insert_run(true)?;
Ok(cursor.prev()?.map(|entry| entry.to_guards()))
})
}
pub(crate) fn insert_before(&mut self, key: &[u8], value: &[u8]) -> Result<bool> {
self.settle_reseek()?;
self.with_cursor(|cursor| cursor.insert_before(key, value))
}
pub(crate) fn insert_after(&mut self, key: &[u8], value: &[u8]) -> Result<bool> {
self.settle_reseek()?;
self.with_cursor(|cursor| cursor.insert_after(key, value))
}
#[allow(clippy::type_complexity)]
pub(crate) fn remove_next(
&mut self,
) -> Result<Option<(AccessGuard<'_, K>, AccessGuard<'_, V>)>> {
self.settle_reseek()?;
let removed = self.with_cursor(|cursor| {
cursor.flush_insert_run(true)?;
cursor.remove_next_taking_key()
})?;
Ok(removed.map(|(key, key_guard, value_guard)| {
self.pending_reseek = Some(key);
(key_guard, value_guard)
}))
}
#[allow(clippy::type_complexity)]
pub(crate) fn remove_prev(
&mut self,
) -> Result<Option<(AccessGuard<'_, K>, AccessGuard<'_, V>)>> {
self.settle_reseek()?;
let removed = self.with_cursor(|cursor| {
cursor.flush_insert_run(true)?;
cursor.remove_prev_taking_key()
})?;
Ok(removed.map(|(key, key_guard, value_guard)| {
self.pending_reseek = Some(key);
(key_guard, value_guard)
}))
}
pub(crate) fn apply_pending_inserts(&mut self) -> Result {
self.with_cursor(|cursor| cursor.flush_insert_run(true))
}
pub(crate) fn finish(&mut self) -> Result {
self.with_cursor(|cursor| cursor.flush_insert_run(false))
}
pub(crate) fn poisoned(&self) -> bool {
self.state.poisoned
}
fn with_cursor<R>(
&mut self,
operation: impl FnOnce(&mut CursorMut<'a, '_, K, V>) -> Result<R>,
) -> Result<R> {
let mut cursor = self.tree.cursor(core::mem::take(&mut self.state));
let result = operation(&mut cursor);
self.state = cursor.into_state();
self.tree.drain_freed();
result
}
}
enum EndState {
Parked(Bound<Vec<u8>>),
Pending(ParkedBatch),
Live(CursorState),
}
struct ParkedBatch {
bound: Bound<Vec<u8>>,
leaf_bytes: Arc<[u8]>,
removed_indexes: Vec<usize>,
}
pub(super) struct RangeMut<'a, K: Key + 'static, V: Value + 'static> {
tree: CursorTree<'a, K, V>,
front: EndState,
back: EndState,
settled: Option<Direction>,
poisoned: bool,
}
impl<'a, K: Key + 'static, V: Value + 'static> RangeMut<'a, K, V> {
pub(super) fn new(
root: &'a mut Option<BtreeHeader>,
lower_bound: Bound<Vec<u8>>,
upper_bound: Bound<Vec<u8>>,
page_allocator: PageAllocator,
master_free_list: Arc<Mutex<Vec<PageNumber>>>,
allocated: Arc<PageTracker>,
) -> Self {
Self {
tree: CursorTree::new(root, page_allocator, master_free_list, allocated),
front: EndState::Parked(lower_bound),
back: EndState::Parked(upper_bound),
settled: None,
poisoned: false,
}
}
pub(super) fn peek_next(&mut self) -> Result<Option<EntryRef<'_, K, V>>> {
self.peek(Direction::Next)
}
pub(super) fn peek_prev(&mut self) -> Result<Option<EntryRef<'_, K, V>>> {
self.peek(Direction::Previous)
}
pub(super) fn next(&mut self) -> Result<bool> {
self.advance(Direction::Next)
}
pub(super) fn prev(&mut self) -> Result<bool> {
self.advance(Direction::Previous)
}
pub(super) fn remove_next(
&mut self,
) -> Result<Option<(AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
self.remove(Direction::Next)
}
pub(super) fn remove_prev(
&mut self,
) -> Result<Option<(AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
self.remove(Direction::Previous)
}
pub(super) fn close(&mut self) -> Result {
let front = self.flush_end(Direction::Next);
let back = self.flush_end(Direction::Previous);
front.and(back)
}
pub(super) fn poisoned(&self) -> bool {
self.poisoned
}
fn check_not_poisoned(&self) -> Result {
if self.poisoned {
return Err(StorageError::PreviousIo);
}
Ok(())
}
fn flush_end(&mut self, direction: Direction) -> Result {
self.check_not_poisoned()?;
if matches!(self.end_ref(direction), EndState::Live(_)) {
let result =
self.with_live_cursor(direction, |cursor| cursor.finish_pending_removals());
result.and(self.park(direction))
} else {
self.apply_pending(direction)
}
}
fn peek(&mut self, direction: Direction) -> Result<Option<EntryRef<'_, K, V>>> {
if !self.settle(direction)? {
return Ok(None);
}
let EndState::Live(state) = self.end_ref(direction) else {
unreachable!("settled end must be live");
};
let position = state.position.as_ref().expect("settled end is positioned");
Ok(Some(entry_ref(
&position.leaf,
position.entry_index(direction),
)))
}
fn advance(&mut self, direction: Direction) -> Result<bool> {
if !self.settle(direction)? {
return Ok(false);
}
self.settled = None;
let EndState::Live(state) = self.end_mut(direction) else {
unreachable!("settled end must be live");
};
state
.position
.as_mut()
.expect("settled end is positioned")
.move_once(direction);
Ok(true)
}
fn remove(
&mut self,
direction: Direction,
) -> Result<Option<(AccessGuard<'a, K>, AccessGuard<'a, V>)>> {
if !self.settle(direction)? {
return Ok(None);
}
self.settled = None;
let result = self.with_live_cursor(direction, |cursor| match direction {
Direction::Next => cursor.remove_next_deferred(),
Direction::Previous => cursor.remove_prev_deferred(),
})?;
Ok(Some(result.expect("settled entry must be removable")))
}
fn settle(&mut self, direction: Direction) -> Result<bool> {
self.check_not_poisoned()?;
if self.settled == Some(direction) {
return Ok(true);
}
self.activate(direction)?;
let has_entry = self.with_live_cursor(direction, |cursor| match direction {
Direction::Next => Ok(cursor.peek_next()?.is_some()),
Direction::Previous => Ok(cursor.peek_prev()?.is_some()),
})?;
if !has_entry {
return Ok(false);
}
if !self.entry_in_range(direction) {
return Ok(false);
}
self.settled = Some(direction);
Ok(true)
}
fn activate(&mut self, direction: Direction) -> Result {
if matches!(self.end_ref(direction), EndState::Live(_)) {
return Ok(());
}
self.park(direction.opposite())?;
let mut state = self.seek_end(direction)?;
if matches!(self.end_ref(direction), EndState::Pending(_)) {
let EndState::Pending(batch) =
core::mem::replace(self.end_mut(direction), EndState::Parked(Unbounded))
else {
unreachable!();
};
if snapshot_matches(&state, &batch.leaf_bytes) {
state.removed_indexes = batch.removed_indexes;
state.detached_guards = true;
} else {
drop(state);
*self.end_mut(direction) = EndState::Parked(batch.bound.clone());
self.resolve_batch(batch)?;
state = self.seek_end(direction)?;
}
}
*self.end_mut(direction) = EndState::Live(state);
Ok(())
}
fn seek_end(&mut self, direction: Direction) -> Result<CursorState> {
let bound = match self.end_ref(direction) {
EndState::Parked(bound) => bound.clone(),
EndState::Pending(batch) => batch.bound.clone(),
EndState::Live(_) => unreachable!("end must be parked"),
};
let bound = bound.as_ref().map(Vec::as_slice);
let target = match direction {
Direction::Next => Position::from_lower_bound(bound),
Direction::Previous => Position::from_upper_bound(bound),
};
let mut cursor = self.tree.cursor(CursorState::default());
let result = cursor.seek_to(target);
let state = cursor.into_state();
result?;
Ok(state)
}
fn park(&mut self, direction: Direction) -> Result {
self.settled = None;
let end = self.end_mut(direction);
let EndState::Live(state) = end else {
return Ok(());
};
let bound = park_bound::<K, V>(state, direction);
let parked = if state.removed_indexes.is_empty() {
EndState::Parked(bound)
} else {
debug_assert!(state.detached_guards);
let position = state
.position
.as_ref()
.expect("pending removals require a position");
EndState::Pending(ParkedBatch {
bound,
leaf_bytes: position.leaf.page.to_arc(),
removed_indexes: core::mem::take(&mut state.removed_indexes),
})
};
let result = if state.leaf_run_rewrite.is_some() {
self.with_live_cursor(direction, |cursor| cursor.splice_open_run())
} else {
Ok(())
};
*self.end_mut(direction) = parked;
self.tree.drain_freed();
result
}
fn apply_pending(&mut self, direction: Direction) -> Result {
if !matches!(self.end_ref(direction), EndState::Pending(_)) {
return Ok(());
}
self.activate(direction)?;
self.with_live_cursor(direction, |cursor| cursor.finish_pending_removals())?;
self.park(direction)
}
fn resolve_batch(&mut self, batch: ParkedBatch) -> Result {
for direction in [Direction::Next, Direction::Previous] {
if let EndState::Live(state) = self.end_ref(direction) {
assert!(state.leaf_run_rewrite.is_none());
}
}
for &index in &batch.removed_indexes {
let key = LeafAccessor::new(&batch.leaf_bytes, K::fixed_width(), V::fixed_width())
.entry(index)
.expect("snapshot entry must exist")
.key();
let mut helper: MutateHelper<'_, '_, K, V> = MutateHelper::new(
&mut *self.tree.root,
&self.tree.page_allocator,
&mut self.tree.freed,
&self.tree.allocated,
);
let result = helper.delete_key(key, false);
self.tree.drain_freed();
match result {
Ok(removed) => debug_assert!(removed.is_some()),
Err(err) => {
self.poisoned = true;
return Err(err);
}
}
}
Ok(())
}
fn with_live_cursor<R>(
&mut self,
direction: Direction,
operation: impl FnOnce(&mut CursorMut<'a, '_, K, V>) -> Result<R>,
) -> Result<R> {
let end = self.end_mut(direction);
let EndState::Live(state) = core::mem::replace(end, EndState::Parked(Unbounded)) else {
unreachable!("end must be live");
};
let mut cursor = self.tree.cursor(state);
let result = operation(&mut cursor);
let state = cursor.into_state();
if state.poisoned {
self.poisoned = true;
}
*self.end_mut(direction) = EndState::Live(state);
self.tree.drain_freed();
result
}
fn entry_in_range(&self, direction: Direction) -> bool {
let EndState::Live(state) = self.end_ref(direction) else {
unreachable!("end must be live");
};
let position = state.position.as_ref().expect("end is positioned");
let entry = entry_ref::<K, V>(&position.leaf, position.entry_index(direction));
let key = entry.key_bytes();
let bound = match self.end_ref(direction.opposite()) {
EndState::Parked(bound) => bound,
EndState::Pending(batch) => &batch.bound,
EndState::Live(_) => unreachable!("peer end must be parked while this end is live"),
};
match direction {
Direction::Next => match bound {
Included(bound) => K::compare(key, bound).is_le(),
Excluded(bound) => K::compare(key, bound).is_lt(),
Unbounded => true,
},
Direction::Previous => match bound {
Included(bound) => K::compare(key, bound).is_ge(),
Excluded(bound) => K::compare(key, bound).is_gt(),
Unbounded => true,
},
}
}
fn end_ref(&self, direction: Direction) -> &EndState {
match direction {
Direction::Next => &self.front,
Direction::Previous => &self.back,
}
}
fn end_mut(&mut self, direction: Direction) -> &mut EndState {
match direction {
Direction::Next => &mut self.front,
Direction::Previous => &mut self.back,
}
}
}
fn snapshot_matches(state: &CursorState, snapshot: &[u8]) -> bool {
state
.position
.as_ref()
.is_some_and(|position| position.leaf.page.memory() == snapshot)
}
fn park_bound<K: Key + 'static, V: Value + 'static>(
state: &CursorState,
direction: Direction,
) -> Bound<Vec<u8>> {
let Some(position) = state.position.as_ref() else {
return Unbounded;
};
let leaf = &position.leaf;
match direction {
Direction::Next if leaf.position < leaf.len => {
Included(key_data::<K, V>(leaf, leaf.position))
}
Direction::Previous if leaf.position > 0 => {
Included(key_data::<K, V>(leaf, leaf.position - 1))
}
_ => Excluded(scan_boundary_key::<K, V>(leaf, direction)),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tree_store::btree_base::{DEFERRED, LeafBuilder};
use crate::tree_store::{AllocationPolicy, InMemoryBackend, PAGE_SIZE, TransactionalMemory};
fn test_page_allocator() -> PageAllocator {
let mem = TransactionalMemory::new(
Box::new(InMemoryBackend::new()),
true,
PAGE_SIZE,
None,
0,
false,
)
.unwrap();
mem.reset_allocator_state().unwrap();
PageAllocator::new(Arc::new(mem), AllocationPolicy::Default)
}
fn leaf_root_with_entries(entries: &[u64]) -> (PageAllocator, PageNumber, Arc<PageTracker>) {
let page_allocator = test_page_allocator();
let allocated_pages = Arc::new(PageTracker::new_tracking());
let keys_and_values: Vec<_> = entries
.iter()
.map(|entry| {
(
u64::as_bytes(entry).as_ref().to_vec(),
u64::as_bytes(entry).as_ref().to_vec(),
)
})
.collect();
let mut builder = LeafBuilder::new(
&page_allocator,
&allocated_pages,
entries.len(),
u64::fixed_width(),
u64::fixed_width(),
);
for (key, value) in &keys_and_values {
builder.push(key, value);
}
let page = builder.build().unwrap();
let root = page.get_page_number();
drop(page);
(page_allocator, root, allocated_pages)
}
fn cursor_with_entries(entries: &[u64]) -> Cursor<u64, u64> {
let (page_allocator, root, _) = leaf_root_with_entries(entries);
let mut cursor = Cursor::<u64, u64>::new(root, page_allocator.resolver(), PageHint::None);
cursor.seek_to(Position::Start).unwrap();
cursor
}
#[test]
fn cursor_preserves_boundary_gap_after_failed_next() {
let mut cursor = cursor_with_entries(&[1, 2, 3]);
for expected in [1, 2, 3] {
assert_eq!(cursor.next().unwrap().unwrap().key(), expected);
}
assert!(cursor.next().unwrap().is_none());
assert_eq!(cursor.prev().unwrap().unwrap().key(), 3);
assert_eq!(cursor.prev().unwrap().unwrap().key(), 2);
}
#[test]
fn cursor_preserves_boundary_gap_after_failed_prev() {
let mut cursor = cursor_with_entries(&[1, 2, 3]);
assert!(cursor.prev().unwrap().is_none());
assert_eq!(cursor.next().unwrap().unwrap().key(), 1);
assert_eq!(cursor.next().unwrap().unwrap().key(), 2);
}
#[test]
fn cursor_mut_stays_parked_at_tree_edge() {
let (page_allocator, root_page, allocated) = leaf_root_with_entries(&[1, 2, 3]);
let mut root = Some(BtreeHeader::new(root_page, DEFERRED, 3));
let mut freed = vec![];
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor.seek_to(Position::Start).unwrap();
for _ in 0..3 {
assert!(cursor.move_next().unwrap());
}
assert!(cursor.peek_next().unwrap().is_none());
assert!(cursor.state.position.is_some());
assert!(matches!(
park_bound::<u64, u64>(&cursor.state, Direction::Next),
Excluded(_)
));
cursor.seek_to(Position::End).unwrap();
for _ in 0..3 {
assert!(cursor.peek_prev().unwrap().is_some());
cursor
.state
.position
.as_mut()
.unwrap()
.move_once(Direction::Previous);
}
assert!(cursor.peek_prev().unwrap().is_none());
assert!(cursor.state.position.is_some());
assert!(matches!(
park_bound::<u64, u64>(&cursor.state, Direction::Previous),
Excluded(_)
));
}
#[cfg(feature = "experimental_cursor")]
mod insert_tests {
use super::*;
use crate::tree_store::RawBtree;
use crate::tree_store::btree::UntypedBtreeMut;
fn insert(cursor: &mut CursorMut<'_, '_, u64, u64>, key: u64, value: u64) -> bool {
cursor
.insert_before(u64::as_bytes(&key).as_ref(), u64::as_bytes(&value).as_ref())
.unwrap()
}
fn insert_after(cursor: &mut CursorMut<'_, '_, u64, u64>, key: u64, value: u64) -> bool {
cursor
.insert_after(u64::as_bytes(&key).as_ref(), u64::as_bytes(&value).as_ref())
.unwrap()
}
fn scan(root: Option<BtreeHeader>, page_allocator: &PageAllocator) -> Vec<(u64, u64)> {
let Some(header) = root else {
return vec![];
};
let mut cursor =
Cursor::<u64, u64>::new(header.root, page_allocator.resolver(), PageHint::None);
cursor.seek_to(Position::Start).unwrap();
let mut entries = vec![];
while let Some(entry) = cursor.next().unwrap() {
entries.push((entry.key(), entry.value()));
}
entries
}
fn assert_tree(
root: Option<BtreeHeader>,
page_allocator: &PageAllocator,
expected: &[(u64, u64)],
) {
assert_eq!(scan(root, page_allocator), expected);
assert_eq!(
root.map_or(0, |header| header.length),
expected.len() as u64
);
for (key, value) in expected {
let mut cursor = Cursor::<u64, u64>::new(
root.unwrap().root,
page_allocator.resolver(),
PageHint::None,
);
cursor
.seek_to(Position::Before(u64::as_bytes(key).as_ref()))
.unwrap();
let entry = cursor.next().unwrap().expect("key must route to its leaf");
assert_eq!(entry.key(), *key);
assert_eq!(entry.value(), *value);
}
let mut untyped = UntypedBtreeMut::new(
root,
page_allocator.clone(),
Arc::new(Mutex::new(vec![])),
u64::fixed_width(),
u64::fixed_width(),
);
let finalized = untyped.finalize_dirty_checksums().unwrap();
let raw = RawBtree::new(
finalized,
u64::fixed_width(),
u64::fixed_width(),
page_allocator.resolver(),
PageHint::None,
);
assert!(raw.verify_checksum().unwrap());
}
#[test]
fn insert_run_builds_tree_from_empty() {
for flush_every in [1, 3, 64] {
let page_allocator = test_page_allocator();
let mut root = None;
let mut freed = vec![];
let allocated = Arc::new(PageTracker::new_tracking());
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor.seek_to(Position::End).unwrap();
for key in 0..2000 {
assert!(insert(&mut cursor, key, key * 3));
if key % flush_every == 0 {
cursor.flush_insert_run(true).unwrap();
}
}
cursor.flush_insert_run(true).unwrap();
drop(cursor);
let expected: Vec<_> = (0..2000).map(|key| (key, key * 3)).collect();
assert_tree(root, &page_allocator, &expected);
}
}
#[test]
fn insert_run_grows_multiple_levels() {
let page_allocator = test_page_allocator();
let mut root = None;
let mut freed = vec![];
let allocated = Arc::new(PageTracker::new_tracking());
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor.seek_to(Position::End).unwrap();
for key in 0..60_000 {
assert!(insert(&mut cursor, key, key));
if key % 1000 == 999 {
cursor.flush_insert_run(true).unwrap();
}
}
cursor.flush_insert_run(false).unwrap();
drop(cursor);
let stats = crate::tree_store::btree::btree_stats(
root.map(|header| header.root),
&page_allocator.resolver(),
u64::fixed_width(),
u64::fixed_width(),
PageHint::None,
)
.unwrap();
assert!(stats.tree_height >= 3, "height {}", stats.tree_height);
let expected: Vec<_> = (0..60_000).map(|key| (key, key)).collect();
assert_tree(root, &page_allocator, &expected);
}
#[test]
fn insert_run_into_leaf_middle() {
let existing: Vec<u64> = (0..100).map(|i| i * 10).collect();
let (page_allocator, root_page, allocated) = leaf_root_with_entries(&existing);
let mut root = Some(BtreeHeader::new(root_page, DEFERRED, existing.len() as u64));
let mut freed = vec![];
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor
.seek_to(Position::Before(u64::as_bytes(&500).as_ref()))
.unwrap();
for key in 491..500 {
assert!(insert(&mut cursor, key, key));
}
cursor.flush_insert_run(false).unwrap();
drop(cursor);
let mut expected: Vec<_> = existing.iter().map(|&key| (key, key)).collect();
expected.extend((491..500).map(|key| (key, key)));
expected.sort_unstable();
assert_tree(root, &page_allocator, &expected);
}
#[test]
fn insert_run_splits_root_leaf() {
let (page_allocator, root_page, allocated) = leaf_root_with_entries(&[0, 1_000_000]);
let mut root = Some(BtreeHeader::new(root_page, DEFERRED, 2));
let mut freed = vec![];
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor
.seek_to(Position::After(u64::as_bytes(&0).as_ref()))
.unwrap();
for key in 1..=10_000 {
assert!(insert(&mut cursor, key, key));
}
cursor.flush_insert_run(false).unwrap();
drop(cursor);
let mut expected = vec![(0, 0), (1_000_000, 1_000_000)];
expected.extend((1..=10_000).map(|key| (key, key)));
expected.sort_unstable();
assert_tree(root, &page_allocator, &expected);
}
#[test]
fn insert_run_into_middle_of_tall_tree() {
let page_allocator = test_page_allocator();
let mut root = None;
let mut freed = vec![];
let allocated = Arc::new(PageTracker::new_tracking());
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor.seek_to(Position::End).unwrap();
for key in 0..60_000 {
assert!(insert(&mut cursor, key * 2, key * 2));
if key % 1000 == 999 {
cursor.flush_insert_run(true).unwrap();
}
}
cursor.flush_insert_run(false).unwrap();
drop(cursor);
let stats = crate::tree_store::btree::btree_stats(
root.map(|header| header.root),
&page_allocator.resolver(),
u64::fixed_width(),
u64::fixed_width(),
PageHint::None,
)
.unwrap();
assert!(stats.tree_height >= 3, "height {}", stats.tree_height);
let mut expected: Vec<(u64, u64)> = (0..60_000).map(|key| (key * 2, key * 2)).collect();
for target in [1001u64, 30_001, 60_001, 90_001, 119_001] {
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor
.seek_to(Position::Before(u64::as_bytes(&target).as_ref()))
.unwrap();
assert!(insert(&mut cursor, target, target));
cursor.flush_insert_run(false).unwrap();
drop(cursor);
expected.push((target, target));
}
expected.sort_unstable();
assert_tree(root, &page_allocator, &expected);
}
#[test]
fn insert_run_swaps_ancestors_in_place() {
let page_allocator = test_page_allocator();
let mut root = None;
let mut freed = vec![];
let allocated = Arc::new(PageTracker::new_tracking());
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor.seek_to(Position::End).unwrap();
for key in 0..60_000 {
assert!(insert(&mut cursor, key * 2, key * 2));
if key % 1000 == 999 {
cursor.flush_insert_run(true).unwrap();
}
}
cursor.flush_insert_run(false).unwrap();
drop(cursor);
let stats = crate::tree_store::btree::btree_stats(
root.map(|header| header.root),
&page_allocator.resolver(),
u64::fixed_width(),
u64::fixed_width(),
PageHint::None,
)
.unwrap();
assert!(stats.tree_height >= 3, "height {}", stats.tree_height);
for target in [30_001u64, 30_003] {
let root_before = root.map(|header| header.root);
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor
.seek_to(Position::Before(u64::as_bytes(&target).as_ref()))
.unwrap();
assert!(insert(&mut cursor, target, target));
cursor.flush_insert_run(false).unwrap();
drop(cursor);
if target == 30_003 {
assert_eq!(root.map(|header| header.root), root_before);
}
}
let mut expected: Vec<(u64, u64)> = (0..60_000).map(|key| (key * 2, key * 2)).collect();
expected.push((30_001, 30_001));
expected.push((30_003, 30_003));
expected.sort_unstable();
assert_tree(root, &page_allocator, &expected);
}
#[test]
fn insert_run_rejects_unordered_keys() {
let (page_allocator, root_page, allocated) = leaf_root_with_entries(&[10, 20, 30]);
let mut root = Some(BtreeHeader::new(root_page, DEFERRED, 3));
let mut freed = vec![];
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor
.seek_to(Position::After(u64::as_bytes(&20).as_ref()))
.unwrap();
assert!(!insert(&mut cursor, 20, 20));
assert!(!insert(&mut cursor, 15, 15));
assert!(!insert(&mut cursor, 30, 30));
assert!(!insert(&mut cursor, 35, 35));
assert!(cursor.state.insert_run.is_none());
cursor
.seek_to(Position::After(u64::as_bytes(&20).as_ref()))
.unwrap();
assert!(insert(&mut cursor, 25, 25));
assert!(!insert(&mut cursor, 25, 25));
assert!(!insert(&mut cursor, 24, 24));
assert!(cursor.state.insert_run.is_some());
assert!(insert(&mut cursor, 26, 26));
cursor.flush_insert_run(false).unwrap();
drop(cursor);
let expected = [10, 20, 25, 26, 30].map(|key| (key, key));
assert_tree(root, &page_allocator, &expected);
}
#[test]
fn insert_run_descending_from_empty() {
for flush_every in [1, 3, 64] {
let page_allocator = test_page_allocator();
let mut root = None;
let mut freed = vec![];
let allocated = Arc::new(PageTracker::new_tracking());
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor.seek_to(Position::Start).unwrap();
for key in (0..2000).rev() {
assert!(insert_after(&mut cursor, key, key * 3));
if key % flush_every == 0 {
cursor.flush_insert_run(true).unwrap();
}
}
cursor.flush_insert_run(true).unwrap();
drop(cursor);
let expected: Vec<_> = (0..2000).map(|key| (key, key * 3)).collect();
assert_tree(root, &page_allocator, &expected);
}
}
#[test]
fn insert_run_mixed_directions() {
let (page_allocator, root_page, allocated) = leaf_root_with_entries(&[10, 20, 30]);
let mut root = Some(BtreeHeader::new(root_page, DEFERRED, 3));
let mut freed = vec![];
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor
.seek_to(Position::After(u64::as_bytes(&20).as_ref()))
.unwrap();
assert!(insert(&mut cursor, 21, 21));
assert!(insert_after(&mut cursor, 29, 29));
assert!(insert_after(&mut cursor, 25, 25));
assert!(insert(&mut cursor, 22, 22));
assert!(!insert(&mut cursor, 25, 25));
assert!(!insert_after(&mut cursor, 22, 22));
assert!(!insert_after(&mut cursor, 25, 25));
assert!(insert(&mut cursor, 24, 24));
cursor.flush_insert_run(true).unwrap();
assert_eq!(cursor.peek_prev().unwrap().unwrap().key(), 24);
assert_eq!(cursor.peek_next().unwrap().unwrap().key(), 25);
drop(cursor);
let expected = [10, 20, 21, 22, 24, 25, 29, 30].map(|key| (key, key));
assert_tree(root, &page_allocator, &expected);
}
#[test]
fn insert_run_after_rejects_unordered_keys() {
let (page_allocator, root_page, allocated) = leaf_root_with_entries(&[10, 20, 30]);
let mut root = Some(BtreeHeader::new(root_page, DEFERRED, 3));
let mut freed = vec![];
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor
.seek_to(Position::After(u64::as_bytes(&20).as_ref()))
.unwrap();
assert!(!insert_after(&mut cursor, 20, 20));
assert!(!insert_after(&mut cursor, 15, 15));
assert!(!insert_after(&mut cursor, 30, 30));
assert!(!insert_after(&mut cursor, 35, 35));
assert!(cursor.state.insert_run.is_none());
cursor
.seek_to(Position::After(u64::as_bytes(&20).as_ref()))
.unwrap();
assert!(insert_after(&mut cursor, 25, 25));
assert!(!insert_after(&mut cursor, 25, 25));
assert!(!insert_after(&mut cursor, 26, 26));
assert!(cursor.state.insert_run.is_some());
assert!(!insert(&mut cursor, 25, 25));
assert!(cursor.state.insert_run.is_none());
assert!(insert(&mut cursor, 21, 21));
cursor.flush_insert_run(false).unwrap();
drop(cursor);
let expected = [10, 20, 21, 25, 30].map(|key| (key, key));
assert_tree(root, &page_allocator, &expected);
}
}
#[test]
fn poisoned_cursor_mut_re_raises() {
let (page_allocator, root_page, allocated) = leaf_root_with_entries(&[1, 2, 3]);
let mut root = Some(BtreeHeader::new(root_page, DEFERRED, 3));
let mut freed = vec![];
let mut cursor: CursorMut<'_, '_, u64, u64> =
CursorMut::new(&mut root, &page_allocator, &mut freed, &allocated);
cursor.seek_to(Position::Start).unwrap();
cursor.poison();
assert!(cursor.poisoned());
assert!(matches!(cursor.peek_next(), Err(StorageError::PreviousIo)));
assert!(matches!(
cursor.seek_to(Position::Start),
Err(StorageError::PreviousIo)
));
assert!(matches!(
cursor.finish_pending_removals(),
Err(StorageError::PreviousIo)
));
assert!(matches!(
cursor.splice_open_run(),
Err(StorageError::PreviousIo)
));
}
}