Skip to main content

box2d_rust/dynamic_tree/
query.rs

1// AABB query, ray cast, and box cast traversals from dynamic_tree.c.
2//
3// The C callbacks take a `void* context`; the Rust versions take closures,
4// which capture their context directly.
5//
6// SPDX-FileCopyrightText: 2023 Erin Catto
7// SPDX-License-Identifier: MIT
8
9use super::{BoxCastInput, DynamicTree, TreeStats, TREE_STACK_SIZE};
10use crate::collision::RayCastInput;
11use crate::core::NULL_INDEX;
12use crate::math_functions::{
13    aabb_center, aabb_extents, aabb_overlaps, abs, abs_float, add, cross_sv, distance_squared, dot,
14    max, min, mul_add, mul_sv, normalize, sub, Aabb,
15};
16
17impl DynamicTree {
18    /// Query an AABB for overlapping proxies. The callback is called for each
19    /// proxy that overlaps the supplied AABB and passes the mask-bits filter;
20    /// return false from the callback to stop. (b2DynamicTree_Query)
21    pub fn query(
22        &self,
23        aabb: Aabb,
24        mask_bits: u64,
25        mut callback: impl FnMut(i32, u64) -> bool,
26    ) -> TreeStats {
27        let mut result = TreeStats::default();
28
29        if self.node_count == 0 {
30            return result;
31        }
32
33        let mut stack = [0i32; TREE_STACK_SIZE];
34        let mut stack_count = 0usize;
35        stack[stack_count] = self.root;
36        stack_count += 1;
37
38        while stack_count > 0 {
39            stack_count -= 1;
40            let node_id = stack[stack_count];
41
42            let node = &self.nodes[node_id as usize];
43            result.node_visits += 1;
44
45            if aabb_overlaps(node.aabb, aabb) && (node.category_bits & mask_bits) != 0 {
46                if node.is_leaf() {
47                    // callback to user code with proxy id
48                    let proceed = callback(node_id, node.user_data);
49                    result.leaf_visits += 1;
50
51                    if !proceed {
52                        return result;
53                    }
54                } else if stack_count < TREE_STACK_SIZE - 1 {
55                    stack[stack_count] = node.child1;
56                    stack_count += 1;
57                    stack[stack_count] = node.child2;
58                    stack_count += 1;
59                } else {
60                    debug_assert!(stack_count < TREE_STACK_SIZE - 1);
61                }
62            }
63        }
64
65        result
66    }
67
68    /// Query an AABB for overlapping proxies with no filtering.
69    /// (b2DynamicTree_QueryAll)
70    pub fn query_all(&self, aabb: Aabb, mut callback: impl FnMut(i32, u64) -> bool) -> TreeStats {
71        let mut result = TreeStats::default();
72
73        if self.node_count == 0 {
74            return result;
75        }
76
77        let mut stack = [0i32; TREE_STACK_SIZE];
78        let mut stack_count = 0usize;
79        stack[stack_count] = self.root;
80        stack_count += 1;
81
82        while stack_count > 0 {
83            stack_count -= 1;
84            let node_id = stack[stack_count];
85
86            let node = &self.nodes[node_id as usize];
87            result.node_visits += 1;
88
89            if aabb_overlaps(node.aabb, aabb) {
90                if node.is_leaf() {
91                    // callback to user code with proxy id
92                    let proceed = callback(node_id, node.user_data);
93                    result.leaf_visits += 1;
94
95                    if !proceed {
96                        return result;
97                    }
98                } else if stack_count < TREE_STACK_SIZE - 1 {
99                    stack[stack_count] = node.child1;
100                    stack_count += 1;
101                    stack[stack_count] = node.child2;
102                    stack_count += 1;
103                } else {
104                    debug_assert!(stack_count < TREE_STACK_SIZE - 1);
105                }
106            }
107        }
108
109        result
110    }
111
112    /// Ray cast against the proxies in the tree. The callback performs an
113    /// exact ray cast when the proxy contains a shape, and returns the new
114    /// ray fraction:
115    /// - return 0 to terminate the ray cast
116    /// - return a value less than the input max_fraction to clip the ray
117    /// - return the input max_fraction to continue without clipping
118    ///
119    /// (b2DynamicTree_RayCast)
120    pub fn ray_cast(
121        &self,
122        input: &RayCastInput,
123        mask_bits: u64,
124        mut callback: impl FnMut(&RayCastInput, i32, u64) -> f32,
125    ) -> TreeStats {
126        let mut result = TreeStats::default();
127
128        if self.node_count == 0 {
129            return result;
130        }
131
132        let p1 = input.origin;
133        let d = input.translation;
134
135        let r = normalize(d);
136
137        // v is perpendicular to the segment.
138        let v = cross_sv(1.0, r);
139        let abs_v = abs(v);
140
141        // Separating axis for segment (Gino, p80).
142        // |dot(v, p1 - c)| > dot(|v|, h)
143
144        let mut max_fraction = input.max_fraction;
145
146        let mut p2 = mul_add(p1, max_fraction, d);
147
148        // Build a bounding box for the segment.
149        let mut segment_aabb = Aabb {
150            lower_bound: min(p1, p2),
151            upper_bound: max(p1, p2),
152        };
153
154        let mut stack = [0i32; TREE_STACK_SIZE];
155        let mut stack_count = 0usize;
156        stack[stack_count] = self.root;
157        stack_count += 1;
158
159        let mut sub_input = *input;
160
161        while stack_count > 0 {
162            stack_count -= 1;
163            let node_id = stack[stack_count];
164            if node_id == NULL_INDEX {
165                debug_assert!(false);
166                continue;
167            }
168
169            let node = &self.nodes[node_id as usize];
170            result.node_visits += 1;
171
172            let node_aabb = node.aabb;
173
174            if (node.category_bits & mask_bits) == 0 || !aabb_overlaps(node_aabb, segment_aabb) {
175                continue;
176            }
177
178            // Separating axis for segment (Gino, p80).
179            // |dot(v, p1 - c)| > dot(|v|, h)
180            // radius extension is added to the node in this case
181            let c = aabb_center(node_aabb);
182            let h = aabb_extents(node_aabb);
183            let term1 = abs_float(dot(v, sub(p1, c)));
184            let term2 = dot(abs_v, h);
185            if term2 < term1 {
186                continue;
187            }
188
189            if node.is_leaf() {
190                sub_input.max_fraction = max_fraction;
191
192                let value = callback(&sub_input, node_id, node.user_data);
193                result.leaf_visits += 1;
194
195                // The user may return -1 to indicate this shape should be skipped
196
197                if value == 0.0 {
198                    // The client has terminated the ray cast.
199                    return result;
200                }
201
202                if 0.0 < value && value <= max_fraction {
203                    // Update segment bounding box.
204                    max_fraction = value;
205                    p2 = mul_add(p1, max_fraction, d);
206                    segment_aabb.lower_bound = min(p1, p2);
207                    segment_aabb.upper_bound = max(p1, p2);
208                }
209            } else if stack_count < TREE_STACK_SIZE - 1 {
210                let c1 = aabb_center(self.nodes[node.child1 as usize].aabb);
211                let c2 = aabb_center(self.nodes[node.child2 as usize].aabb);
212                if distance_squared(c1, p1) < distance_squared(c2, p1) {
213                    stack[stack_count] = node.child2;
214                    stack_count += 1;
215                    stack[stack_count] = node.child1;
216                    stack_count += 1;
217                } else {
218                    stack[stack_count] = node.child1;
219                    stack_count += 1;
220                    stack[stack_count] = node.child2;
221                    stack_count += 1;
222                }
223            } else {
224                debug_assert!(stack_count < TREE_STACK_SIZE - 1);
225            }
226        }
227
228        result
229    }
230
231    /// Cast a swept AABB through the tree. The callback returns the new cast
232    /// fraction, with the same semantics as [`DynamicTree::ray_cast`].
233    /// (b2DynamicTree_BoxCast)
234    pub fn box_cast(
235        &self,
236        input: &BoxCastInput,
237        mask_bits: u64,
238        mut callback: impl FnMut(&BoxCastInput, i32, u64) -> f32,
239    ) -> TreeStats {
240        let mut stats = TreeStats::default();
241
242        if self.node_count == 0 {
243            return stats;
244        }
245
246        // The caller folds the shape radius into the box
247        let origin_aabb = input.box_;
248
249        let p1 = aabb_center(origin_aabb);
250        let extension = aabb_extents(origin_aabb);
251
252        // v is perpendicular to the segment.
253        let r = input.translation;
254        let v = cross_sv(1.0, r);
255        let abs_v = abs(v);
256
257        // Separating axis for segment (Gino, p80).
258        // |dot(v, p1 - c)| > dot(|v|, h)
259
260        let mut max_fraction = input.max_fraction;
261
262        // Build total box for the cast
263        let mut t = mul_sv(max_fraction, input.translation);
264        let mut total_aabb = Aabb {
265            lower_bound: min(origin_aabb.lower_bound, add(origin_aabb.lower_bound, t)),
266            upper_bound: max(origin_aabb.upper_bound, add(origin_aabb.upper_bound, t)),
267        };
268
269        let mut sub_input = *input;
270
271        let mut stack = [0i32; TREE_STACK_SIZE];
272        let mut stack_count = 0usize;
273        stack[stack_count] = self.root;
274        stack_count += 1;
275
276        while stack_count > 0 {
277            stack_count -= 1;
278            let node_id = stack[stack_count];
279            if node_id == NULL_INDEX {
280                debug_assert!(false);
281                continue;
282            }
283
284            let node = &self.nodes[node_id as usize];
285            stats.node_visits += 1;
286
287            if (node.category_bits & mask_bits) == 0 || !aabb_overlaps(node.aabb, total_aabb) {
288                continue;
289            }
290
291            // Separating axis for segment (Gino, p80).
292            // |dot(v, p1 - c)| > dot(|v|, h)
293            // radius extension is added to the node in this case
294            let c = aabb_center(node.aabb);
295            let h = add(aabb_extents(node.aabb), extension);
296            let term1 = abs_float(dot(v, sub(p1, c)));
297            let term2 = dot(abs_v, h);
298            if term2 < term1 {
299                continue;
300            }
301
302            if node.is_leaf() {
303                sub_input.max_fraction = max_fraction;
304
305                let value = callback(&sub_input, node_id, node.user_data);
306                stats.leaf_visits += 1;
307
308                if value == 0.0 {
309                    // The client has terminated the ray cast.
310                    return stats;
311                }
312
313                if 0.0 < value && value < max_fraction {
314                    // Update segment bounding box.
315                    max_fraction = value;
316                    t = mul_sv(max_fraction, input.translation);
317                    total_aabb.lower_bound =
318                        min(origin_aabb.lower_bound, add(origin_aabb.lower_bound, t));
319                    total_aabb.upper_bound =
320                        max(origin_aabb.upper_bound, add(origin_aabb.upper_bound, t));
321                }
322            } else if stack_count < TREE_STACK_SIZE - 1 {
323                let c1 = aabb_center(self.nodes[node.child1 as usize].aabb);
324                let c2 = aabb_center(self.nodes[node.child2 as usize].aabb);
325                if distance_squared(c1, p1) < distance_squared(c2, p1) {
326                    stack[stack_count] = node.child2;
327                    stack_count += 1;
328                    stack[stack_count] = node.child1;
329                    stack_count += 1;
330                } else {
331                    stack[stack_count] = node.child1;
332                    stack_count += 1;
333                    stack[stack_count] = node.child2;
334                    stack_count += 1;
335                }
336            } else {
337                debug_assert!(stack_count < TREE_STACK_SIZE - 1);
338            }
339        }
340
341        stats
342    }
343}