use super::{slice_insert, slice_shift, Forest, Node, SetValue, INNER_SIZE};
use core::borrow::{Borrow, BorrowMut};
use core::fmt;
#[allow(dead_code)]
pub(super) enum NodeData<F: Forest> {
Inner {
size: u8,
keys: [F::Key; INNER_SIZE - 1],
tree: [Node; INNER_SIZE],
},
Leaf {
size: u8,
keys: F::LeafKeys,
vals: F::LeafValues,
},
Free { next: Option<Node> },
}
impl<F: Forest> Copy for NodeData<F> {}
impl<F: Forest> Clone for NodeData<F> {
fn clone(&self) -> Self {
*self
}
}
impl<F: Forest> NodeData<F> {
pub fn is_free(&self) -> bool {
match *self {
Self::Free { .. } => true,
_ => false,
}
}
pub fn entries(&self) -> usize {
match *self {
Self::Inner { size, .. } => usize::from(size) + 1,
Self::Leaf { size, .. } => usize::from(size),
Self::Free { .. } => panic!("freed node"),
}
}
pub fn inner(left: Node, key: F::Key, right: Node) -> Self {
let mut tree = [right; INNER_SIZE];
tree[0] = left;
Self::Inner {
size: 1,
keys: [key; INNER_SIZE - 1],
tree,
}
}
pub fn leaf(key: F::Key, value: F::Value) -> Self {
Self::Leaf {
size: 1,
keys: F::splat_key(key),
vals: F::splat_value(value),
}
}
pub fn unwrap_inner(&self) -> (&[F::Key], &[Node]) {
match *self {
Self::Inner {
size,
ref keys,
ref tree,
} => {
let size = usize::from(size);
(&keys[0..size], &tree[0..=size])
}
_ => panic!("Expected inner node"),
}
}
pub fn unwrap_leaf(&self) -> (&[F::Key], &[F::Value]) {
match *self {
Self::Leaf {
size,
ref keys,
ref vals,
} => {
let size = usize::from(size);
let keys = keys.borrow();
let vals = vals.borrow();
(&keys[0..size], &vals[0..size])
}
_ => panic!("Expected leaf node"),
}
}
pub fn unwrap_leaf_mut(&mut self) -> (&mut [F::Key], &mut [F::Value]) {
match *self {
Self::Leaf {
size,
ref mut keys,
ref mut vals,
} => {
let size = usize::from(size);
let keys = keys.borrow_mut();
let vals = vals.borrow_mut();
(&mut keys[0..size], &mut vals[0..size])
}
_ => panic!("Expected leaf node"),
}
}
pub fn leaf_crit_key(&self) -> F::Key {
match *self {
Self::Leaf { size, ref keys, .. } => {
debug_assert!(size > 0, "Empty leaf node");
keys.borrow()[0]
}
_ => panic!("Expected leaf node"),
}
}
pub fn try_inner_insert(&mut self, index: usize, key: F::Key, node: Node) -> bool {
match *self {
Self::Inner {
ref mut size,
ref mut keys,
ref mut tree,
} => {
let sz = usize::from(*size);
debug_assert!(sz <= keys.len());
debug_assert!(index <= sz, "Can't insert at {} with {} keys", index, sz);
if let Some(ks) = keys.get_mut(0..=sz) {
*size = (sz + 1) as u8;
slice_insert(ks, index, key);
slice_insert(&mut tree[1..=sz + 1], index, node);
true
} else {
false
}
}
_ => panic!("Expected inner node"),
}
}
pub fn try_leaf_insert(&mut self, index: usize, key: F::Key, value: F::Value) -> bool {
match *self {
Self::Leaf {
ref mut size,
ref mut keys,
ref mut vals,
} => {
let sz = usize::from(*size);
let keys = keys.borrow_mut();
let vals = vals.borrow_mut();
debug_assert!(sz <= keys.len());
debug_assert!(index <= sz);
if let Some(ks) = keys.get_mut(0..=sz) {
*size = (sz + 1) as u8;
slice_insert(ks, index, key);
slice_insert(&mut vals[0..=sz], index, value);
true
} else {
false
}
}
_ => panic!("Expected leaf node"),
}
}
pub fn split(&mut self, insert_index: usize) -> SplitOff<F> {
match *self {
Self::Inner {
ref mut size,
ref keys,
ref tree,
} => {
debug_assert_eq!(usize::from(*size), keys.len(), "Node not full");
let l_ents = split_pos(tree.len(), insert_index + 1);
let r_ents = tree.len() - l_ents;
*size = (l_ents - 1) as u8;
let mut r_keys = *keys;
r_keys[0..r_ents - 1].copy_from_slice(&keys[l_ents..]);
let mut r_tree = *tree;
r_tree[0..r_ents].copy_from_slice(&tree[l_ents..]);
SplitOff {
lhs_entries: l_ents,
rhs_entries: r_ents,
crit_key: keys[l_ents - 1],
rhs_data: Self::Inner {
size: (r_ents - 1) as u8,
keys: r_keys,
tree: r_tree,
},
}
}
Self::Leaf {
ref mut size,
ref keys,
ref vals,
} => {
let o_keys = keys.borrow();
let o_vals = vals.borrow();
debug_assert_eq!(usize::from(*size), o_keys.len(), "Node not full");
let l_size = split_pos(o_keys.len(), insert_index);
let r_size = o_keys.len() - l_size;
*size = l_size as u8;
let mut r_keys = *keys;
r_keys.borrow_mut()[0..r_size].copy_from_slice(&o_keys[l_size..]);
let mut r_vals = *vals;
r_vals.borrow_mut()[0..r_size].copy_from_slice(&o_vals[l_size..]);
SplitOff {
lhs_entries: l_size,
rhs_entries: r_size,
crit_key: o_keys[l_size],
rhs_data: Self::Leaf {
size: r_size as u8,
keys: r_keys,
vals: r_vals,
},
}
}
_ => panic!("Expected leaf node"),
}
}
pub fn inner_remove(&mut self, index: usize) -> Removed {
match *self {
Self::Inner {
ref mut size,
ref mut keys,
ref mut tree,
} => {
let ents = usize::from(*size) + 1;
debug_assert!(ents <= tree.len());
debug_assert!(index < ents);
*size = ents.wrapping_sub(2) as u8;
if ents > 1 {
slice_shift(&mut keys[index.saturating_sub(1)..ents - 1], 1);
}
slice_shift(&mut tree[index..ents], 1);
Removed::new(index, ents - 1, tree.len())
}
_ => panic!("Expected inner node"),
}
}
pub fn leaf_remove(&mut self, index: usize) -> Removed {
match *self {
Self::Leaf {
ref mut size,
ref mut keys,
ref mut vals,
} => {
let sz = usize::from(*size);
let keys = keys.borrow_mut();
let vals = vals.borrow_mut();
*size -= 1;
slice_shift(&mut keys[index..sz], 1);
slice_shift(&mut vals[index..sz], 1);
Removed::new(index, sz - 1, keys.len())
}
_ => panic!("Expected leaf node"),
}
}
pub fn balance(&mut self, crit_key: F::Key, rhs: &mut Self) -> Option<F::Key> {
match (self, rhs) {
(
&mut Self::Inner {
size: ref mut l_size,
keys: ref mut l_keys,
tree: ref mut l_tree,
},
&mut Self::Inner {
size: ref mut r_size,
keys: ref mut r_keys,
tree: ref mut r_tree,
},
) => {
let l_ents = usize::from(*l_size) + 1;
let r_ents = usize::from(*r_size) + 1;
let ents = l_ents + r_ents;
if ents <= r_tree.len() {
*l_size = 0;
l_keys[l_ents - 1] = crit_key;
l_keys[l_ents..ents - 1].copy_from_slice(&r_keys[0..r_ents - 1]);
r_keys[0..ents - 1].copy_from_slice(&l_keys[0..ents - 1]);
l_tree[l_ents..ents].copy_from_slice(&r_tree[0..r_ents]);
r_tree[0..ents].copy_from_slice(&l_tree[0..ents]);
*r_size = (ents - 1) as u8;
None
} else {
let r_goal = ents / 2;
let l_goal = ents - r_goal;
debug_assert!(l_goal > l_ents, "Node must be underflowed");
l_keys[l_ents - 1] = crit_key;
l_keys[l_ents..l_goal - 1].copy_from_slice(&r_keys[0..l_goal - 1 - l_ents]);
l_tree[l_ents..l_goal].copy_from_slice(&r_tree[0..l_goal - l_ents]);
*l_size = (l_goal - 1) as u8;
let new_crit = r_keys[r_ents - r_goal - 1];
slice_shift(&mut r_keys[0..r_ents - 1], r_ents - r_goal);
slice_shift(&mut r_tree[0..r_ents], r_ents - r_goal);
*r_size = (r_goal - 1) as u8;
Some(new_crit)
}
}
(
&mut Self::Leaf {
size: ref mut l_size,
keys: ref mut l_keys,
vals: ref mut l_vals,
},
&mut Self::Leaf {
size: ref mut r_size,
keys: ref mut r_keys,
vals: ref mut r_vals,
},
) => {
let l_ents = usize::from(*l_size);
let l_keys = l_keys.borrow_mut();
let l_vals = l_vals.borrow_mut();
let r_ents = usize::from(*r_size);
let r_keys = r_keys.borrow_mut();
let r_vals = r_vals.borrow_mut();
let ents = l_ents + r_ents;
if ents <= r_vals.len() {
*l_size = 0;
l_keys[l_ents..ents].copy_from_slice(&r_keys[0..r_ents]);
r_keys[0..ents].copy_from_slice(&l_keys[0..ents]);
l_vals[l_ents..ents].copy_from_slice(&r_vals[0..r_ents]);
r_vals[0..ents].copy_from_slice(&l_vals[0..ents]);
*r_size = ents as u8;
None
} else {
let r_goal = ents / 2;
let l_goal = ents - r_goal;
debug_assert!(l_goal > l_ents, "Node must be underflowed");
l_keys[l_ents..l_goal].copy_from_slice(&r_keys[0..l_goal - l_ents]);
l_vals[l_ents..l_goal].copy_from_slice(&r_vals[0..l_goal - l_ents]);
*l_size = l_goal as u8;
slice_shift(&mut r_keys[0..r_ents], r_ents - r_goal);
slice_shift(&mut r_vals[0..r_ents], r_ents - r_goal);
*r_size = r_goal as u8;
Some(r_keys[0])
}
}
_ => panic!("Mismatched nodes"),
}
}
}
fn split_pos(len: usize, ins: usize) -> usize {
if ins <= len / 2 {
len / 2
} else {
(len + 1) / 2
}
}
pub(super) struct SplitOff<F: Forest> {
pub lhs_entries: usize,
pub rhs_entries: usize,
pub crit_key: F::Key,
pub rhs_data: NodeData<F>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum Removed {
Healthy,
Rightmost,
Underflow,
Empty,
}
impl Removed {
fn new(removed: usize, new_size: usize, capacity: usize) -> Self {
if 2 * new_size >= capacity {
if removed == new_size {
Self::Rightmost
} else {
Self::Healthy
}
} else if new_size > 0 {
Self::Underflow
} else {
Self::Empty
}
}
}
pub(super) trait ValDisp {
fn valfmt(&self, f: &mut fmt::Formatter) -> fmt::Result;
}
impl ValDisp for SetValue {
fn valfmt(&self, _: &mut fmt::Formatter) -> fmt::Result {
Ok(())
}
}
impl<T: fmt::Display> ValDisp for T {
fn valfmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, ":{}", self)
}
}
impl<F> fmt::Display for NodeData<F>
where
F: Forest,
F::Key: fmt::Display,
F::Value: ValDisp,
{
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match *self {
Self::Inner { size, keys, tree } => {
write!(f, "[ {}", tree[0])?;
for i in 0..usize::from(size) {
write!(f, " {} {}", keys[i], tree[i + 1])?;
}
write!(f, " ]")
}
Self::Leaf { size, keys, vals } => {
let keys = keys.borrow();
let vals = vals.borrow();
write!(f, "[")?;
for i in 0..usize::from(size) {
write!(f, " {}", keys[i])?;
vals[i].valfmt(f)?;
}
write!(f, " ]")
}
Self::Free { next: Some(n) } => write!(f, "[ free -> {} ]", n),
Self::Free { next: None } => write!(f, "[ free ]"),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::string::ToString;
use core::mem;
struct TF();
impl Forest for TF {
type Key = char;
type Value = SetValue;
type LeafKeys = [char; 15];
type LeafValues = [SetValue; 15];
fn splat_key(key: Self::Key) -> Self::LeafKeys {
[key; 15]
}
fn splat_value(value: Self::Value) -> Self::LeafValues {
[value; 15]
}
}
#[test]
fn inner() {
let n1 = Node(1);
let n2 = Node(2);
let n3 = Node(3);
let n4 = Node(4);
let mut inner = NodeData::<TF>::inner(n1, 'c', n4);
assert_eq!(mem::size_of_val(&inner), 64);
assert_eq!(inner.to_string(), "[ node1 c node4 ]");
assert!(inner.try_inner_insert(0, 'a', n2));
assert_eq!(inner.to_string(), "[ node1 a node2 c node4 ]");
assert!(inner.try_inner_insert(1, 'b', n3));
assert_eq!(inner.to_string(), "[ node1 a node2 b node3 c node4 ]");
for i in 3..7 {
assert!(inner.try_inner_insert(
usize::from(i),
('a' as u8 + i) as char,
Node(i as u32 + 2),
));
}
assert_eq!(
inner.to_string(),
"[ node1 a node2 b node3 c node4 d node5 e node6 f node7 g node8 ]"
);
assert!(!inner.try_inner_insert(0, 'x', n3));
assert!(!inner.try_inner_insert(4, 'x', n3));
assert!(!inner.try_inner_insert(7, 'x', n3));
let saved = inner.clone();
let sp = inner.split(1);
assert_eq!(sp.lhs_entries, 4);
assert_eq!(sp.rhs_entries, 4);
assert_eq!(sp.crit_key, 'd');
assert_eq!(inner.to_string(), "[ node1 a node2 b node3 c node4 ]");
assert_eq!(sp.rhs_data.to_string(), "[ node5 e node6 f node7 g node8 ]");
assert_eq!(inner.inner_remove(0), Removed::Underflow);
assert_eq!(inner.to_string(), "[ node2 b node3 c node4 ]");
assert_eq!(inner.inner_remove(1), Removed::Underflow);
assert_eq!(inner.to_string(), "[ node2 c node4 ]");
assert_eq!(inner.inner_remove(1), Removed::Underflow);
assert_eq!(inner.to_string(), "[ node2 ]");
assert_eq!(inner.inner_remove(0), Removed::Empty);
inner = saved;
let sp = inner.split(6);
assert_eq!(sp.lhs_entries, 4);
assert_eq!(sp.rhs_entries, 4);
assert_eq!(sp.crit_key, 'd');
assert_eq!(inner.to_string(), "[ node1 a node2 b node3 c node4 ]");
assert_eq!(sp.rhs_data.to_string(), "[ node5 e node6 f node7 g node8 ]");
}
#[test]
fn leaf() {
let mut leaf = NodeData::<TF>::leaf('d', SetValue());
assert_eq!(leaf.to_string(), "[ d ]");
assert!(leaf.try_leaf_insert(0, 'a', SetValue()));
assert_eq!(leaf.to_string(), "[ a d ]");
assert!(leaf.try_leaf_insert(1, 'b', SetValue()));
assert!(leaf.try_leaf_insert(2, 'c', SetValue()));
assert_eq!(leaf.to_string(), "[ a b c d ]");
for i in 4..15 {
assert!(leaf.try_leaf_insert(usize::from(i), ('a' as u8 + i) as char, SetValue()));
}
assert_eq!(leaf.to_string(), "[ a b c d e f g h i j k l m n o ]");
assert!(!leaf.try_leaf_insert(0, 'x', SetValue()));
assert!(!leaf.try_leaf_insert(8, 'x', SetValue()));
assert!(!leaf.try_leaf_insert(15, 'x', SetValue()));
let saved = leaf.clone();
let sp = leaf.split(12);
assert_eq!(sp.lhs_entries, 8);
assert_eq!(sp.rhs_entries, 7);
assert_eq!(sp.crit_key, 'i');
assert_eq!(leaf.to_string(), "[ a b c d e f g h ]");
assert_eq!(sp.rhs_data.to_string(), "[ i j k l m n o ]");
assert!(leaf.try_leaf_insert(8, 'i', SetValue()));
assert_eq!(leaf.leaf_remove(2), Removed::Healthy);
assert_eq!(leaf.to_string(), "[ a b d e f g h i ]");
assert_eq!(leaf.leaf_remove(7), Removed::Underflow);
assert_eq!(leaf.to_string(), "[ a b d e f g h ]");
leaf = saved;
let sp = leaf.split(7);
assert_eq!(sp.lhs_entries, 7);
assert_eq!(sp.rhs_entries, 8);
assert_eq!(sp.crit_key, 'h');
assert_eq!(leaf.to_string(), "[ a b c d e f g ]");
assert_eq!(sp.rhs_data.to_string(), "[ h i j k l m n o ]");
}
#[test]
fn optimal_split_pos() {
assert_eq!(split_pos(8, 0), 4);
assert_eq!(split_pos(8, 8), 4);
assert_eq!(split_pos(7, 0), 3);
assert_eq!(split_pos(7, 7), 4);
assert_eq!(split_pos(7, 3), 3);
assert_eq!(split_pos(7, 4), 4);
}
#[test]
fn inner_balance() {
let n1 = Node(1);
let n2 = Node(2);
let n3 = Node(3);
let mut lhs = NodeData::<TF>::inner(n1, 'a', n2);
assert!(lhs.try_inner_insert(1, 'b', n3));
assert_eq!(lhs.to_string(), "[ node1 a node2 b node3 ]");
let n11 = Node(11);
let n12 = Node(12);
let mut rhs = NodeData::<TF>::inner(n11, 'p', n12);
for i in 1..4 {
assert!(rhs.try_inner_insert(
usize::from(i),
('p' as u8 + i) as char,
Node(i as u32 + 12),
));
}
assert_eq!(
rhs.to_string(),
"[ node11 p node12 q node13 r node14 s node15 ]"
);
assert_eq!(lhs.balance('o', &mut rhs), None);
assert_eq!(
rhs.to_string(),
"[ node1 a node2 b node3 o node11 p node12 q node13 r node14 s node15 ]"
);
lhs = NodeData::<TF>::inner(Node(20), 'x', Node(21));
assert_eq!(lhs.balance('y', &mut rhs), Some('o'));
assert_eq!(
lhs.to_string(),
"[ node20 x node21 y node1 a node2 b node3 ]"
);
assert_eq!(
rhs.to_string(),
"[ node11 p node12 q node13 r node14 s node15 ]"
);
}
#[test]
fn leaf_balance() {
let mut lhs = NodeData::<TF>::leaf('a', SetValue());
for i in 1..6 {
assert!(lhs.try_leaf_insert(usize::from(i), ('a' as u8 + i) as char, SetValue()));
}
assert_eq!(lhs.to_string(), "[ a b c d e f ]");
let mut rhs = NodeData::<TF>::leaf('0', SetValue());
for i in 1..8 {
assert!(rhs.try_leaf_insert(usize::from(i), ('0' as u8 + i) as char, SetValue()));
}
assert_eq!(rhs.to_string(), "[ 0 1 2 3 4 5 6 7 ]");
assert_eq!(lhs.balance('0', &mut rhs), None);
assert_eq!(rhs.to_string(), "[ a b c d e f 0 1 2 3 4 5 6 7 ]");
assert!(lhs.try_leaf_insert(0, 'x', SetValue()));
assert!(lhs.try_leaf_insert(1, 'y', SetValue()));
assert!(lhs.try_leaf_insert(2, 'z', SetValue()));
assert_eq!(lhs.to_string(), "[ x y z ]");
assert_eq!(lhs.balance('a', &mut rhs), Some('0'));
assert_eq!(lhs.to_string(), "[ x y z a b c d e f ]");
assert_eq!(rhs.to_string(), "[ 0 1 2 3 4 5 6 7 ]");
}
}