use alloc::format;
use alloc::vec::Vec;
use crate::domain::{BoundingBox, Cut};
use crate::error::{RcfError, RcfResult};
use crate::tree::node::{InternalData, LeafData, NodeRef, NodeView, NodeViewMut};
#[derive(Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct NodeStore<const D: usize> {
internals: Vec<Option<InternalData<D>>>,
leaves: Vec<Option<LeafData>>,
internal_free: Vec<u32>,
leaf_free: Vec<u32>,
capacity: u32,
}
impl<const D: usize> NodeStore<D> {
pub fn new(capacity: u32) -> RcfResult<Self> {
if capacity == 0 {
return Err(RcfError::InvalidConfig(
"NodeStore capacity must be > 0".into(),
));
}
if capacity > NodeRef::MAX_INDEX {
return Err(RcfError::InvalidConfig(
format!(
"NodeStore capacity {capacity} exceeds NodeRef::MAX_INDEX {}",
NodeRef::MAX_INDEX
)
.into(),
));
}
let cap = capacity as usize;
let internals = (0..cap).map(|_| None).collect();
let leaves = (0..cap).map(|_| None).collect();
let internal_free = (0..capacity).rev().collect();
let leaf_free = (0..capacity).rev().collect();
Ok(Self {
internals,
leaves,
internal_free,
leaf_free,
capacity,
})
}
#[must_use]
pub fn capacity(&self) -> u32 {
self.capacity
}
#[must_use]
pub fn live_count(&self) -> usize {
self.live_internal_count() + self.live_leaf_count()
}
#[must_use]
pub fn live_internal_count(&self) -> usize {
self.capacity as usize - self.internal_free.len()
}
#[must_use]
pub fn live_leaf_count(&self) -> usize {
self.capacity as usize - self.leaf_free.len()
}
pub fn add_internal(
&mut self,
cut: Cut,
bbox: BoundingBox<D>,
left: NodeRef,
right: NodeRef,
parent: Option<NodeRef>,
mass: u64,
) -> RcfResult<NodeRef> {
let idx = self
.internal_free
.pop()
.ok_or_else(|| RcfError::InvalidConfig("NodeStore internal arena exhausted".into()))?;
self.internals[idx as usize] = Some(InternalData {
cut,
bbox,
left,
right,
parent,
mass,
});
Ok(NodeRef::internal(idx))
}
pub fn add_leaf(
&mut self,
point_idx: usize,
parent: Option<NodeRef>,
mass: u64,
) -> RcfResult<NodeRef> {
let idx = self
.leaf_free
.pop()
.ok_or_else(|| RcfError::InvalidConfig("NodeStore leaf arena exhausted".into()))?;
self.leaves[idx as usize] = Some(LeafData {
point_idx,
parent,
mass,
});
Ok(NodeRef::leaf(idx))
}
pub fn delete(&mut self, n: NodeRef) -> RcfResult<()> {
let idx = n.index();
if idx >= self.capacity as usize {
return Err(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
});
}
let was_live = if n.is_leaf() {
self.leaves[idx].take().is_some()
} else {
self.internals[idx].take().is_some()
};
if !was_live {
return Err(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
});
}
if n.is_leaf() {
#[allow(clippy::cast_possible_truncation)]
self.leaf_free.push(idx as u32);
} else {
#[allow(clippy::cast_possible_truncation)]
self.internal_free.push(idx as u32);
}
Ok(())
}
pub fn view(&self, n: NodeRef) -> RcfResult<NodeView<'_, D>> {
let idx = n.index();
if idx >= self.capacity as usize {
return Err(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
});
}
if n.is_leaf() {
self.leaves[idx]
.as_ref()
.map(NodeView::Leaf)
.ok_or(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
})
} else {
self.internals[idx]
.as_ref()
.map(NodeView::Internal)
.ok_or(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
})
}
}
pub fn view_mut(&mut self, n: NodeRef) -> RcfResult<NodeViewMut<'_, D>> {
let idx = n.index();
if idx >= self.capacity as usize {
return Err(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
});
}
if n.is_leaf() {
self.leaves[idx]
.as_mut()
.map(NodeViewMut::Leaf)
.ok_or(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
})
} else {
self.internals[idx]
.as_mut()
.map(NodeViewMut::Internal)
.ok_or(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
})
}
}
pub fn internal(&self, n: NodeRef) -> RcfResult<&InternalData<D>> {
if n.is_leaf() {
return Err(RcfError::InvalidConfig(
"NodeStore::internal: called on a leaf reference".into(),
));
}
let idx = n.index();
if idx >= self.capacity as usize {
return Err(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
});
}
self.internals[idx].as_ref().ok_or(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
})
}
pub fn internal_mut(&mut self, n: NodeRef) -> RcfResult<&mut InternalData<D>> {
if n.is_leaf() {
return Err(RcfError::InvalidConfig(
"NodeStore::internal_mut: called on a leaf reference".into(),
));
}
let idx = n.index();
if idx >= self.capacity as usize {
return Err(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
});
}
self.internals[idx].as_mut().ok_or(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
})
}
pub fn leaf(&self, n: NodeRef) -> RcfResult<&LeafData> {
if !n.is_leaf() {
return Err(RcfError::InvalidConfig(
"NodeStore::leaf: called on an internal reference".into(),
));
}
let idx = n.index();
if idx >= self.capacity as usize {
return Err(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
});
}
self.leaves[idx].as_ref().ok_or(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
})
}
pub fn leaf_mut(&mut self, n: NodeRef) -> RcfResult<&mut LeafData> {
if !n.is_leaf() {
return Err(RcfError::InvalidConfig(
"NodeStore::leaf_mut: called on an internal reference".into(),
));
}
let idx = n.index();
if idx >= self.capacity as usize {
return Err(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
});
}
self.leaves[idx].as_mut().ok_or(RcfError::OutOfBounds {
index: idx,
len: self.capacity as usize,
})
}
pub fn parent(&self, n: NodeRef) -> RcfResult<Option<NodeRef>> {
Ok(self.view(n)?.parent())
}
pub fn sibling(&self, n: NodeRef) -> RcfResult<Option<NodeRef>> {
let Some(parent_ref) = self.parent(n)? else {
return Ok(None);
};
if parent_ref.is_leaf() {
return Err(RcfError::InvalidConfig(
"NodeStore::sibling: parent is a leaf — invariant violated".into(),
));
}
let parent = self.internal(parent_ref)?;
let n_raw = n.raw();
if parent.left.raw() == n_raw {
Ok(Some(parent.right))
} else if parent.right.raw() == n_raw {
Ok(Some(parent.left))
} else {
Err(RcfError::InvalidConfig(
"NodeStore::sibling: child not registered with parent".into(),
))
}
}
pub fn internal_bbox(&self, n: NodeRef) -> RcfResult<&BoundingBox<D>> {
Ok(&self.internal(n)?.bbox)
}
pub fn set_parent(&mut self, n: NodeRef, parent: Option<NodeRef>) -> RcfResult<()> {
match self.view_mut(n)? {
NodeViewMut::Internal(i) => {
i.parent = parent;
}
NodeViewMut::Leaf(l) => {
l.parent = parent;
}
}
Ok(())
}
pub fn set_mass(&mut self, n: NodeRef, mass: u64) -> RcfResult<()> {
match self.view_mut(n)? {
NodeViewMut::Internal(i) => {
i.mass = mass;
}
NodeViewMut::Leaf(l) => {
l.mass = mass;
}
}
Ok(())
}
pub fn set_internal_bbox(&mut self, n: NodeRef, bbox: BoundingBox<D>) -> RcfResult<()> {
self.internal_mut(n)?.bbox = bbox;
Ok(())
}
pub fn set_internal_children(
&mut self,
n: NodeRef,
new_left: NodeRef,
new_right: NodeRef,
) -> RcfResult<()> {
let i = self.internal_mut(n)?;
i.left = new_left;
i.right = new_right;
Ok(())
}
pub fn set_internal_cut(&mut self, n: NodeRef, new_cut: Cut) -> RcfResult<()> {
self.internal_mut(n)?.cut = new_cut;
Ok(())
}
}
#[cfg(test)]
#[allow(clippy::float_cmp)] mod tests {
use super::*;
use proptest::prelude::*;
fn unit_bbox<const D: usize>() -> BoundingBox<D> {
let mut b = BoundingBox::<D>::from_point(&vec![0.0; D]).unwrap();
b.extend(&vec![1.0; D]).unwrap();
b
}
#[test]
fn new_rejects_zero_capacity() {
assert!(matches!(
NodeStore::<2>::new(0).unwrap_err(),
RcfError::InvalidConfig(_)
));
}
#[test]
fn new_rejects_capacity_above_max() {
assert!(matches!(
NodeStore::<2>::new(u32::MAX).unwrap_err(),
RcfError::InvalidConfig(_)
));
}
#[test]
fn new_starts_empty() {
let s = NodeStore::<2>::new(8).unwrap();
assert_eq!(s.capacity(), 8);
assert_eq!(s.live_count(), 0);
assert_eq!(s.live_internal_count(), 0);
assert_eq!(s.live_leaf_count(), 0);
}
#[test]
fn add_leaf_increments_live() {
let mut s = NodeStore::<2>::new(4).unwrap();
let l = s.add_leaf(7, None, 1).unwrap();
assert!(l.is_leaf());
assert_eq!(s.live_leaf_count(), 1);
assert_eq!(s.live_internal_count(), 0);
}
#[test]
fn add_internal_increments_live() {
let mut s = NodeStore::<2>::new(4).unwrap();
let l1 = s.add_leaf(0, None, 1).unwrap();
let l2 = s.add_leaf(1, None, 1).unwrap();
let cut = Cut::new(0, 0.5);
let i = s
.add_internal(cut, unit_bbox::<2>(), l1, l2, None, 2)
.unwrap();
assert!(i.is_internal());
assert_eq!(s.live_internal_count(), 1);
assert_eq!(s.live_leaf_count(), 2);
assert_eq!(s.live_count(), 3);
}
#[test]
fn add_leaf_capacity_exhausted() {
let mut s = NodeStore::<2>::new(2).unwrap();
s.add_leaf(0, None, 1).unwrap();
s.add_leaf(1, None, 1).unwrap();
assert!(matches!(
s.add_leaf(2, None, 1).unwrap_err(),
RcfError::InvalidConfig(_)
));
}
#[test]
fn add_internal_capacity_exhausted() {
let mut s = NodeStore::<2>::new(1).unwrap();
let l1 = s.add_leaf(0, None, 1).unwrap();
let l2 = s.add_leaf(1, None, 1);
assert!(matches!(l2.unwrap_err(), RcfError::InvalidConfig(_)));
let i = s
.add_internal(Cut::new(0, 0.0), unit_bbox::<2>(), l1, l1, None, 1)
.unwrap();
assert!(
s.add_internal(Cut::new(0, 0.0), unit_bbox::<2>(), i, i, None, 1)
.is_err()
);
}
#[test]
fn delete_frees_slot_for_reuse() {
let mut s = NodeStore::<2>::new(2).unwrap();
let l = s.add_leaf(7, None, 1).unwrap();
let l_idx = l.index();
s.delete(l).unwrap();
assert_eq!(s.live_leaf_count(), 0);
let l2 = s.add_leaf(99, None, 1).unwrap();
assert_eq!(l2.index(), l_idx);
}
#[test]
fn delete_oob_index_rejected() {
let mut s = NodeStore::<2>::new(2).unwrap();
let bogus = NodeRef::leaf(99);
assert!(matches!(
s.delete(bogus).unwrap_err(),
RcfError::OutOfBounds { .. }
));
}
#[test]
fn delete_double_free_rejected() {
let mut s = NodeStore::<2>::new(2).unwrap();
let l = s.add_leaf(0, None, 1).unwrap();
s.delete(l).unwrap();
assert!(matches!(
s.delete(l).unwrap_err(),
RcfError::OutOfBounds { .. }
));
}
#[test]
fn leaf_returns_inserted_value() {
let mut s = NodeStore::<2>::new(2).unwrap();
let l = s.add_leaf(42, None, 1).unwrap();
let data = s.leaf(l).unwrap();
assert_eq!(data.point_idx, 42);
assert_eq!(data.mass, 1);
}
#[test]
fn leaf_mut_allows_inplace_update() {
let mut s = NodeStore::<2>::new(2).unwrap();
let l = s.add_leaf(1, None, 1).unwrap();
s.leaf_mut(l).unwrap().mass = 5;
assert_eq!(s.view(l).unwrap().mass(), 5);
}
#[test]
fn view_oob_returns_err() {
let s = NodeStore::<2>::new(2).unwrap();
assert!(matches!(
s.view(NodeRef::leaf(7)).unwrap_err(),
RcfError::OutOfBounds { .. }
));
}
#[test]
fn parent_and_sibling_lookup() {
let mut s = NodeStore::<2>::new(4).unwrap();
let l1 = s.add_leaf(0, None, 1).unwrap();
let l2 = s.add_leaf(1, None, 1).unwrap();
let i = s
.add_internal(Cut::new(0, 0.5), unit_bbox::<2>(), l1, l2, None, 2)
.unwrap();
s.set_parent(l1, Some(i)).unwrap();
s.set_parent(l2, Some(i)).unwrap();
assert_eq!(s.parent(l1).unwrap(), Some(i));
assert_eq!(s.parent(i).unwrap(), None);
assert_eq!(s.sibling(l1).unwrap(), Some(l2));
assert_eq!(s.sibling(l2).unwrap(), Some(l1));
assert_eq!(s.sibling(i).unwrap(), None);
}
#[test]
fn sibling_orphan_detected() {
let mut s = NodeStore::<2>::new(4).unwrap();
let real_l = s.add_leaf(0, None, 1).unwrap();
let fake_l = s.add_leaf(1, None, 1).unwrap();
let other = s.add_leaf(2, None, 1).unwrap();
let i = s
.add_internal(Cut::new(0, 0.5), unit_bbox::<2>(), real_l, other, None, 2)
.unwrap();
s.set_parent(fake_l, Some(i)).unwrap();
assert!(matches!(
s.sibling(fake_l).unwrap_err(),
RcfError::InvalidConfig(_)
));
}
#[test]
fn sibling_parent_is_leaf_rejected() {
let mut s = NodeStore::<2>::new(4).unwrap();
let leaf_parent = s.add_leaf(0, None, 1).unwrap();
let child = s.add_leaf(1, None, 1).unwrap();
s.set_parent(child, Some(leaf_parent)).unwrap();
assert!(matches!(
s.sibling(child).unwrap_err(),
RcfError::InvalidConfig(_)
));
}
#[test]
fn internal_bbox_returns_cached_box() {
let mut s = NodeStore::<2>::new(2).unwrap();
let l1 = s.add_leaf(0, None, 1).unwrap();
let l2 = s.add_leaf(1, None, 1).unwrap();
let bbox = unit_bbox::<2>();
let i = s
.add_internal(Cut::new(0, 0.5), bbox.clone(), l1, l2, None, 2)
.unwrap();
assert_eq!(s.internal_bbox(i).unwrap(), &bbox);
}
#[test]
fn internal_bbox_on_leaf_rejected() {
let mut s = NodeStore::<2>::new(2).unwrap();
let l = s.add_leaf(0, None, 1).unwrap();
assert!(matches!(
s.internal_bbox(l).unwrap_err(),
RcfError::InvalidConfig(_)
));
}
#[test]
fn set_mass_updates_leaf_and_internal() {
let mut s = NodeStore::<2>::new(2).unwrap();
let l = s.add_leaf(0, None, 1).unwrap();
s.set_mass(l, 9).unwrap();
assert_eq!(s.view(l).unwrap().mass(), 9);
}
#[test]
fn set_internal_bbox_replaces_cached() {
let mut s = NodeStore::<2>::new(2).unwrap();
let l1 = s.add_leaf(0, None, 1).unwrap();
let l2 = s.add_leaf(1, None, 1).unwrap();
let i = s
.add_internal(Cut::new(0, 0.5), unit_bbox::<2>(), l1, l2, None, 2)
.unwrap();
let mut new_bbox = BoundingBox::<2>::from_point(&[0.0, 0.0]).unwrap();
new_bbox.extend(&[10.0, 10.0]).unwrap();
s.set_internal_bbox(i, new_bbox.clone()).unwrap();
assert_eq!(s.internal_bbox(i).unwrap(), &new_bbox);
}
#[test]
fn set_internal_children_swaps_refs() {
let mut s = NodeStore::<2>::new(4).unwrap();
let l1 = s.add_leaf(0, None, 1).unwrap();
let l2 = s.add_leaf(1, None, 1).unwrap();
let l3 = s.add_leaf(2, None, 1).unwrap();
let i = s
.add_internal(Cut::new(0, 0.5), unit_bbox::<2>(), l1, l2, None, 2)
.unwrap();
s.set_internal_children(i, l1, l3).unwrap();
let data = s.internal(i).unwrap();
assert_eq!(data.left, l1);
assert_eq!(data.right, l3);
}
#[test]
fn set_internal_cut_replaces_cut() {
let mut s = NodeStore::<2>::new(4).unwrap();
let l1 = s.add_leaf(0, None, 1).unwrap();
let l2 = s.add_leaf(1, None, 1).unwrap();
let i = s
.add_internal(Cut::new(0, 0.5), unit_bbox::<2>(), l1, l2, None, 2)
.unwrap();
s.set_internal_cut(i, Cut::new(1, 9.0)).unwrap();
let data = s.internal(i).unwrap();
assert_eq!(data.cut.dim(), 1);
assert_eq!(data.cut.value(), 9.0);
}
#[test]
fn setters_reject_oob_or_wrong_kind() {
let mut s = NodeStore::<2>::new(2).unwrap();
let l = s.add_leaf(0, None, 1).unwrap();
assert!(matches!(
s.set_internal_bbox(l, unit_bbox::<2>()).unwrap_err(),
RcfError::InvalidConfig(_)
));
assert!(matches!(
s.set_internal_children(l, l, l).unwrap_err(),
RcfError::InvalidConfig(_)
));
assert!(matches!(
s.set_internal_cut(l, Cut::new(0, 0.0)).unwrap_err(),
RcfError::InvalidConfig(_)
));
assert!(matches!(
s.set_parent(NodeRef::leaf(99), None).unwrap_err(),
RcfError::OutOfBounds { .. }
));
assert!(matches!(
s.set_mass(NodeRef::leaf(99), 1).unwrap_err(),
RcfError::OutOfBounds { .. }
));
}
proptest! {
#![proptest_config(ProptestConfig { cases: 64, ..ProptestConfig::default() })]
#[test]
fn invariants_under_random_ops(ops in proptest::collection::vec(0u32..20, 0..200)) {
const CAP: u32 = 16;
let mut s = NodeStore::<2>::new(CAP).unwrap();
let mut live_leaves: Vec<NodeRef> = Vec::new();
let mut live_internals: Vec<NodeRef> = Vec::new();
for op in ops {
match op % 4 {
0 => {
if let Ok(r) = s.add_leaf(0, None, 1) {
prop_assert!(!live_leaves.iter().any(|x| x.raw() == r.raw()));
live_leaves.push(r);
}
}
1 => {
if live_leaves.len() >= 2 {
let l1 = live_leaves[live_leaves.len() - 1];
let l2 = live_leaves[live_leaves.len() - 2];
if let Ok(r) = s.add_internal(
Cut::new(0, 0.0), unit_bbox::<2>(), l1, l2, None, 2,
) {
prop_assert!(!live_internals.iter().any(|x| x.raw() == r.raw()));
live_internals.push(r);
}
}
}
2 => {
if let Some(l) = live_leaves.pop() {
s.delete(l).unwrap();
}
}
_ => {
if let Some(i) = live_internals.pop() {
s.delete(i).unwrap();
}
}
}
prop_assert_eq!(s.live_leaf_count(), live_leaves.len());
prop_assert_eq!(s.live_internal_count(), live_internals.len());
prop_assert_eq!(s.live_count(), live_leaves.len() + live_internals.len());
}
}
}
}