1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
use crate::bounding_volume::{Aabb, SimdAabb};
use crate::math::{Real, Vector};

use na::SimdValue;

#[cfg(feature = "rkyv")]
use rkyv::{bytecheck, CheckBytes};

/// A data to which an index is associated.
pub trait IndexedData: Copy {
    /// Creates a new default instance of `Self`.
    fn default() -> Self;
    /// Gets the index associated to `self`.
    fn index(&self) -> usize;
}

impl IndexedData for usize {
    fn default() -> Self {
        // NOTE: we use u32::MAX for compatibility
        // between 32 and 64 bit platforms.
        u32::MAX as usize
    }
    fn index(&self) -> usize {
        *self
    }
}

impl IndexedData for u32 {
    fn default() -> Self {
        u32::MAX
    }
    fn index(&self) -> usize {
        *self as usize
    }
}

impl IndexedData for u64 {
    fn default() -> Self {
        u64::MAX
    }
    fn index(&self) -> usize {
        *self as usize
    }
}

/// The index of an internal SIMD node of a Qbvh.
pub type SimdNodeIndex = u32;

/// The index of a node part of a Qbvh.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
#[cfg_attr(
    feature = "rkyv",
    derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, CheckBytes),
    archive(as = "Self")
)]
/// The index of one specific node of a Qbvh.
pub struct NodeIndex {
    /// The index of the SIMD node containing the addressed node.
    pub index: SimdNodeIndex, // Index of the addressed node in the `nodes` array.
    /// The SIMD lane the addressed node is associated to.
    pub lane: u8, // SIMD lane of the addressed node.
}

impl NodeIndex {
    pub(super) fn new(index: u32, lane: u8) -> Self {
        Self { index, lane }
    }

    pub(super) fn invalid() -> Self {
        Self {
            index: u32::MAX,
            lane: 0,
        }
    }

    pub(super) fn is_invalid(&self) -> bool {
        self.index == u32::MAX
    }
}

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(
    feature = "rkyv",
    derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize),
    archive(as = "Self")
)]
#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
/// The status of a QBVH node.
pub struct QbvhNodeFlags(u8);

bitflags::bitflags! {
    impl QbvhNodeFlags: u8 {
        /// If this bit is set, the node is a leaf.
        const LEAF = 0b0001;
        /// If this bit is set, this node was recently changed.
        const CHANGED = 0b0010;
        /// Does this node need an update?
        const DIRTY = 0b0100;
    }
}

/// A SIMD node of an SIMD Qbvh.
///
/// This groups four nodes of the Qbvh.
#[derive(Copy, Clone, Debug)]
#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
#[cfg_attr(
    feature = "rkyv",
    derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize),
    archive(check_bytes)
)]
pub struct QbvhNode {
    /// The Aabbs of the qbvh nodes represented by this node.
    pub simd_aabb: SimdAabb,
    /// Index of the nodes of the 4 nodes represented by `self`.
    /// If this is a leaf, it contains the proxy ids instead.
    pub children: [u32; 4],
    /// The index of the node parent to the 4 nodes represented by `self`.
    pub parent: NodeIndex,
    /// Status flags for this node.
    pub flags: QbvhNodeFlags,
}

impl QbvhNode {
    #[inline]
    /// Is this node a leaf?
    pub fn is_leaf(&self) -> bool {
        self.flags.contains(QbvhNodeFlags::LEAF)
    }

    #[inline]
    /// Does the AABB of this node needs to be updated?
    pub fn is_dirty(&self) -> bool {
        self.flags.contains(QbvhNodeFlags::DIRTY)
    }

    #[inline]
    /// Sets if the AABB of this node needs to be updated.
    pub fn set_dirty(&mut self, dirty: bool) {
        self.flags.set(QbvhNodeFlags::DIRTY, dirty);
    }

    #[inline]
    /// Was the AABB of this node changed since the last rebalancing?
    pub fn is_changed(&self) -> bool {
        self.flags.contains(QbvhNodeFlags::CHANGED)
    }

    #[inline]
    /// Sets if the AABB of this node changed since the last rebalancing.
    pub fn set_changed(&mut self, changed: bool) {
        self.flags.set(QbvhNodeFlags::CHANGED, changed);
    }

    #[inline]
    /// An empty internal node.
    pub fn empty() -> Self {
        Self {
            simd_aabb: SimdAabb::new_invalid(),
            children: [u32::MAX; 4],
            parent: NodeIndex::invalid(),
            flags: QbvhNodeFlags::default(),
        }
    }

    #[inline]
    /// An empty leaf.
    pub fn empty_leaf_with_parent(parent: NodeIndex) -> Self {
        Self {
            simd_aabb: SimdAabb::new_invalid(),
            children: [u32::MAX; 4],
            parent,
            flags: QbvhNodeFlags::LEAF,
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
#[cfg_attr(
    feature = "rkyv",
    derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize),
    archive(check_bytes)
)]
/// Combination of a leaf data and its associated node’s index.
pub struct QbvhProxy<LeafData> {
    /// Index of the leaf node the leaf data is associated to.
    pub node: NodeIndex,
    /// The data contained in this node.
    pub data: LeafData, // The collider data. TODO: only set the collider generation here?
}

impl<LeafData> QbvhProxy<LeafData> {
    pub(super) fn invalid() -> Self
    where
        LeafData: IndexedData,
    {
        Self {
            node: NodeIndex::invalid(),
            data: LeafData::default(),
        }
    }

    pub(super) fn detached(data: LeafData) -> Self {
        Self {
            node: NodeIndex::invalid(),
            data,
        }
    }

    pub(super) fn is_detached(&self) -> bool {
        self.node.is_invalid()
    }
}

/// A quaternary bounding-volume-hierarchy.
///
/// This is a bounding-volume-hierarchy where each node has either four children or none.
#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
#[cfg_attr(
    feature = "rkyv",
    derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize),
    archive(check_bytes)
)]
#[repr(C)]
#[derive(Debug, Clone)]
pub struct Qbvh<LeafData> {
    pub(super) root_aabb: Aabb,
    pub(super) nodes: Vec<QbvhNode>,
    pub(super) dirty_nodes: Vec<u32>,
    pub(super) free_list: Vec<u32>,
    pub(super) proxies: Vec<QbvhProxy<LeafData>>,
}

impl<LeafData: IndexedData> Default for Qbvh<LeafData> {
    fn default() -> Self {
        Self::new()
    }
}

impl<LeafData: IndexedData> Qbvh<LeafData> {
    /// Initialize an empty Qbvh.
    pub fn new() -> Self {
        Qbvh {
            root_aabb: Aabb::new_invalid(),
            nodes: Vec::new(),
            dirty_nodes: Vec::new(),
            free_list: Vec::new(),
            proxies: Vec::new(),
        }
    }

    /// Iterates mutably through all the leaf data in this Qbvh.
    pub fn iter_data_mut(&mut self) -> impl Iterator<Item = (NodeIndex, &mut LeafData)> {
        self.proxies.iter_mut().map(|p| (p.node, &mut p.data))
    }

    /// Iterate through all the leaf data in this Qbvh.
    pub fn iter_data(&self) -> impl Iterator<Item = (NodeIndex, &LeafData)> {
        self.proxies.iter().map(|p| (p.node, &p.data))
    }

    /// The Aabb of the given node.
    pub fn node_aabb(&self, node_id: NodeIndex) -> Option<Aabb> {
        self.nodes
            .get(node_id.index as usize)
            .map(|n| n.simd_aabb.extract(node_id.lane as usize))
    }

    /// Returns the data associated to a given leaf.
    ///
    /// Returns `None` if the provided node ID does not identify a leaf.
    pub fn leaf_data(&self, node_id: NodeIndex) -> Option<LeafData> {
        let node = self.nodes.get(node_id.index as usize)?;

        if !node.is_leaf() {
            return None;
        }

        let proxy = self
            .proxies
            .get(node.children[node_id.lane as usize] as usize)?;
        Some(proxy.data)
    }

    /// The raw nodes of this BVH.
    ///
    /// If this Qbvh isn’t empty, the first element of the returned slice is the root of the
    /// tree. The other elements are not arranged in any particular order.
    /// The more high-level traversal methods should be used instead of this.
    pub fn raw_nodes(&self) -> &[QbvhNode] {
        &self.nodes
    }

    /// The raw proxies of this BVH.
    ///
    /// If this Qbvh isn’t empty, the first element of the returned slice is the root of the
    /// tree. The other elements are not arranged in any particular order.
    /// The more high-level traversal methods should be used instead of this.
    pub fn raw_proxies(&self) -> &[QbvhProxy<LeafData>] {
        &self.proxies
    }

    /// Computes a scaled version of this Qbvh.
    ///
    /// This will apply the scale to each Aabb on this BVH.
    pub fn scaled(mut self, scale: &Vector<Real>) -> Self {
        self.root_aabb = self.root_aabb.scaled(scale);
        for node in &mut self.nodes {
            node.simd_aabb = node.simd_aabb.scaled(&Vector::splat(*scale));
        }
        self
    }
}

impl<LeafData: IndexedData> Qbvh<LeafData> {
    /// The Aabb of the root of this tree.
    pub fn root_aabb(&self) -> &Aabb {
        &self.root_aabb
    }
}

#[cfg(test)]
mod test {
    use crate::bounding_volume::Aabb;
    use crate::math::{Point, Vector};
    use crate::partitioning::Qbvh;

    #[test]
    fn multiple_identical_aabb_stack_overflow() {
        // A stack overflow was caused during the construction of the
        // Qbvh with more than four Aabb with the same center.
        let aabb = Aabb::new(Point::origin(), Vector::repeat(1.0).into());

        for k in 0u32..20 {
            let mut tree = Qbvh::new();
            tree.clear_and_rebuild((0..k).map(|i| (i, aabb)), 0.0);
        }
    }
}