avian3d/collider_tree/
tree.rs

1use bevy::{
2    ecs::{entity::Entity, resource::Resource},
3    reflect::prelude::*,
4};
5use obvhs::{
6    aabb::Aabb,
7    bvh2::{Bvh2, reinsertion::ReinsertionOptimizer},
8    ploc::{PlocBuilder, PlocSearchDistance, SortPrecision, rebuild::compute_rebuild_path_flags},
9};
10
11use crate::{
12    collider_tree::ProxyId,
13    data_structures::stable_vec::StableVec,
14    prelude::{ActiveCollisionHooks, CollisionLayers},
15};
16
17/// A [Bounding Volume Hierarchy (BVH)][BVH] for accelerating queries on a set of colliders.
18///
19/// See the [`collider_tree`](crate::collider_tree) module for more information.
20///
21/// [BVH]: https://en.wikipedia.org/wiki/Bounding_volume_hierarchy
22#[derive(Clone, Default)]
23pub struct ColliderTree {
24    /// The underlying BVH structure.
25    pub bvh: Bvh2,
26    /// The proxies stored in the tree.
27    pub proxies: StableVec<ColliderTreeProxy>,
28    /// A list of moved proxies since the last update.
29    ///
30    /// This is used during tree optimization to determine
31    /// which proxies need to be reinserted or rebuilt.
32    pub moved_proxies: MovedProxyList,
33    /// A workspace for reusing allocations across tree operations.
34    pub workspace: ColliderTreeWorkspace,
35}
36
37/// A set of [`ProxyId`]s that have moved since the last update, in insertion order.
38///
39/// Supports constant-time insertion, removal, and membership tests by keeping
40/// a sparse index from [`ProxyId`] to the position of the ID in the dense list.
41#[derive(Clone, Default, Debug)]
42pub struct MovedProxyList {
43    /// The moved proxy IDs, in insertion order.
44    proxies: Vec<ProxyId>,
45    /// Maps a [`ProxyId`] to its index in `proxies`, or [`PROXY_NOT_PRESENT`] if it is not in the list.
46    indices: Vec<u32>,
47}
48
49/// The sentinel stored in [`MovedProxyList::indices`] for proxies that are not in the list.
50const PROXY_NOT_PRESENT: u32 = u32::MAX;
51
52impl MovedProxyList {
53    /// Adds a proxy to the list if it is not already present.
54    ///
55    /// Returns `true` if the proxy was added.
56    #[inline]
57    pub fn push(&mut self, proxy_id: ProxyId) -> bool {
58        if self.indices.len() <= proxy_id.index() {
59            self.indices.resize(proxy_id.index() + 1, PROXY_NOT_PRESENT);
60        } else if self.indices[proxy_id.index()] != PROXY_NOT_PRESENT {
61            return false;
62        }
63
64        self.indices[proxy_id.index()] = self.proxies.len() as u32;
65        self.proxies.push(proxy_id);
66
67        true
68    }
69
70    /// Removes a proxy from the list. This may change the order of the remaining proxies.
71    ///
72    /// If the proxy is not present, nothing happens.
73    #[inline]
74    pub fn remove(&mut self, proxy_id: ProxyId) {
75        let Some(index) = self.indices.get_mut(proxy_id.index()) else {
76            return;
77        };
78        if *index == PROXY_NOT_PRESENT {
79            return;
80        }
81
82        let index = core::mem::replace(index, PROXY_NOT_PRESENT) as usize;
83        self.proxies.swap_remove(index);
84
85        // The proxy that was swapped into the removed slot needs its index updated.
86        if let Some(&swapped) = self.proxies.get(index) {
87            self.indices[swapped.index()] = index as u32;
88        }
89    }
90
91    /// Removes all proxies from the list.
92    #[inline]
93    pub fn clear(&mut self) {
94        for proxy_id in self.proxies.drain(..) {
95            self.indices[proxy_id.index()] = PROXY_NOT_PRESENT;
96        }
97    }
98}
99
100impl core::ops::Deref for MovedProxyList {
101    type Target = [ProxyId];
102
103    #[inline]
104    fn deref(&self) -> &[ProxyId] {
105        &self.proxies
106    }
107}
108
109/// A proxy representing a collider in the [`ColliderTree`].
110#[derive(Clone, Debug)]
111pub struct ColliderTreeProxy {
112    /// The collider entity this proxy represents.
113    pub collider: Entity,
114    /// The body this collider is attached to.
115    pub body: Option<Entity>,
116    /// The collision layers of the collider.
117    pub layers: CollisionLayers,
118    /// Flags for the proxy.
119    pub flags: ColliderTreeProxyFlags,
120}
121
122/// Flags for a [`ColliderTreeProxy`].
123#[repr(transparent)]
124#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Reflect)]
125#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
126#[cfg_attr(feature = "serialize", reflect(Serialize, Deserialize))]
127#[reflect(Debug, PartialEq)]
128pub struct ColliderTreeProxyFlags(u32);
129
130bitflags::bitflags! {
131    impl ColliderTreeProxyFlags: u32 {
132        /// Set if the collider is a sensor.
133        const SENSOR = 1 << 0;
134        /// Set if the body this collider is attached to has [`RigidBodyDisabled`](crate::dynamics::rigid_body::RigidBodyDisabled).
135        const BODY_DISABLED = 1 << 1;
136        /// Set if the custom filtering hook is active for this collider.
137        const CUSTOM_FILTER = 1 << 2;
138        /// Set if the contact modification hook is active for this collider.
139        const MODIFY_CONTACTS = 1 << 3;
140        /// Set if contact events are enabled for this collider.
141        const CONTACT_EVENTS = 1 << 4;
142    }
143}
144
145impl ColliderTreeProxyFlags {
146    /// Creates [`ColliderTreeProxyFlags`] from the given sensor status and active collision hooks.
147    #[inline]
148    pub fn new(
149        is_sensor: bool,
150        is_body_disabled: bool,
151        events_enabled: bool,
152        active_hooks: ActiveCollisionHooks,
153    ) -> Self {
154        let mut flags = ColliderTreeProxyFlags::empty();
155        if is_sensor {
156            flags |= ColliderTreeProxyFlags::SENSOR;
157        }
158        if is_body_disabled {
159            flags |= ColliderTreeProxyFlags::BODY_DISABLED;
160        }
161        if active_hooks.contains(ActiveCollisionHooks::FILTER_PAIRS) {
162            flags |= ColliderTreeProxyFlags::CUSTOM_FILTER;
163        }
164        if active_hooks.contains(ActiveCollisionHooks::MODIFY_CONTACTS) {
165            flags |= ColliderTreeProxyFlags::MODIFY_CONTACTS;
166        }
167        if events_enabled {
168            flags |= ColliderTreeProxyFlags::CONTACT_EVENTS;
169        }
170        flags
171    }
172}
173
174impl ColliderTreeProxy {
175    /// Returns `true` if the collider is a sensor.
176    #[inline]
177    pub fn is_sensor(&self) -> bool {
178        self.flags.contains(ColliderTreeProxyFlags::SENSOR)
179    }
180
181    /// Returns `true` if the custom filtering hook is active.
182    #[inline]
183    pub fn has_custom_filter(&self) -> bool {
184        self.flags.contains(ColliderTreeProxyFlags::CUSTOM_FILTER)
185    }
186
187    /// Returns `true` if the contact modification hook is active.
188    #[inline]
189    pub fn has_contact_modification(&self) -> bool {
190        self.flags.contains(ColliderTreeProxyFlags::MODIFY_CONTACTS)
191    }
192}
193
194/// A workspace for performing operations on a [`ColliderTree`].
195///
196/// This stores temporary data structures and working memory used when modifying the tree.
197/// It is recommended to reuse a single instance of this struct for all operations on a tree
198/// to avoid unnecessary allocations.
199#[derive(Resource, Default)]
200pub struct ColliderTreeWorkspace {
201    /// Builds the tree using PLOC (*Parallel, Locally Ordered Clustering*).
202    pub ploc_builder: PlocBuilder,
203    /// Restructures the BVH, optimizing node locations within the BVH hierarchy per SAH cost.
204    pub reinsertion_optimizer: ReinsertionOptimizer,
205    /// Temporary flagged nodes for partial rebuilds.
206    pub temp_flags: Vec<bool>,
207}
208
209impl Clone for ColliderTreeWorkspace {
210    fn clone(&self) -> Self {
211        Self {
212            ploc_builder: self.ploc_builder.clone(),
213            reinsertion_optimizer: ReinsertionOptimizer::default(),
214            temp_flags: Vec::new(),
215        }
216    }
217}
218
219impl ColliderTree {
220    /// Adds a proxy to the tree, returning its index.
221    #[inline]
222    pub fn add_proxy(&mut self, aabb: Aabb, proxy: ColliderTreeProxy) -> ProxyId {
223        let id = self.proxies.push(proxy) as u32;
224
225        // Insert the proxy into the BVH.
226        self.bvh.insert_primitive(aabb, id);
227
228        // Add to moved proxies.
229        self.moved_proxies.push(ProxyId::new(id));
230
231        ProxyId::new(id)
232    }
233
234    /// Removes a proxy from the tree.
235    ///
236    /// Returns `true` if the proxy was successfully removed, or `false` if the proxy ID was invalid.
237    #[inline]
238    pub fn remove_proxy(&mut self, proxy_id: ProxyId) -> Option<ColliderTreeProxy> {
239        if let Some(proxy) = self.proxies.try_remove(proxy_id.index()) {
240            // Remove from the BVH.
241            self.bvh.remove_primitive(proxy_id.id());
242
243            // Remove from moved proxies.
244            self.moved_proxies.remove(proxy_id);
245
246            Some(proxy)
247        } else {
248            None
249        }
250    }
251
252    /// Gets a proxy from the tree by its ID.
253    ///
254    /// Returns `None` if the proxy ID is invalid.
255    #[inline]
256    pub fn get_proxy(&self, proxy_id: ProxyId) -> Option<&ColliderTreeProxy> {
257        self.proxies.get(proxy_id.index())
258    }
259
260    /// Gets a mutable reference to a proxy from the tree by its ID.
261    ///
262    /// Returns `None` if the proxy ID is invalid.
263    #[inline]
264    pub fn get_proxy_mut(&mut self, proxy_id: ProxyId) -> Option<&mut ColliderTreeProxy> {
265        self.proxies.get_mut(proxy_id.index())
266    }
267
268    /// Gets a proxy from the tree by its ID without bounds checking.
269    ///
270    /// # Safety
271    ///
272    /// The caller must ensure that the `proxy_id` is valid.
273    #[inline]
274    pub unsafe fn get_proxy_unchecked(&self, proxy_id: ProxyId) -> &ColliderTreeProxy {
275        unsafe { self.proxies.get_unchecked(proxy_id.index()) }
276    }
277
278    /// Gets a mutable reference to a proxy from the tree by its ID without bounds checking.
279    ///
280    /// # Safety
281    ///
282    /// The caller must ensure that the `proxy_id` is valid.
283    #[inline]
284    pub unsafe fn get_proxy_unchecked_mut(&mut self, proxy_id: ProxyId) -> &mut ColliderTreeProxy {
285        unsafe { self.proxies.get_unchecked_mut(proxy_id.index()) }
286    }
287
288    /// Gets the AABB of a proxy in the tree.
289    ///
290    /// Returns `None` if the proxy ID is invalid.
291    #[inline]
292    pub fn get_proxy_aabb(&self, proxy_id: ProxyId) -> Option<Aabb> {
293        let node_id = self.bvh.primitives_to_nodes.get(proxy_id.index())?;
294        self.bvh.nodes.get(*node_id as usize).map(|node| node.aabb)
295    }
296
297    /// Gets the AABB of a proxy in the tree without bounds checking.
298    ///
299    /// # Safety
300    ///
301    /// The caller must ensure that the `proxy_id` is valid.
302    #[inline]
303    pub unsafe fn get_proxy_aabb_unchecked(&self, proxy_id: ProxyId) -> Aabb {
304        unsafe {
305            let node_id = *self.bvh.primitives_to_nodes.get_unchecked(proxy_id.index()) as usize;
306            self.bvh.nodes.get_unchecked(node_id).aabb
307        }
308    }
309
310    /// Updates the AABB of a proxy in the tree.
311    ///
312    /// If the BVH should be refitted at the same time, consider using
313    /// [`resize_proxy_aabb`](Self::resize_proxy_aabb) instead.
314    ///
315    /// If resizing a large number of proxies, consider calling this method
316    /// for each proxy and then calling [`refit_all`](Self::refit_all) once at the end.
317    #[inline]
318    pub fn set_proxy_aabb(&mut self, proxy_id: ProxyId, aabb: Aabb) {
319        // Get the node index for the proxy.
320        let node_index = self.bvh.primitives_to_nodes[proxy_id.index()] as usize;
321
322        // Update the proxy's AABB in the BVH.
323        self.bvh.nodes[node_index].set_aabb(aabb);
324    }
325
326    /// Resizes the AABB of a proxy in the tree.
327    ///
328    /// This is equivalent to calling [`set_proxy_aabb`](Self::set_proxy_aabb)
329    /// and then refitting the BVH working up from the resized node.
330    ///
331    /// For resizing a large number of proxies, consider calling [`set_proxy_aabb`](Self::set_proxy_aabb)
332    /// for each proxy and then calling [`refit_all`](Self::refit_all) once at the end.
333    #[inline]
334    pub fn resize_proxy_aabb(&mut self, proxy_id: ProxyId, aabb: Aabb) {
335        let node_index = self.bvh.primitives_to_nodes[proxy_id.index()] as usize;
336        self.bvh.resize_node(node_index, aabb);
337    }
338
339    /// Updates the AABB of a proxy and reinserts it at an optimal place in the tree.
340    #[inline]
341    pub fn reinsert_proxy(&mut self, proxy_id: ProxyId, aabb: Aabb) {
342        // Reinsert the node into the BVH.
343        let node_id = self.bvh.primitives_to_nodes[proxy_id.index()];
344        self.bvh.resize_node(node_id as usize, aabb);
345        self.bvh.reinsert_node(node_id as usize);
346    }
347
348    /// Refits the entire tree from the leaves up.
349    #[inline]
350    pub fn refit_all(&mut self) {
351        self.bvh.refit_all();
352    }
353
354    /// Fully rebuilds the tree from the given list of AABBs.
355    #[inline]
356    pub fn rebuild_full(&mut self) {
357        self.workspace.ploc_builder.full_rebuild(
358            &mut self.bvh,
359            PlocSearchDistance::Minimum,
360            SortPrecision::U64,
361            0,
362        );
363    }
364
365    /// Rebuilds parts of the tree corresponding to the given list of leaf node indices.
366    #[inline]
367    pub fn rebuild_partial(&mut self, leaves: &[u32]) {
368        self.bvh.init_parents_if_uninit();
369
370        // TODO: We could maybe get flagged nodes while refitting the tree.
371        compute_rebuild_path_flags(&self.bvh, leaves, &mut self.workspace.temp_flags);
372
373        self.workspace.ploc_builder.partial_rebuild(
374            &mut self.bvh,
375            |node_id| self.workspace.temp_flags[node_id],
376            PlocSearchDistance::Minimum,
377            SortPrecision::U64,
378            0,
379        );
380    }
381
382    /// Restructures the tree using parallel reinsertion, optimizing node locations based on SAH cost.
383    ///
384    /// This can be used to improve query performance after the tree quality has degraded,
385    /// for example after many proxy insertions and removals.
386    #[inline]
387    pub fn optimize(&mut self, batch_size_ratio: f32) {
388        self.workspace
389            .reinsertion_optimizer
390            .run(&mut self.bvh, batch_size_ratio, None);
391    }
392
393    /// Restructures the tree using parallel reinsertion, optimizing node locations based on SAH cost.
394    ///
395    /// Only the specified candidate proxies are considered for reinsertion.
396    #[inline]
397    pub fn optimize_candidates(&mut self, candidates: &[u32], iterations: u32) {
398        self.workspace.reinsertion_optimizer.run_with_candidates(
399            &mut self.bvh,
400            candidates,
401            iterations,
402        );
403    }
404}
avian3d/collider_tree/tree.rs

avian3d/collider_tree/
tree.rs
