Remove KdTree struct to simplify geometry ownership

kdtree-cli/src/main.rs

@@ -4,9 +4,9 @@ use geometry::{
 };
 use kdtree::{
     build::build_kdtree,
-    format::{write_tree_dot, write_tree_json, write_tree_pretty, write_tree_rust},
+    format::{write_node_pretty, write_tree_dot, write_tree_json, write_tree_rust},
     sah::SahCost,
-    KdNode, KdTree, MAX_DEPTH,
+    KdNode, MAX_DEPTH,
 };
 use std::{
     fs::File,
@@ -89,15 +89,21 @@ fn node_cost(
     }
 }
 
-fn tree_cost(kdtree: &KdTree, cost_traverse: f32, cost_intersect: f32, empty_factor: f32) -> f32 {
-    let bounding_box = geometries_bounding_box(&kdtree.geometries);
+fn tree_cost(
+    geometries: &[Geometry],
+    node: &KdNode,
+    cost_traverse: f32,
+    cost_intersect: f32,
+    empty_factor: f32,
+) -> f32 {
+    let bounding_box = geometries_bounding_box(geometries);
     node_cost(
         cost_traverse,
         cost_intersect,
         empty_factor,
         bounding_box.surface_area(),
         bounding_box,
-        kdtree.root.as_ref(),
+        node,
     )
 }
 
@@ -138,8 +144,8 @@ struct KdTreeStatistics {
     leaf_geometries: Statistics,
 }
 
-fn statistics(tree: &KdTree) -> KdTreeStatistics {
-    let geometries = tree.geometries.len();
+fn statistics(geometries: &[Geometry], tree: &KdNode) -> KdTreeStatistics {
+    let geometries = geometries.len();
     let node_count = tree.iter_nodes().map(|_| 1).sum();
     let leaf_count = tree.iter_leafs().map(|_| 1).sum();
     let leaf_depth = Statistics::compute(tree.iter_leafs().map(|(depth, _)| depth).collect());
@@ -195,17 +201,18 @@ fn main() {
         intersect_cost: args.intersect_cost,
         empty_factor: args.empty_factor,
     };
-    let kdtree = build_kdtree(geometries, args.max_depth, &cost);
+    let kdtree = build_kdtree(&geometries, args.max_depth, &cost);
     let duration = Instant::now().duration_since(start_time);
     let duration = Duration::new(duration.as_secs() as i64, duration.as_nanos() as i32);
 
     let cost = tree_cost(
+        &geometries,
         &kdtree,
         args.traverse_cost,
         args.intersect_cost,
         args.empty_factor,
     );
-    let stats = statistics(&kdtree);
+    let stats = statistics(&geometries, &kdtree);
     eprintln!("Done...");
     eprintln!("Tree statistics:");
     eprintln!("  Build time: {:.3}", duration);
@@ -226,12 +233,12 @@ fn main() {
     eprintln!("    Median: {}", stats.leaf_geometries.median);
 
     if args.json {
-        write_tree_json(&mut io::stdout().lock(), &kdtree).unwrap();
+        write_tree_json(&mut io::stdout().lock(), &geometries, &kdtree).unwrap();
     } else if args.rust {
-        write_tree_rust(&mut io::stdout().lock(), &kdtree).unwrap();
+        write_tree_rust(&mut io::stdout().lock(), &geometries, &kdtree).unwrap();
     } else if args.dot {
         write_tree_dot(&mut io::stdout().lock(), &kdtree).unwrap();
     } else {
-        write_tree_pretty(&mut io::stdout().lock(), &kdtree).unwrap();
+        write_node_pretty(&mut io::stdout().lock(), &kdtree).unwrap();
     }
 }

kdtree-tester/src/bin/kdtree-reducer-cli.rs

@@ -11,17 +11,22 @@ use rand::{rngs::SmallRng, seq::SliceRandom, SeedableRng};
 use geometry::{geometry::Geometry, intersection::RayIntersection, ray::Ray, triangle::Triangle};
 use kdtree::{
     build::build_kdtree, format::write_tree_json, intersection::KdIntersection, sah::SahCost,
-    KdTree, MAX_DEPTH,
+    KdNode, MAX_DEPTH,
 };
 use wavefront::obj;
 
-fn build_test_tree(geometries: Vec<Geometry>) -> kdtree::KdTree {
+fn build_test_tree(geometries: &[Geometry]) -> KdNode {
     build_kdtree(geometries, MAX_DEPTH as u32, &SahCost::default())
 }
 
-fn verify_removal(ray: &Ray, actual: &(Geometry, RayIntersection), tree: &KdTree) -> bool {
-    let intersection = tree.intersect(ray, 0.0..=f32::MAX).unwrap();
-    let same_geometry = tree.geometries[intersection.index as usize] == actual.0;
+fn verify_removal(
+    geometries: &[Geometry],
+    ray: &Ray,
+    actual: &(Geometry, RayIntersection),
+    tree: &KdNode,
+) -> bool {
+    let intersection = tree.intersect(geometries, ray, 0.0..=f32::MAX).unwrap();
+    let same_geometry = geometries[intersection.index as usize] == actual.0;
     let same_intersection = intersection.intersection == actual.1;
     same_geometry && same_intersection
 }
@@ -36,11 +41,15 @@ fn try_removing(
     let mut reduced = Vec::with_capacity(geometries.len() - try_count);
     reduced.extend_from_slice(&geometries[0..try_index]);
     reduced.extend_from_slice(&geometries[try_index + try_count..]);
-    let tree = build_test_tree(reduced);
-    verify_removal(ray, actual, &tree).then_some(tree.geometries)
+    let tree = build_test_tree(&reduced);
+    verify_removal(&reduced, ray, actual, &tree).then_some(reduced)
 }
 
-fn reduce_tree(seed: u64, intersection: &CheckedIntersection, geometries: Vec<Geometry>) -> KdTree {
+fn reduce_tree(
+    seed: u64,
+    intersection: &CheckedIntersection,
+    geometries: Vec<Geometry>,
+) -> (Vec<Geometry>, KdNode) {
     let actual_intersection = intersection.kdtree.as_ref().unwrap();
     let actual_geometry = geometries[actual_intersection.index as usize].clone();
     let actual = (actual_geometry, actual_intersection.intersection);
@@ -87,7 +96,8 @@ fn reduce_tree(seed: u64, intersection: &CheckedIntersection, geometries: Vec<Ge
             );
         }
     }
-    build_test_tree(geometries)
+    let tree = build_test_tree(&geometries);
+    (geometries, tree)
 }
 
 #[derive(Parser, Debug)]
@@ -179,11 +189,12 @@ fn main() {
     }
 
     eprintln!("Reducing tree...");
-    let tree = reduce_tree(args.seed, &intersection, geometries);
+    let (geometries, tree) = reduce_tree(args.seed, &intersection, geometries);
 
     eprintln!("Writing reduced tree to {:?}...", args.output);
     write_tree_json(
         &mut BufWriter::new(File::create(args.output).unwrap()),
+        &geometries,
         &tree,
     )
     .unwrap();

kdtree-tester/src/bin/kdtree-tester-cli.rs

@@ -82,7 +82,7 @@ fn main() {
 
     println!("Building kdtree...");
     let kdtree = build_kdtree(
-        scene.geometries.clone(),
+        &scene.geometries,
         args.max_depth,
         &SahCost {
             traverse_cost: args.traverse_cost,

kdtree-tester/src/ray_bouncer.rs

@@ -3,7 +3,7 @@ use geometry::ray::Ray;
 use glam::{UVec2, Vec2};
 use kdtree::{
     intersection::{intersect_closest_geometry, KdIntersection},
-    KdTree,
+    KdNode,
 };
 use rand::{rngs::SmallRng, SeedableRng};
 use scene::{camera::Pinhole, Scene};
@@ -15,7 +15,7 @@ use tracing::{
 
 pub struct RayBouncer {
     pub scene: Scene,
-    pub kdtree: KdTree,
+    pub kdtree: KdNode,
     pub camera: Pinhole,
     pub bounces: u32,
     pub size: UVec2,
@@ -27,16 +27,18 @@ impl RayBouncer {
         ray: &Ray,
         t_range: RangeInclusive<f32>,
     ) -> Option<KdIntersection> {
-        let indices = 0u32..self.kdtree.geometries.len() as u32;
-        intersect_closest_geometry(&self.kdtree.geometries, indices, ray, t_range)
+        let indices = 0u32..self.scene.geometries.len() as u32;
+        intersect_closest_geometry(&self.scene.geometries, indices, ray, t_range)
     }
 
     fn checked_ray_intersect(
         &self,
         ray: &Ray,
         t_range: RangeInclusive<f32>,
     ) -> CheckedIntersection {
-        let kdtree = self.kdtree.intersect(ray, t_range.clone());
+        let kdtree = self
+            .kdtree
+            .intersect(&self.scene.geometries, ray, t_range.clone());
         let reference = self.reference_ray_intersect(ray, t_range);
         CheckedIntersection {
             ray: *ray,

kdtree/src/build.rs

@@ -4,9 +4,10 @@ use glam::Vec3;
 use crate::{
     cell::KdCell,
     sah::{find_best_split, should_terminate, SahCost},
+    MAX_DEPTH,
 };
 
-use super::{KdNode, KdTree};
+use super::KdNode;
 
 fn starting_box(geometries: &[Geometry]) -> KdCell {
     KdCell::new(
@@ -40,16 +41,14 @@ fn build_helper(
     }
 }
 
-pub fn build_kdtree(geometries: Vec<Geometry>, max_depth: u32, cost: &SahCost) -> KdTree {
-    if max_depth as usize > super::MAX_DEPTH {
+pub fn build_kdtree(geometries: &[Geometry], max_depth: u32, cost: &SahCost) -> KdNode {
+    if max_depth as usize > MAX_DEPTH {
         panic!(
             "Max depth ({}) must be smaller than hard coded value ({}).",
-            max_depth,
-            super::MAX_DEPTH
+            max_depth, MAX_DEPTH
         );
     }
-    let root = build_helper(&geometries, cost, max_depth, 1, starting_box(&geometries));
-    KdTree { root, geometries }
+    *build_helper(geometries, cost, max_depth, 1, starting_box(geometries))
 }
 
 #[cfg(test)]
@@ -73,7 +72,7 @@ mod tests {
             intersect_cost: 1.0,
             empty_factor: 0.8,
         };
-        let tree = build_kdtree(vec![triangle.into()], 7, &cost);
+        let actual = build_kdtree(&vec![triangle.into()], 7, &cost);
 
         let expected = KdNode::new_node(
             Aap::new_x(0.0),
@@ -101,9 +100,9 @@ mod tests {
             ),
         );
         assert_eq!(
-            tree.root, expected,
+            actual, *expected,
             "\n   actual: {}\n expected: {}",
-            tree.root, expected
+            actual, *expected
         );
     }
 
@@ -119,7 +118,7 @@ mod tests {
             intersect_cost: 10.0,
             empty_factor: 0.8,
         };
-        let tree = build_kdtree(vec![triangle.into()], 6, &cost);
+        let actual = build_kdtree(&vec![triangle.into()], 6, &cost);
 
         let expected = KdNode::new_node(
             Aap::new_x(0.0),
@@ -135,9 +134,9 @@ mod tests {
             ),
         );
         assert_eq!(
-            tree.root, expected,
+            actual, *expected,
             "\n   actual: {}\n expected: {}",
-            tree.root, expected
+            actual, *expected
         );
     }
 }

kdtree/src/format.rs

@@ -1,4 +1,4 @@
-use crate::{KdNode, KdTree};
+use crate::KdNode;
 use geometry::{axis::Axis, geometry::Geometry};
 use std::io::{self};
 
@@ -46,13 +46,6 @@ where
     Ok(())
 }
 
-pub fn write_tree_pretty<W>(write: &mut W, tree: &KdTree) -> Result<(), io::Error>
-where
-    W: io::Write,
-{
-    write_node_pretty(write, &tree.root)
-}
-
 pub fn write_node_rust<W>(write: &mut W, node: &KdNode) -> Result<(), io::Error>
 where
     W: io::Write,
@@ -81,15 +74,19 @@ where
     Ok(())
 }
 
-pub fn write_tree_rust<W>(write: &mut W, tree: &KdTree) -> Result<(), io::Error>
+pub fn write_tree_rust<W>(
+    write: &mut W,
+    geometries: &[Geometry],
+    tree: &KdNode,
+) -> Result<(), io::Error>
 where
     W: io::Write,
 {
     write!(write, "let geometries = ")?;
-    write_triangle_bracketed(write, &tree.geometries)?;
+    write_triangle_bracketed(write, geometries)?;
     writeln!(write, ";")?;
     write!(write, "let root = ")?;
-    write_node_rust(write, &tree.root)?;
+    write_node_rust(write, tree)?;
     writeln!(write, ";")?;
     writeln!(write, "let tree = KdTree {{ root, geometries }};")?;
     Ok(())
@@ -117,14 +114,18 @@ where
     Ok(())
 }
 
-pub fn write_tree_json<W>(write: &mut W, tree: &KdTree) -> Result<(), io::Error>
+pub fn write_tree_json<W>(
+    write: &mut W,
+    geometries: &[Geometry],
+    tree: &KdNode,
+) -> Result<(), io::Error>
 where
     W: io::Write,
 {
     write!(write, "{{\"triangles\": ")?;
-    write_triangle_bracketed(write, &tree.geometries)?;
+    write_triangle_bracketed(write, geometries)?;
     write!(write, ", \"root\": ")?;
-    write_node_json(write, &tree.root)?;
+    write_node_json(write, tree)?;
     writeln!(write, "}}")?;
     Ok(())
 }
@@ -156,14 +157,14 @@ where
     Ok(())
 }
 
-pub fn write_tree_dot<W>(write: &mut W, tree: &KdTree) -> Result<(), io::Error>
+pub fn write_tree_dot<W>(write: &mut W, tree: &KdNode) -> Result<(), io::Error>
 where
     W: io::Write,
 {
     writeln!(write, "digraph {{")?;
     writeln!(write, "  rankdir=\"LR\";")?;
     writeln!(write, "  node [shape=\"box\"];")?;
-    write_node_dot(write, "t".to_string(), &tree.root)?;
+    write_node_dot(write, "t".to_string(), tree)?;
     writeln!(write, "}}")?;
     Ok(())
 }