#![forbid(unsafe_code)]
use core::borrow::Borrow;
use super::node::inline_keys::{SearchKey, WindowedKey};
use super::node::internal::Internal;
use super::node::leaf::Leaf;
use super::node::{Child, Node};
use super::node_stack::{FirstEdge, MAX_DEPTH, NodeStackMut};
use super::root::Tree;
use crate::allocator::Allocator;
use crate::inline_key::InlineKeyStrategy;
impl<K, V, S: InlineKeyStrategy<K>> Tree<K, V, S> {
pub(crate) fn remove_entry<Q>(
&mut self,
key: &Q,
alloc: &impl Allocator,
length: &mut usize,
) -> Option<(K, V)>
where
K: Borrow<Q>,
Q: Ord + ?Sized,
{
if !self.has_root() {
return None;
}
let mut stack = NodeStackMut::new(self);
if stack.search(&SearchKey::new(key)) {
Some(stack.remove(alloc, length))
} else {
None
}
}
pub(crate) fn pop_first(
&mut self,
alloc: &impl Allocator,
length: &mut usize,
) -> Option<(K, V)> {
if !self.has_root() {
return None;
}
let mut stack = NodeStackMut::new(self);
if stack.descend_to_first_key() {
Some(stack.remove(alloc, length))
} else {
None
}
}
pub(crate) fn pop_last(
&mut self,
alloc: &impl Allocator,
length: &mut usize,
) -> Option<(K, V)> {
if !self.has_root() {
return None;
}
let mut stack = NodeStackMut::new(self);
if stack.descend_to_last_key() {
Some(stack.remove(alloc, length))
} else {
None
}
}
}
impl<K, V, S: InlineKeyStrategy<K>> NodeStackMut<'_, K, V, S> {
pub(crate) fn remove(&mut self, alloc: &impl Allocator, length: &mut usize) -> (K, V) {
let entry_depth = self.depth();
let entry_index = self.index();
let (entry, changed) = if let Node::Leaf(leaf) = self.node() {
let entry = leaf.remove(entry_index);
(entry, self.fix_up_the_stack(None, alloc, |_| {}))
} else {
self.refit_left_of_separator();
self.set_index(entry_index + 1);
self.descend_to_leaf::<FirstEdge>();
let succession = self
.leaf()
.take_head_and_first_key()
.expect("the leaf right of a separator has keys");
let mut successor = Some(succession.successor);
let mut replaced = None;
let changed = self.fix_up_the_stack(Some(entry_depth), alloc, |node| {
let successor = successor.take().expect("the separator is replaced once");
replaced = Some(node.replace_separator_from_below(entry_index, successor));
});
let replaced = replaced.expect("the node stack goes up past the separator");
let key = (succession.removed_key)
.or_else(|| replaced.into_windowed_key().map(WindowedKey::into_key))
.expect("a removed separator's key is in its head or its node");
((key, succession.value), changed)
};
*length -= 1;
self.refit_changed_fences(&changed);
self.detach_lowering_an_emptied_root(changed.reached_root(), alloc);
entry
}
fn fix_up_the_stack(
&mut self,
must_reach: Option<usize>,
alloc: &impl Allocator,
mut at_must_reach: impl FnMut(Internal<K, V, S>),
) -> ChangedFences {
let mut changed = ChangedFences::new(self.depth(), must_reach);
let mut child_is_underfull = self.leaf().is_underfull();
while !self.is_at_root()
&& (child_is_underfull || must_reach.is_some_and(|depth| self.depth() > depth))
{
let node = self.ascend();
let (depth, index) = (self.depth(), self.index());
changed.indexes[depth] = index;
if Some(depth) == must_reach {
at_must_reach(node);
}
if child_is_underfull {
let fix = if node.has_leaf_children() {
node.fix_underfull_child::<Leaf<K, V, S>>(index, alloc)
} else {
node.fix_underfull_child::<Internal<K, V, S>>(index, alloc)
};
let merged = matches!(
fix,
ChildFix::MergedIntoLeft { .. } | ChildFix::MergedRightIn
);
changed.record(depth, fix);
child_is_underfull = merged && node.is_underfull();
}
}
changed
}
fn refit_changed_fences(&mut self, changed: &ChangedFences) {
let Some(top_depth) = changed.top_depth else {
return;
};
while self.depth() + 1 < changed.leaf_depth {
self.set_index(changed.indexes[self.depth()]);
self.descend();
}
loop {
let depth = self.depth();
let child = changed.indexes[depth];
self.set_index(child);
match changed.donors[depth] {
Some(donor) => self.refit_children_to_fences([child, donor]),
None => self.refit_children_to_fences([child]),
}
if depth == top_depth {
return;
}
self.ascend();
}
}
}
enum ChildFix {
Nothing,
MergedIntoLeft { left_len: usize },
MergedRightIn,
TookFromLeft,
TookFromRight,
}
struct ChangedFences {
indexes: [usize; MAX_DEPTH + 1],
donors: [Option<usize>; MAX_DEPTH + 1],
leaf_depth: usize,
top_depth: Option<usize>,
}
impl ChangedFences {
fn reached_root(&self) -> bool {
self.top_depth == Some(0)
}
fn new(leaf_depth: usize, top_depth: Option<usize>) -> Self {
Self {
indexes: [0; MAX_DEPTH + 1],
donors: [None; MAX_DEPTH + 1],
leaf_depth,
top_depth,
}
}
fn record(&mut self, depth: usize, fix: ChildFix) {
let child = self.indexes[depth];
match fix {
ChildFix::Nothing => return,
ChildFix::MergedIntoLeft { left_len } => {
self.indexes[depth] = child - 1;
self.indexes[depth + 1] += left_len + 1;
}
ChildFix::MergedRightIn => {}
ChildFix::TookFromLeft => {
self.indexes[depth + 1] += 1;
self.donors[depth] = Some(child - 1);
}
ChildFix::TookFromRight => self.donors[depth] = Some(child + 1),
}
self.top_depth = Some(
self.top_depth
.map_or(depth, |top_depth| top_depth.min(depth)),
);
}
}
impl<K, V, S: InlineKeyStrategy<K>> Internal<K, V, S> {
fn fix_underfull_child<C: Child<K, V, S>>(
self,
index: usize,
alloc: &impl Allocator,
) -> ChildFix {
if C::child_of(self, index).len() >= C::min_len() {
return ChildFix::Nothing;
}
let (separator_index, child_is_right) = if index > 0 {
(index - 1, true)
} else if index < self.len() {
(index, false)
} else {
unreachable!("an internal node a removal went through has a separator")
};
let mut siblings = self.siblings::<C>(separator_index);
let left_len = siblings.left_len();
if siblings.can_merge() {
siblings.merge().deallocate(alloc);
return if child_is_right {
ChildFix::MergedIntoLeft { left_len }
} else {
ChildFix::MergedRightIn
};
}
let stolen = if child_is_right {
siblings.steal_left()
} else {
siblings.steal_right()
};
stolen.expect("a sibling too full to merge with can spare an entry");
if child_is_right {
ChildFix::TookFromLeft
} else {
ChildFix::TookFromRight
}
}
}