avian2d/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    /// Temporary flagged nodes for partial rebuilds.
137    pub temp_flags: Vec<bool>,
138}
139
140impl Clone for ColliderTreeWorkspace {
141    fn clone(&self) -> Self {
142        Self {
143            ploc_builder: self.ploc_builder.clone(),
144            reinsertion_optimizer: ReinsertionOptimizer::default(),
145            insertion_stack: self.insertion_stack.clone(),
146            temp_flags: Vec::new(),
147        }
148    }
149}
150
151impl Default for ColliderTreeWorkspace {
152    fn default() -> Self {
153        Self {
154            ploc_builder: PlocBuilder::default(),
155            reinsertion_optimizer: ReinsertionOptimizer::default(),
156            insertion_stack: HeapStack::new_with_capacity(2000),
157            temp_flags: Vec::new(),
158        }
159    }
160}
161
162impl ColliderTree {
163    /// Adds a proxy to the tree, returning its index.
164    #[inline]
165    pub fn add_proxy(&mut self, aabb: Aabb, proxy: ColliderTreeProxy) -> ProxyId {
166        let id = self.proxies.push(proxy) as u32;
167
168        // Insert the proxy into the BVH.
169        self.bvh
170            .insert_primitive(aabb, id, &mut self.workspace.insertion_stack);
171
172        // Add to moved proxies.
173        self.moved_proxies.push(ProxyId::new(id));
174
175        ProxyId::new(id)
176    }
177
178    /// Removes a proxy from the tree.
179    ///
180    /// Returns `true` if the proxy was successfully removed, or `false` if the proxy ID was invalid.
181    #[inline]
182    pub fn remove_proxy(&mut self, proxy_id: ProxyId) -> Option<ColliderTreeProxy> {
183        if let Some(proxy) = self.proxies.try_remove(proxy_id.index()) {
184            // Remove from the BVH.
185            self.bvh.remove_primitive(proxy_id.id());
186
187            // Remove from moved proxies.
188            for i in 0..self.moved_proxies.len() {
189                if self.moved_proxies[i] == proxy_id {
190                    self.moved_proxies.swap_remove(i);
191                    break;
192                }
193            }
194
195            Some(proxy)
196        } else {
197            None
198        }
199    }
200
201    /// Gets a proxy from the tree by its ID.
202    ///
203    /// Returns `None` if the proxy ID is invalid.
204    #[inline]
205    pub fn get_proxy(&self, proxy_id: ProxyId) -> Option<&ColliderTreeProxy> {
206        self.proxies.get(proxy_id.index())
207    }
208
209    /// Gets a mutable reference to a proxy from the tree by its ID.
210    ///
211    /// Returns `None` if the proxy ID is invalid.
212    #[inline]
213    pub fn get_proxy_mut(&mut self, proxy_id: ProxyId) -> Option<&mut ColliderTreeProxy> {
214        self.proxies.get_mut(proxy_id.index())
215    }
216
217    /// Gets a proxy from the tree by its ID without bounds checking.
218    ///
219    /// # Safety
220    ///
221    /// The caller must ensure that the `proxy_id` is valid.
222    #[inline]
223    pub unsafe fn get_proxy_unchecked(&self, proxy_id: ProxyId) -> &ColliderTreeProxy {
224        unsafe { self.proxies.get_unchecked(proxy_id.index()) }
225    }
226
227    /// Gets a mutable reference to a proxy from the tree by its ID without bounds checking.
228    ///
229    /// # Safety
230    ///
231    /// The caller must ensure that the `proxy_id` is valid.
232    #[inline]
233    pub unsafe fn get_proxy_unchecked_mut(&mut self, proxy_id: ProxyId) -> &mut ColliderTreeProxy {
234        unsafe { self.proxies.get_unchecked_mut(proxy_id.index()) }
235    }
236
237    /// Gets the AABB of a proxy in the tree.
238    ///
239    /// Returns `None` if the proxy ID is invalid.
240    #[inline]
241    pub fn get_proxy_aabb(&self, proxy_id: ProxyId) -> Option<Aabb> {
242        let node_id = self.bvh.primitives_to_nodes.get(proxy_id.index())?;
243        self.bvh.nodes.get(*node_id as usize).map(|node| node.aabb)
244    }
245
246    /// Gets the AABB of a proxy in the tree without bounds checking.
247    ///
248    /// # Safety
249    ///
250    /// The caller must ensure that the `proxy_id` is valid.
251    #[inline]
252    pub unsafe fn get_proxy_aabb_unchecked(&self, proxy_id: ProxyId) -> Aabb {
253        unsafe {
254            let node_id = *self.bvh.primitives_to_nodes.get_unchecked(proxy_id.index()) as usize;
255            self.bvh.nodes.get_unchecked(node_id).aabb
256        }
257    }
258
259    /// Updates the AABB of a proxy in the tree.
260    ///
261    /// If the BVH should be refitted at the same time, consider using
262    /// [`resize_proxy_aabb`](Self::resize_proxy_aabb) instead.
263    ///
264    /// If resizing a large number of proxies, consider calling this method
265    /// for each proxy and then calling [`refit_all`](Self::refit_all) once at the end.
266    #[inline]
267    pub fn set_proxy_aabb(&mut self, proxy_id: ProxyId, aabb: Aabb) {
268        // Get the node index for the proxy.
269        let node_index = self.bvh.primitives_to_nodes[proxy_id.index()] as usize;
270
271        // Update the proxy's AABB in the BVH.
272        self.bvh.nodes[node_index].set_aabb(aabb);
273    }
274
275    /// Resizes the AABB of a proxy in the tree.
276    ///
277    /// This is equivalent to calling [`set_proxy_aabb`](Self::set_proxy_aabb)
278    /// and then refitting the BVH working up from the resized node.
279    ///
280    /// For resizing a large number of proxies, consider calling [`set_proxy_aabb`](Self::set_proxy_aabb)
281    /// for each proxy and then calling [`refit_all`](Self::refit_all) once at the end.
282    #[inline]
283    pub fn resize_proxy_aabb(&mut self, proxy_id: ProxyId, aabb: Aabb) {
284        let node_index = self.bvh.primitives_to_nodes[proxy_id.index()] as usize;
285        self.bvh.resize_node(node_index, aabb);
286    }
287
288    /// Updates the AABB of a proxy and reinserts it at an optimal place in the tree.
289    #[inline]
290    pub fn reinsert_proxy(&mut self, proxy_id: ProxyId, aabb: Aabb) {
291        // Reinsert the node into the BVH.
292        let node_id = self.bvh.primitives_to_nodes[proxy_id.index()];
293        self.bvh.resize_node(node_id as usize, aabb);
294        self.bvh.reinsert_node(node_id as usize);
295    }
296
297    /// Refits the entire tree from the leaves up.
298    #[inline]
299    pub fn refit_all(&mut self) {
300        self.bvh.refit_all();
301    }
302
303    /// Fully rebuilds the tree from the given list of AABBs.
304    #[inline]
305    pub fn rebuild_full(&mut self) {
306        self.workspace.ploc_builder.full_rebuild(
307            &mut self.bvh,
308            PlocSearchDistance::Minimum,
309            SortPrecision::U64,
310            0,
311        );
312    }
313
314    /// Rebuilds parts of the tree corresponding to the given list of leaf node indices.
315    #[inline]
316    pub fn rebuild_partial(&mut self, leaves: &[u32]) {
317        self.bvh.init_parents_if_uninit();
318
319        // TODO: We could maybe get flagged nodes while refitting the tree.
320        compute_rebuild_path_flags(&self.bvh, leaves, &mut self.workspace.temp_flags);
321
322        self.workspace.ploc_builder.partial_rebuild(
323            &mut self.bvh,
324            |node_id| self.workspace.temp_flags[node_id],
325            PlocSearchDistance::Minimum,
326            SortPrecision::U64,
327            0,
328        );
329    }
330
331    /// Restructures the tree using parallel reinsertion, optimizing node locations based on SAH cost.
332    ///
333    /// This can be used to improve query performance after the tree quality has degraded,
334    /// for example after many proxy insertions and removals.
335    #[inline]
336    pub fn optimize(&mut self, batch_size_ratio: f32) {
337        self.workspace
338            .reinsertion_optimizer
339            .run(&mut self.bvh, batch_size_ratio, None);
340    }
341
342    /// Restructures the tree using parallel reinsertion, optimizing node locations based on SAH cost.
343    ///
344    /// Only the specified candidate proxies are considered for reinsertion.
345    #[inline]
346    pub fn optimize_candidates(&mut self, candidates: &[u32], iterations: u32) {
347        self.workspace.reinsertion_optimizer.run_with_candidates(
348            &mut self.bvh,
349            candidates,
350            iterations,
351        );
352    }
353}
avian2d/collider_tree/tree.rs

avian2d/collider_tree/
tree.rs
