#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use super::bounding_box::BoundingBox;
pub(super) const NULL: usize = usize::MAX;
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub(super) enum Node {
Leaf {
point_idx: usize,
mass: usize,
},
Internal {
left: usize,
right: usize,
cut_dim: usize,
cut_val: f32,
mass: usize,
bbox: BoundingBox,
},
}
impl Node {
pub(super) fn mass(&self) -> usize {
match self {
Node::Leaf { mass, .. } => *mass,
Node::Internal { mass, .. } => *mass,
}
}
}
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub(super) struct NodeArena {
nodes: Vec<Option<Node>>,
free: Vec<usize>,
}
impl NodeArena {
pub(super) fn new(initial_capacity: usize) -> Self {
NodeArena {
nodes: Vec::with_capacity(initial_capacity),
free: Vec::new(),
}
}
pub(super) fn alloc(&mut self, node: Node) -> usize {
if let Some(id) = self.free.pop() {
self.nodes[id] = Some(node);
id
} else {
let id = self.nodes.len();
self.nodes.push(Some(node));
id
}
}
pub(super) fn get(&self, id: usize) -> &Node {
self.nodes[id]
.as_ref()
.expect("accessed freed or uninitialized node")
}
pub(super) fn get_mut(&mut self, id: usize) -> &mut Node {
self.nodes[id]
.as_mut()
.expect("accessed freed or uninitialized node")
}
pub(super) fn free(&mut self, id: usize) {
debug_assert!(id < self.nodes.len());
self.nodes[id] = None;
self.free.push(id);
}
#[cfg(all(test, feature = "std"))]
fn slot_count(&self) -> usize {
self.nodes.len()
}
}
#[cfg(all(test, feature = "std"))]
mod tests {
use super::*;
use proptest::prelude::*;
fn leaf_node() -> Node {
Node::Leaf {
point_idx: 0,
mass: 1,
}
}
proptest! {
#[test]
fn slot_count_increases_with_alloc(n in 1usize..=32) {
let mut arena = NodeArena::new(4);
let before = arena.slot_count();
for _ in 0..n {
arena.alloc(leaf_node());
}
prop_assert!(arena.slot_count() > before);
}
#[test]
fn alloc_then_free_restores_slot(n in 1usize..=16) {
let mut arena = NodeArena::new(n);
let ids: Vec<usize> = (0..n).map(|_| arena.alloc(leaf_node())).collect();
let count_after_alloc = arena.slot_count();
for id in ids {
arena.free(id);
}
for _ in 0..n {
arena.alloc(leaf_node());
}
prop_assert!(arena.slot_count() <= count_after_alloc + 1);
}
}
}