diff --git a/kdtree-cli/src/main.rs b/kdtree-cli/src/main.rs index 78487d6..633a7c2 100644 --- a/kdtree-cli/src/main.rs +++ b/kdtree-cli/src/main.rs @@ -199,6 +199,7 @@ fn main() { let kdtree = build_kdtree(&geometries, &cost); let duration = Instant::now().duration_since(start_time); let duration = Duration::new(duration.as_secs() as i64, duration.as_nanos() as i32); + eprintln!("Done in {duration:.3}..."); let cost = tree_cost( &geometries, @@ -208,9 +209,7 @@ fn main() { args.empty_factor, ); let stats = statistics(&geometries, &kdtree); - eprintln!("Done..."); eprintln!("Tree statistics:"); - eprintln!(" Build time: {duration:.3}"); eprintln!(" SAH cost: {cost:.3}"); eprintln!(" Geometries: {}", stats.geometries); eprintln!(" Node count: {}", stats.node_count); diff --git a/kdtree/src/lib.rs b/kdtree/src/lib.rs index 6e018f3..9e0aef6 100644 --- a/kdtree/src/lib.rs +++ b/kdtree/src/lib.rs @@ -65,10 +65,8 @@ impl KdNode { } #[inline] - pub fn iter_nodes(&self) -> KdTreeNodeIter { - KdTreeNodeIter { - stack: vec![(1, self)], - } + pub fn iter_nodes(&self) -> KdNodeIter { + KdNodeIter::new(self) } #[inline] @@ -152,13 +150,25 @@ impl Display for KdNode { } } -pub struct KdTreeNodeIter<'a> { - pub(crate) stack: Vec<(usize, &'a KdNode)>, +pub struct KdNodeIter<'a> { + pub(crate) stack: ArrayVec<(usize, &'a KdNode), MAX_DEPTH>, } -impl<'a> Iterator for KdTreeNodeIter<'a> { +impl<'a> KdNodeIter<'a> { + #[inline] + fn new(node: &'a KdNode) -> Self { + let mut stack = ArrayVec::<(usize, &'a KdNode), MAX_DEPTH>::new(); + unsafe { + stack.push_unchecked((1, node)); + } + Self { stack } + } +} + +impl<'a> Iterator for KdNodeIter<'a> { type Item = (usize, &'a KdNode); + #[inline] fn next(&mut self) -> Option { if let Some((depth, node)) = self.stack.pop() { match node {