kdtree_ray/
kdtree.rs

1use rayon::ThreadPoolBuilder;
2
3use crate::aabb::*;
4use crate::candidate::*;
5use crate::config::BuilderConfig;
6use crate::kdnode::{build_tree, KDTreeNode};
7use crate::ray::Ray;
8use crate::Vector3;
9
10/// The KD-tree data structure.
11#[derive(Clone, Debug)]
12pub struct KDTree {
13    tree: Vec<KDTreeNode>,
14    space: AABB,
15    depth: usize,
16}
17
18impl KDTree {
19    /// This function is used to build a new KD-tree. You need to provide a
20    /// `Vec` of shapes that implement `Bounded` trait.
21    /// You also should give a configuration.
22    /// Panic if the `shapes` is empty.
23    pub fn build_config<S: Bounded>(shapes: &Vec<S>, config: &BuilderConfig) -> Self {
24        assert!(!shapes.is_empty());
25        let mut space = AABB::default();
26        let mut candidates = Candidates::with_capacity(shapes.len() * 6);
27        for (index, v) in shapes.iter().enumerate() {
28            // Create items from values
29            let bb = v.bound();
30            candidates.extend(Candidate::gen_candidates(index, &bb));
31
32            // Update space with the bounding box of the item
33            space.merge(&bb);
34        }
35
36        // Sort candidates only once at the begining
37        candidates.sort();
38
39        let nb_shapes = shapes.len();
40
41        // Build the tree
42        let pool = ThreadPoolBuilder::new().build().unwrap();
43        let (depth, tree) = pool.install(|| build_tree(config, &space, candidates, nb_shapes));
44
45        KDTree { space, tree, depth }
46    }
47
48    /// This function is used to build a new KD-tree. You need to provide a
49    /// `Vec` of shapes that implement `Bounded` trait.
50    /// Take a default configuration.
51    /// Panic if the `shapes` is empty.
52    pub fn build<S: Bounded>(shapes: &Vec<S>) -> Self {
53        Self::build_config(shapes, &BuilderConfig::default())
54    }
55
56    /// This function takes a ray and return a reduced list of shapes that
57    /// can be intersected by the ray.
58    pub fn intersect(&self, ray_origin: &Vector3, ray_direction: &Vector3) -> Vec<usize> {
59        let ray = Ray::new(ray_origin, ray_direction);
60        let mut result = vec![];
61        let mut stack = vec![0];
62        stack.reserve_exact(self.depth);
63        while !stack.is_empty() {
64            let node = &self.tree[stack.pop().unwrap()];
65            match node {
66                KDTreeNode::Leaf { shapes } => result.extend(shapes),
67                KDTreeNode::Node {
68                    l_child,
69                    l_space,
70                    r_child,
71                    r_space,
72                } => {
73                    if ray.intersect(r_space) {
74                        stack.push(*r_child)
75                    }
76                    if ray.intersect(l_space) {
77                        stack.push(*l_child)
78                    }
79                }
80            }
81        }
82        // Dedup duplicated shapes
83        result.sort();
84        result.dedup();
85        result
86    }
87}
88
89impl Bounded for KDTree {
90    fn bound(&self) -> AABB {
91        self.space.clone()
92    }
93}
kdtree_ray/kdtree.rs

kdtree_ray/
kdtree.rs
