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

avian3d/collider_tree/
tree.rs
