fyrox_math/
octree.rs

1// Copyright (c) 2019-present Dmitry Stepanov and Fyrox Engine contributors.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a copy
4// of this software and associated documentation files (the "Software"), to deal
5// in the Software without restriction, including without limitation the rights
6// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7// copies of the Software, and to permit persons to whom the Software is
8// furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in all
11// copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19// SOFTWARE.
20
21use crate::{aabb::AxisAlignedBoundingBox, ray::Ray};
22use arrayvec::ArrayVec;
23use nalgebra::Vector3;
24
25#[derive(Clone, Debug)]
26pub enum OctreeNode {
27    Leaf {
28        indices: Vec<u32>,
29        bounds: AxisAlignedBoundingBox,
30    },
31    Branch {
32        bounds: AxisAlignedBoundingBox,
33        leaves: [usize; 8],
34    },
35}
36
37#[derive(Default, Clone, Debug)]
38pub struct Octree {
39    nodes: Vec<OctreeNode>,
40    root: usize,
41}
42
43impl Octree {
44    pub fn new(triangles: &[[Vector3<f32>; 3]], split_threshold: usize) -> Self {
45        // Calculate bounds.
46        let mut bounds = AxisAlignedBoundingBox::default();
47        for triangle in triangles {
48            for pt in triangle.iter() {
49                bounds.add_point(*pt);
50            }
51        }
52
53        // Inflate initial bounds by very low value to fix floating-point calculation
54        // issues when splitting and checking intersection later on.
55        let inflation = 2.0 * f32::EPSILON;
56        bounds.inflate(Vector3::new(inflation, inflation, inflation));
57
58        // Get initial list of indices.
59        let mut indices = Vec::new();
60        for i in 0..triangles.len() {
61            indices.push(i as u32);
62        }
63
64        let mut nodes = Vec::new();
65        let root = build_recursive(&mut nodes, triangles, bounds, indices, split_threshold);
66
67        Self { nodes, root }
68    }
69
70    pub fn sphere_query(&self, position: Vector3<f32>, radius: f32, buffer: &mut Vec<u32>) {
71        buffer.clear();
72        self.sphere_recursive_query(self.root, position, radius, buffer);
73    }
74
75    fn sphere_recursive_query(
76        &self,
77        node: usize,
78        position: Vector3<f32>,
79        radius: f32,
80        buffer: &mut Vec<u32>,
81    ) {
82        match &self.nodes[node] {
83            OctreeNode::Leaf { indices, bounds } => {
84                if bounds.is_intersects_sphere(position, radius) {
85                    buffer.extend_from_slice(indices)
86                }
87            }
88            OctreeNode::Branch { bounds, leaves } => {
89                if bounds.is_intersects_sphere(position, radius) {
90                    for leaf in leaves {
91                        self.sphere_recursive_query(*leaf, position, radius, buffer)
92                    }
93                }
94            }
95        }
96    }
97
98    pub fn aabb_query(&self, aabb: &AxisAlignedBoundingBox, buffer: &mut Vec<u32>) {
99        buffer.clear();
100        self.aabb_recursive_query(self.root, aabb, buffer);
101    }
102
103    fn aabb_recursive_query(
104        &self,
105        node: usize,
106        aabb: &AxisAlignedBoundingBox,
107        buffer: &mut Vec<u32>,
108    ) {
109        match &self.nodes[node] {
110            OctreeNode::Leaf { indices, bounds } => {
111                if bounds.is_intersects_aabb(aabb) {
112                    buffer.extend_from_slice(indices)
113                }
114            }
115            OctreeNode::Branch { bounds, leaves } => {
116                if bounds.is_intersects_aabb(aabb) {
117                    for leaf in leaves {
118                        self.aabb_recursive_query(*leaf, aabb, buffer)
119                    }
120                }
121            }
122        }
123    }
124
125    pub fn ray_query(&self, ray: &Ray, buffer: &mut Vec<u32>) {
126        buffer.clear();
127        self.ray_recursive_query(self.root, ray, buffer);
128    }
129
130    fn ray_recursive_query(&self, node: usize, ray: &Ray, buffer: &mut Vec<u32>) {
131        match &self.nodes[node] {
132            OctreeNode::Leaf { indices, bounds } => {
133                if ray.box_intersection(&bounds.min, &bounds.max).is_some() {
134                    buffer.extend_from_slice(indices)
135                }
136            }
137            OctreeNode::Branch { bounds, leaves } => {
138                if ray.box_intersection(&bounds.min, &bounds.max).is_some() {
139                    for leaf in leaves {
140                        self.ray_recursive_query(*leaf, ray, buffer)
141                    }
142                }
143            }
144        }
145    }
146
147    pub fn node(&self, handle: usize) -> &OctreeNode {
148        &self.nodes[handle]
149    }
150
151    pub fn nodes(&self) -> &Vec<OctreeNode> {
152        &self.nodes
153    }
154
155    pub fn ray_query_static<const CAP: usize>(&self, ray: &Ray, buffer: &mut ArrayVec<usize, CAP>) {
156        buffer.clear();
157        self.ray_recursive_query_static(self.root, ray, buffer);
158    }
159
160    fn ray_recursive_query_static<const CAP: usize>(
161        &self,
162        node: usize,
163        ray: &Ray,
164        buffer: &mut ArrayVec<usize, CAP>,
165    ) {
166        match &self.nodes[node] {
167            OctreeNode::Leaf { bounds, .. } => {
168                if ray.box_intersection(&bounds.min, &bounds.max).is_some() {
169                    buffer.push(node);
170                }
171            }
172            OctreeNode::Branch { bounds, leaves } => {
173                if ray.box_intersection(&bounds.min, &bounds.max).is_some() {
174                    for leaf in leaves {
175                        self.ray_recursive_query_static(*leaf, ray, buffer)
176                    }
177                }
178            }
179        }
180    }
181
182    pub fn point_query<C>(&self, point: Vector3<f32>, mut callback: C)
183    where
184        C: FnMut(&[u32]),
185    {
186        self.point_recursive_query(self.root, point, &mut callback);
187    }
188
189    fn point_recursive_query<C>(&self, node: usize, point: Vector3<f32>, callback: &mut C)
190    where
191        C: FnMut(&[u32]),
192    {
193        match &self.nodes[node] {
194            OctreeNode::Leaf { indices, bounds } => {
195                if bounds.is_contains_point(point) {
196                    (callback)(indices)
197                }
198            }
199            OctreeNode::Branch { bounds, leaves } => {
200                if bounds.is_contains_point(point) {
201                    for leaf in leaves {
202                        self.point_recursive_query(*leaf, point, callback)
203                    }
204                }
205            }
206        }
207    }
208}
209
210fn build_recursive(
211    nodes: &mut Vec<OctreeNode>,
212    triangles: &[[Vector3<f32>; 3]],
213    bounds: AxisAlignedBoundingBox,
214    indices: Vec<u32>,
215    split_threshold: usize,
216) -> usize {
217    if indices.len() <= split_threshold {
218        let index = nodes.len();
219        nodes.push(OctreeNode::Leaf { bounds, indices });
220        index
221    } else {
222        let mut leaves = [0; 8];
223        let leaf_bounds = bounds.split();
224
225        for i in 0..8 {
226            let mut leaf_indices = Vec::new();
227
228            for index in indices.iter() {
229                let index = *index;
230
231                let triangle_bounds =
232                    AxisAlignedBoundingBox::from_points(&triangles[index as usize]);
233
234                if triangle_bounds.is_intersects_aabb(&bounds) {
235                    leaf_indices.push(index);
236                }
237            }
238
239            leaves[i] = build_recursive(
240                nodes,
241                triangles,
242                leaf_bounds[i],
243                leaf_indices,
244                split_threshold,
245            );
246        }
247
248        let index = nodes.len();
249        nodes.push(OctreeNode::Branch { leaves, bounds });
250        index
251    }
252}
253
254#[cfg(test)]
255mod test {
256    use super::*;
257
258    fn get_six_triangles() -> [[Vector3<f32>; 3]; 6] {
259        [
260            [
261                Vector3::new(0.0, 0.0, 0.0),
262                Vector3::new(1.0, 0.0, 0.0),
263                Vector3::new(0.0, 1.0, 0.0),
264            ],
265            [
266                Vector3::new(1.0, 1.0, 0.0),
267                Vector3::new(1.0, 0.0, 0.0),
268                Vector3::new(0.0, 1.0, 0.0),
269            ],
270            [
271                Vector3::new(0.0, 1.0, 0.0),
272                Vector3::new(1.0, 1.0, 0.0),
273                Vector3::new(0.0, 2.0, 0.0),
274            ],
275            [
276                Vector3::new(1.0, 2.0, 0.0),
277                Vector3::new(1.0, 1.0, 0.0),
278                Vector3::new(0.0, 2.0, 0.0),
279            ],
280            [
281                Vector3::new(0.0, 2.0, 0.0),
282                Vector3::new(1.0, 2.0, 0.0),
283                Vector3::new(0.0, 3.0, 0.0),
284            ],
285            [
286                Vector3::new(1.0, 3.0, 0.0),
287                Vector3::new(1.0, 2.0, 0.0),
288                Vector3::new(0.0, 3.0, 0.0),
289            ],
290        ]
291    }
292
293    #[test]
294    fn octree_new() {
295        let tree = Octree::new(&get_six_triangles(), 5);
296
297        assert_eq!(tree.root, 72);
298        assert_eq!(tree.nodes().len(), 73);
299    }
300
301    #[test]
302    fn default_for_octree() {
303        let tree = Octree::default();
304        assert_eq!(tree.root, 0);
305        assert_eq!(tree.nodes.len(), 0);
306    }
307
308    #[test]
309    fn octree_point_query() {
310        let tree = Octree::new(&get_six_triangles(), 5);
311        let mut buffer = Vec::new();
312        tree.point_query(Vector3::new(0.0, 0.0, 0.0), |triangles| {
313            buffer.extend_from_slice(triangles)
314        });
315
316        assert_eq!(buffer, [0, 1, 2, 3, 0, 1, 2, 3]);
317    }
318
319    #[test]
320    fn octree_sphere_query() {
321        let tree = Octree::new(&get_six_triangles(), 5);
322        let mut buffer = Vec::new();
323        tree.sphere_query(Vector3::new(0.0, 0.0, 0.0), 1.0, &mut buffer);
324
325        assert_eq!(
326            buffer,
327            [
328                0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3,
329                0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3,
330                0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3,
331                0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3
332            ]
333        );
334    }
335
336    #[test]
337    fn octree_aabb_query() {
338        let tree = Octree::new(&get_six_triangles(), 5);
339        let mut buffer = Vec::new();
340        tree.aabb_query(
341            &AxisAlignedBoundingBox {
342                min: Vector3::new(0.0, 0.0, 0.0),
343                max: Vector3::new(0.5, 0.5, 0.5),
344            },
345            &mut buffer,
346        );
347
348        assert_eq!(
349            buffer,
350            [
351                0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3,
352                0, 1, 2, 3, 0, 1, 2, 3
353            ]
354        );
355    }
356
357    #[test]
358    fn octree_ray_query() {
359        let tree = Octree::new(&get_six_triangles(), 5);
360        let mut buffer = Vec::new();
361        tree.ray_query(
362            &Ray::new(Vector3::new(0.0, 0.0, 0.0), Vector3::new(1.0, 1.0, 0.0)),
363            &mut buffer,
364        );
365
366        assert_eq!(
367            buffer,
368            [
369                0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3,
370                0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3
371            ]
372        );
373    }
374
375    #[test]
376    fn octree_ray_query_static() {
377        const CAP: usize = 10;
378        let tree = Octree::new(&get_six_triangles(), 5);
379        let mut buffer = ArrayVec::<usize, CAP>::new();
380        tree.ray_query_static::<CAP>(
381            &Ray::new(Vector3::new(0.0, 0.0, 0.0), Vector3::new(1.0, 1.0, 0.0)),
382            &mut buffer,
383        );
384
385        assert_eq!(buffer.as_slice(), [2, 3, 11, 15, 16, 18, 19, 27, 31, 32,]);
386    }
387}
fyrox_math/octree.rs

fyrox_math/
octree.rs
