use super::node::Removed;
use super::{slice_insert, slice_shift, Comparator, Forest, Node, NodeData, NodePool, MAX_PATH};
use core::borrow::Borrow;
use core::marker::PhantomData;
#[cfg(test)]
use core::fmt;
pub(super) struct Path<F: Forest> {
size: usize,
node: [Node; MAX_PATH],
entry: [u8; MAX_PATH],
unused: PhantomData<F>,
}
impl<F: Forest> Default for Path<F> {
fn default() -> Self {
Self {
size: 0,
node: [Node(0); MAX_PATH],
entry: [0; MAX_PATH],
unused: PhantomData,
}
}
}
impl<F: Forest> Path<F> {
pub fn find(
&mut self,
key: F::Key,
root: Node,
pool: &NodePool<F>,
comp: &dyn Comparator<F::Key>,
) -> Option<F::Value> {
let mut node = root;
for level in 0.. {
self.size = level + 1;
self.node[level] = node;
match pool[node] {
NodeData::Inner { size, keys, tree } => {
let i = match comp.search(key, &keys[0..size.into()]) {
Ok(i) => i + 1,
Err(i) => i,
};
self.entry[level] = i as u8;
node = tree[i];
}
NodeData::Leaf { size, keys, vals } => {
return match comp.search(key, &keys.borrow()[0..size.into()]) {
Ok(i) => {
self.entry[level] = i as u8;
Some(vals.borrow()[i])
}
Err(i) => {
self.entry[level] = i as u8;
None
}
};
}
NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
}
}
unreachable!();
}
pub fn first(&mut self, root: Node, pool: &NodePool<F>) -> (F::Key, F::Value) {
let mut node = root;
for level in 0.. {
self.size = level + 1;
self.node[level] = node;
self.entry[level] = 0;
match pool[node] {
NodeData::Inner { tree, .. } => node = tree[0],
NodeData::Leaf { keys, vals, .. } => return (keys.borrow()[0], vals.borrow()[0]),
NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
}
}
unreachable!();
}
pub fn next(&mut self, pool: &NodePool<F>) -> Option<(F::Key, F::Value)> {
match self.leaf_pos() {
None => return None,
Some((node, entry)) => {
let (keys, vals) = pool[node].unwrap_leaf();
if entry + 1 < keys.len() {
self.entry[self.size - 1] += 1;
return Some((keys[entry + 1], vals[entry + 1]));
}
}
}
let leaf_level = self.size - 1;
self.next_node(leaf_level, pool).map(|node| {
let (keys, vals) = pool[node].unwrap_leaf();
(keys[0], vals[0])
})
}
pub fn prev(&mut self, root: Node, pool: &NodePool<F>) -> Option<(F::Key, F::Value)> {
if self.size == 0 {
self.goto_subtree_last(0, root, pool);
let (node, entry) = self.leaf_pos().unwrap();
let (keys, vals) = pool[node].unwrap_leaf();
return Some((keys[entry], vals[entry]));
}
match self.leaf_pos() {
None => return None,
Some((node, entry)) => {
if entry > 0 {
self.entry[self.size - 1] -= 1;
let (keys, vals) = pool[node].unwrap_leaf();
return Some((keys[entry - 1], vals[entry - 1]));
}
}
}
self.prev_leaf(pool).map(|node| {
let (keys, vals) = pool[node].unwrap_leaf();
let e = self.leaf_entry();
(keys[e], vals[e])
})
}
fn next_node(&mut self, level: usize, pool: &NodePool<F>) -> Option<Node> {
match self.right_sibling_branch_level(level, pool) {
None => {
self.size = 0;
None
}
Some(bl) => {
let (_, bnodes) = pool[self.node[bl]].unwrap_inner();
self.entry[bl] += 1;
let mut node = bnodes[usize::from(self.entry[bl])];
for l in bl + 1..level {
self.node[l] = node;
self.entry[l] = 0;
node = pool[node].unwrap_inner().1[0];
}
self.node[level] = node;
self.entry[level] = 0;
Some(node)
}
}
}
fn prev_leaf(&mut self, pool: &NodePool<F>) -> Option<Node> {
self.left_sibling_branch_level(self.size - 1).map(|bl| {
let entry = self.entry[bl] - 1;
self.entry[bl] = entry;
let (_, bnodes) = pool[self.node[bl]].unwrap_inner();
self.goto_subtree_last(bl + 1, bnodes[usize::from(entry)], pool)
})
}
fn goto_subtree_last(&mut self, level: usize, root: Node, pool: &NodePool<F>) -> Node {
let mut node = root;
for l in level.. {
self.node[l] = node;
match pool[node] {
NodeData::Inner { size, ref tree, .. } => {
self.entry[l] = size;
node = tree[usize::from(size)];
}
NodeData::Leaf { size, .. } => {
self.entry[l] = size - 1;
self.size = l + 1;
break;
}
NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
}
}
node
}
pub fn set_root_node(&mut self, root: Node) {
self.size = 1;
self.node[0] = root;
self.entry[0] = 0;
}
pub fn leaf_pos(&self) -> Option<(Node, usize)> {
let i = self.size.wrapping_sub(1);
self.node.get(i).map(|&n| (n, self.entry[i].into()))
}
fn leaf_node(&self) -> Node {
self.node[self.size - 1]
}
fn leaf_entry(&self) -> usize {
self.entry[self.size - 1].into()
}
fn at_first_entry(&self) -> bool {
self.entry[0..self.size].iter().all(|&i| i == 0)
}
pub fn value_mut<'a>(&self, pool: &'a mut NodePool<F>) -> &'a mut F::Value {
&mut pool[self.leaf_node()].unwrap_leaf_mut().1[self.leaf_entry()]
}
pub fn insert(&mut self, key: F::Key, value: F::Value, pool: &mut NodePool<F>) -> Node {
if !self.try_leaf_insert(key, value, pool) {
self.split_and_insert(key, value, pool);
}
self.node[0]
}
fn try_leaf_insert(&self, key: F::Key, value: F::Value, pool: &mut NodePool<F>) -> bool {
let index = self.leaf_entry();
debug_assert!(index > 0 || self.at_first_entry());
pool[self.leaf_node()].try_leaf_insert(index, key, value)
}
fn split_and_insert(&mut self, mut key: F::Key, value: F::Value, pool: &mut NodePool<F>) {
let orig_root = self.node[0];
let mut ins_node = None;
let mut split;
for level in (0..self.size).rev() {
let mut node = self.node[level];
let mut entry = self.entry[level].into();
split = pool[node].split(entry);
let rhs_node = pool.alloc_node(split.rhs_data);
if entry > split.lhs_entries
|| (entry == split.lhs_entries
&& (split.lhs_entries > split.rhs_entries || ins_node.is_some()))
{
node = rhs_node;
entry -= split.lhs_entries;
self.node[level] = node;
self.entry[level] = entry as u8;
}
match ins_node {
None => {
let inserted = pool[node].try_leaf_insert(entry, key, value);
debug_assert!(inserted);
if entry == 0 && node == rhs_node {
split.crit_key = key;
}
}
Some(n) => {
let inserted = pool[node].try_inner_insert(entry, key, n);
debug_assert!(inserted);
if n == self.node[level + 1] {
self.entry[level] += 1;
}
}
}
key = split.crit_key;
ins_node = Some(rhs_node);
if level > 0 {
let pnode = &mut pool[self.node[level - 1]];
let pentry = self.entry[level - 1].into();
if pnode.try_inner_insert(pentry, key, rhs_node) {
if node == rhs_node {
self.entry[level - 1] += 1;
}
return;
}
}
}
let rhs_node = ins_node.expect("empty path");
let root = pool.alloc_node(NodeData::inner(orig_root, key, rhs_node));
let entry = if self.node[0] == rhs_node { 1 } else { 0 };
self.size += 1;
slice_insert(&mut self.node[0..self.size], 0, root);
slice_insert(&mut self.entry[0..self.size], 0, entry);
}
pub fn remove(&mut self, pool: &mut NodePool<F>) -> Option<Node> {
let e = self.leaf_entry();
match pool[self.leaf_node()].leaf_remove(e) {
Removed::Healthy => {
if e == 0 {
self.update_crit_key(pool)
}
Some(self.node[0])
}
status => self.balance_nodes(status, pool),
}
}
fn current_crit_key(&self, level: usize, pool: &NodePool<F>) -> Option<F::Key> {
self.left_sibling_branch_level(level).map(|bl| {
let (keys, _) = pool[self.node[bl]].unwrap_inner();
keys[usize::from(self.entry[bl]) - 1]
})
}
fn update_crit_key(&mut self, pool: &mut NodePool<F>) {
let crit_level = match self.left_sibling_branch_level(self.size - 1) {
None => return,
Some(l) => l,
};
let crit_kidx = self.entry[crit_level] - 1;
let crit_key = pool[self.leaf_node()].leaf_crit_key();
let crit_node = self.node[crit_level];
match pool[crit_node] {
NodeData::Inner {
size, ref mut keys, ..
} => {
debug_assert!(crit_kidx < size);
keys[usize::from(crit_kidx)] = crit_key;
}
_ => panic!("Expected inner node"),
}
}
fn balance_nodes(&mut self, status: Removed, pool: &mut NodePool<F>) -> Option<Node> {
if status != Removed::Empty && self.leaf_entry() == 0 {
self.update_crit_key(pool);
}
let leaf_level = self.size - 1;
if self.heal_level(status, leaf_level, pool) {
self.size = 0;
return None;
}
let mut ns = 0;
while let NodeData::Inner {
size: 0, ref tree, ..
} = pool[self.node[ns]]
{
ns += 1;
self.node[ns] = tree[0];
}
if ns > 0 {
for l in 0..ns {
pool.free_node(self.node[l]);
}
slice_shift(&mut self.node, ns);
slice_shift(&mut self.entry, ns);
if self.size > 0 {
self.size -= ns;
}
}
Some(self.node[0])
}
fn heal_level(&mut self, status: Removed, level: usize, pool: &mut NodePool<F>) -> bool {
match status {
Removed::Healthy => {}
Removed::Rightmost => {
debug_assert_eq!(
usize::from(self.entry[level]),
pool[self.node[level]].entries()
);
self.next_node(level, pool);
}
Removed::Underflow => self.underflowed_node(level, pool),
Removed::Empty => return self.empty_node(level, pool),
}
false
}
fn underflowed_node(&mut self, level: usize, pool: &mut NodePool<F>) {
if let Some((crit_key, rhs_node)) = self.right_sibling(level, pool) {
let new_ck: Option<F::Key>;
let empty;
let mut rhs = pool[rhs_node];
match pool[self.node[level]].balance(crit_key, &mut rhs) {
None => {
new_ck = self.current_crit_key(level, pool);
empty = true;
}
Some(key) => {
new_ck = Some(key);
empty = false;
}
}
pool[rhs_node] = rhs;
if let Some(ck) = new_ck {
self.update_right_crit_key(level, ck, pool);
}
if empty {
let empty_tree = self.empty_node(level, pool);
debug_assert!(!empty_tree);
}
debug_assert!(usize::from(self.entry[level]) < pool[self.node[level]].entries());
} else if usize::from(self.entry[level]) >= pool[self.node[level]].entries() {
self.size = 0;
}
}
fn empty_node(&mut self, level: usize, pool: &mut NodePool<F>) -> bool {
pool.free_node(self.node[level]);
if level == 0 {
return true;
}
let rhs_node = self.right_sibling(level, pool).map(|(_, n)| n);
let pl = level - 1;
let pe = self.entry[pl].into();
let status = pool[self.node[pl]].inner_remove(pe);
self.heal_level(status, pl, pool);
match rhs_node {
Some(rhs) => self.node[level] = rhs,
None => self.size = 0,
}
false
}
fn right_sibling_branch_level(&self, level: usize, pool: &NodePool<F>) -> Option<usize> {
(0..level).rposition(|l| match pool[self.node[l]] {
NodeData::Inner { size, .. } => self.entry[l] < size,
_ => panic!("Expected inner node"),
})
}
fn left_sibling_branch_level(&self, level: usize) -> Option<usize> {
self.entry[0..level].iter().rposition(|&e| e != 0)
}
fn right_sibling(&self, level: usize, pool: &NodePool<F>) -> Option<(F::Key, Node)> {
self.right_sibling_branch_level(level, pool).map(|bl| {
let be = usize::from(self.entry[bl]);
let crit_key;
let mut node;
{
let (keys, tree) = pool[self.node[bl]].unwrap_inner();
crit_key = keys[be];
node = tree[be + 1];
}
for _ in bl + 1..level {
node = pool[node].unwrap_inner().1[0];
}
(crit_key, node)
})
}
fn update_right_crit_key(&self, level: usize, crit_key: F::Key, pool: &mut NodePool<F>) {
let bl = self
.right_sibling_branch_level(level, pool)
.expect("No right sibling exists");
match pool[self.node[bl]] {
NodeData::Inner { ref mut keys, .. } => {
keys[usize::from(self.entry[bl])] = crit_key;
}
_ => panic!("Expected inner node"),
}
}
pub fn normalize(&mut self, pool: &mut NodePool<F>) {
if let Some((leaf, entry)) = self.leaf_pos() {
if entry >= pool[leaf].entries() {
let leaf_level = self.size - 1;
self.next_node(leaf_level, pool);
}
}
}
}
#[cfg(test)]
impl<F: Forest> Path<F> {
pub fn verify(&self, pool: &NodePool<F>) {
for level in 0..self.size {
match pool[self.node[level]] {
NodeData::Inner { size, tree, .. } => {
assert!(
level < self.size - 1,
"Expected leaf node at level {}",
level
);
assert!(
self.entry[level] <= size,
"OOB inner entry {}/{} at level {}",
self.entry[level],
size,
level
);
assert_eq!(
self.node[level + 1],
tree[usize::from(self.entry[level])],
"Node mismatch at level {}",
level
);
}
NodeData::Leaf { size, .. } => {
assert_eq!(level, self.size - 1, "Expected inner node");
assert!(
self.entry[level] <= size,
"OOB leaf entry {}/{}",
self.entry[level],
size,
);
}
NodeData::Free { .. } => {
panic!("Free {} in path", self.node[level]);
}
}
}
}
}
#[cfg(test)]
impl<F: Forest> fmt::Display for Path<F> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
if self.size == 0 {
write!(f, "<empty path>")
} else {
write!(f, "{}[{}]", self.node[0], self.entry[0])?;
for i in 1..self.size {
write!(f, "--{}[{}]", self.node[i], self.entry[i])?;
}
Ok(())
}
}
}
#[cfg(test)]
mod tests {
use super::super::{Forest, NodeData, NodePool};
use super::*;
use core::cmp::Ordering;
struct TC();
impl Comparator<i32> for TC {
fn cmp(&self, a: i32, b: i32) -> Ordering {
a.cmp(&b)
}
}
struct TF();
impl Forest for TF {
type Key = i32;
type Value = char;
type LeafKeys = [i32; 7];
type LeafValues = [char; 7];
fn splat_key(key: Self::Key) -> Self::LeafKeys {
[key; 7]
}
fn splat_value(value: Self::Value) -> Self::LeafValues {
[value; 7]
}
}
#[test]
fn search_single_leaf() {
let mut pool = NodePool::<TF>::new();
let root = pool.alloc_node(NodeData::leaf(10, 'a'));
let mut p = Path::default();
let comp = TC();
assert_eq!(p.find(5, root, &pool, &comp), None);
assert_eq!(p.size, 1);
assert_eq!(p.node[0], root);
assert_eq!(p.entry[0], 0);
assert_eq!(p.find(10, root, &pool, &comp), Some('a'));
assert_eq!(p.size, 1);
assert_eq!(p.node[0], root);
assert_eq!(p.entry[0], 0);
assert_eq!(p.find(15, root, &pool, &comp), None);
assert_eq!(p.size, 1);
assert_eq!(p.node[0], root);
assert_eq!(p.entry[0], 1);
match pool[root] {
NodeData::Leaf {
ref mut size,
ref mut keys,
ref mut vals,
} => {
*size = 2;
keys[1] = 20;
vals[1] = 'b';
}
_ => unreachable!(),
}
assert_eq!(p.find(15, root, &pool, &comp), None);
assert_eq!(p.size, 1);
assert_eq!(p.node[0], root);
assert_eq!(p.entry[0], 1);
assert_eq!(p.find(25, root, &pool, &comp), None);
assert_eq!(p.size, 1);
assert_eq!(p.node[0], root);
assert_eq!(p.entry[0], 2);
}
#[test]
fn search_single_inner() {
let mut pool = NodePool::<TF>::new();
let leaf1 = pool.alloc_node(NodeData::leaf(10, 'a'));
let leaf2 = pool.alloc_node(NodeData::leaf(20, 'b'));
let root = pool.alloc_node(NodeData::inner(leaf1, 20, leaf2));
let mut p = Path::default();
let comp = TC();
assert_eq!(p.find(5, root, &pool, &comp), None);
assert_eq!(p.size, 2);
assert_eq!(p.node[0], root);
assert_eq!(p.entry[0], 0);
assert_eq!(p.node[1], leaf1);
assert_eq!(p.entry[1], 0);
assert_eq!(p.find(10, root, &pool, &comp), Some('a'));
assert_eq!(p.size, 2);
assert_eq!(p.node[0], root);
assert_eq!(p.entry[0], 0);
assert_eq!(p.node[1], leaf1);
assert_eq!(p.entry[1], 0);
assert_eq!(p.find(15, root, &pool, &comp), None);
assert_eq!(p.size, 2);
assert_eq!(p.node[0], root);
assert_eq!(p.entry[0], 0);
assert_eq!(p.node[1], leaf1);
assert_eq!(p.entry[1], 1);
assert_eq!(p.find(20, root, &pool, &comp), Some('b'));
assert_eq!(p.size, 2);
assert_eq!(p.node[0], root);
assert_eq!(p.entry[0], 1);
assert_eq!(p.node[1], leaf2);
assert_eq!(p.entry[1], 0);
assert_eq!(p.find(25, root, &pool, &comp), None);
assert_eq!(p.size, 2);
assert_eq!(p.node[0], root);
assert_eq!(p.entry[0], 1);
assert_eq!(p.node[1], leaf2);
assert_eq!(p.entry[1], 1);
}
}