kiddo/kd_tree/
query_stack.rs

1use std::mem::MaybeUninit;
2
3use crate::{Axis, StemStrategy};
4
5const INLINE_QUERY_STACK_CAPACITY: usize = 50;
6#[allow(dead_code)]
7const INLINE_SCALAR_CONTINUATION_STACK_CAPACITY: usize = 64;
8
9pub trait ScalarStackContext<A, S>: Sized {
10    fn from_parts(stem_state: S, old_off: A, rd: A) -> Self;
11    fn into_parts(self) -> (S, A, A);
12
13    #[inline(always)]
14    fn from_parts_with_restore_dim(stem_state: S, _restore_dim: usize, old_off: A, rd: A) -> Self {
15        Self::from_parts(stem_state, old_off, rd)
16    }
17
18    #[inline(always)]
19    fn into_parts_with_restore_dim(self) -> (S, Option<usize>, A, A) {
20        let (stem_state, old_off, rd) = self.into_parts();
21        (stem_state, None, old_off, rd)
22    }
23}
24
25/// Trait for query stack types to enable generic backtracking implementations
26pub trait StackTrait<A, SS: StemStrategy, const K: usize> {
27    fn push(&mut self, item: SS::StackContext<A, K>);
28    fn pop(&mut self) -> Option<SS::StackContext<A, K>>;
29    fn clear(&mut self);
30}
31
32/// Reusable traversal scratch storage for query builder scratch APIs.
33///
34/// The concrete stack representation is selected by the stem strategy and is
35/// intentionally hidden so it can keep changing during the v6 alpha cycle.
36#[derive(Debug)]
37pub struct QueryScratch<SS: StemStrategy, O, const K: usize> {
38    stack: SS::Stack<O, K>,
39}
40
41impl<SS: StemStrategy, O, const K: usize> Default for QueryScratch<SS, O, K> {
42    #[inline]
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl<SS: StemStrategy, O, const K: usize> QueryScratch<SS, O, K> {
49    /// Creates empty query traversal scratch storage.
50    #[inline]
51    pub fn new() -> Self {
52        Self {
53            stack: SS::Stack::<O, K>::default(),
54        }
55    }
56
57    #[inline]
58    pub(crate) fn clear(&mut self) {
59        self.stack.clear();
60    }
61
62    #[inline]
63    pub(crate) fn stack_mut(&mut self) -> &mut SS::Stack<O, K> {
64        &mut self.stack
65    }
66}
67
68#[allow(dead_code)]
69#[derive(Debug)]
70pub(crate) struct ScalarContinuationFar<A, S> {
71    pub stem_state: S,
72    pub far_off: A,
73    pub rd: A,
74}
75
76#[allow(dead_code)]
77#[derive(Debug, Clone, Copy)]
78pub(crate) struct ScalarContinuationRestore<A> {
79    pub old_off: A,
80    pub has_far: bool,
81}
82
83#[allow(dead_code)]
84impl<A> ScalarContinuationRestore<A> {
85    #[inline(always)]
86    pub fn restore_only(old_off: A) -> Self {
87        Self {
88            old_off,
89            has_far: false,
90        }
91    }
92
93    #[inline(always)]
94    pub fn with_far(old_off: A) -> Self {
95        Self {
96            old_off,
97            has_far: true,
98        }
99    }
100}
101
102#[allow(dead_code)]
103#[derive(Debug)]
104pub(crate) struct ScalarContinuationRestoreStack<
105    A,
106    const INLINE_CAPACITY: usize = INLINE_SCALAR_CONTINUATION_STACK_CAPACITY,
107> {
108    stack: [MaybeUninit<ScalarContinuationRestore<A>>; INLINE_CAPACITY],
109    len: usize,
110}
111
112#[allow(dead_code)]
113#[derive(Debug)]
114pub(crate) struct ScalarContinuationFarStack<
115    A,
116    S,
117    const INLINE_CAPACITY: usize = INLINE_SCALAR_CONTINUATION_STACK_CAPACITY,
118> {
119    stack: [MaybeUninit<ScalarContinuationFar<A, S>>; INLINE_CAPACITY],
120    len: usize,
121}
122
123#[derive(Debug)]
124pub struct QueryStack<A, SS: StemStrategy, const K: usize> {
125    stack: [MaybeUninit<SS::StackContext<A, K>>; INLINE_QUERY_STACK_CAPACITY],
126    spill: Vec<SS::StackContext<A, K>>,
127    len: usize,
128}
129
130impl<A, SS: StemStrategy, const K: usize> Default for QueryStack<A, SS, K> {
131    fn default() -> Self {
132        Self::new()
133    }
134}
135
136#[allow(dead_code)]
137impl<A, const INLINE_CAPACITY: usize> Default
138    for ScalarContinuationRestoreStack<A, INLINE_CAPACITY>
139{
140    fn default() -> Self {
141        Self::new()
142    }
143}
144
145#[allow(dead_code)]
146impl<A, S, const INLINE_CAPACITY: usize> Default
147    for ScalarContinuationFarStack<A, S, INLINE_CAPACITY>
148{
149    fn default() -> Self {
150        Self::new()
151    }
152}
153
154#[derive(Debug)]
155pub struct QueryStackContext<A, S> {
156    pub stem_state: S,
157    pub restore_dim: Option<usize>,
158    pub old_off: A,
159    pub rd: A,
160}
161
162impl<A, SS: StemStrategy, const K: usize> StackTrait<A, SS, K> for QueryStack<A, SS, K> {
163    #[inline]
164    fn push(&mut self, item: SS::StackContext<A, K>) {
165        if self.len < INLINE_QUERY_STACK_CAPACITY {
166            unsafe { self.stack.get_unchecked_mut(self.len) }.write(item);
167        } else {
168            self.spill.push(item);
169        }
170        self.len += 1;
171    }
172
173    #[inline]
174    fn pop(&mut self) -> Option<SS::StackContext<A, K>> {
175        if self.len == 0 {
176            None
177        } else {
178            self.len -= 1;
179            if self.len >= INLINE_QUERY_STACK_CAPACITY {
180                Some(self.spill.pop().expect("query stack spill underflow"))
181            } else {
182                Some(unsafe { self.stack.get_unchecked(self.len).assume_init_read() })
183            }
184        }
185    }
186
187    #[inline]
188    fn clear(&mut self) {
189        if self.len > INLINE_QUERY_STACK_CAPACITY {
190            self.spill.clear();
191            self.len = INLINE_QUERY_STACK_CAPACITY;
192        }
193
194        while self.len > 0 {
195            self.len -= 1;
196            unsafe { self.stack.get_unchecked_mut(self.len).assume_init_drop() };
197        }
198    }
199}
200
201impl<A, SS: StemStrategy, const K: usize> QueryStack<A, SS, K> {
202    #[inline]
203    pub fn new() -> Self {
204        Self {
205            stack: [const { MaybeUninit::uninit() }; INLINE_QUERY_STACK_CAPACITY],
206            spill: Vec::new(),
207            len: 0,
208        }
209    }
210
211    #[inline]
212    pub fn push(&mut self, item: SS::StackContext<A, K>) {
213        <Self as StackTrait<A, SS, K>>::push(self, item);
214    }
215
216    #[inline]
217    pub fn pop(&mut self) -> Option<SS::StackContext<A, K>> {
218        <Self as StackTrait<A, SS, K>>::pop(self)
219    }
220}
221
222impl<A, SS: StemStrategy, const K: usize> Drop for QueryStack<A, SS, K> {
223    fn drop(&mut self) {
224        self.clear();
225    }
226}
227
228#[allow(dead_code)]
229impl<A, const INLINE_CAPACITY: usize> ScalarContinuationRestoreStack<A, INLINE_CAPACITY> {
230    #[inline]
231    pub fn new() -> Self {
232        Self {
233            stack: [const { MaybeUninit::uninit() }; INLINE_CAPACITY],
234            len: 0,
235        }
236    }
237
238    #[inline]
239    pub fn push_unchecked_inline(&mut self, item: ScalarContinuationRestore<A>) {
240        debug_assert!(self.len < INLINE_CAPACITY);
241        unsafe { self.stack.get_unchecked_mut(self.len) }.write(item);
242        self.len += 1;
243    }
244
245    #[inline]
246    pub fn pop(&mut self) -> Option<ScalarContinuationRestore<A>> {
247        if self.len == 0 {
248            None
249        } else {
250            self.len -= 1;
251            Some(unsafe { self.stack.get_unchecked(self.len).assume_init_read() })
252        }
253    }
254
255    #[inline]
256    pub fn clear(&mut self) {
257        while self.len > 0 {
258            self.len -= 1;
259            unsafe { self.stack.get_unchecked_mut(self.len).assume_init_drop() };
260        }
261    }
262}
263
264#[allow(dead_code)]
265impl<A, S, const INLINE_CAPACITY: usize> ScalarContinuationFarStack<A, S, INLINE_CAPACITY> {
266    #[inline]
267    pub fn new() -> Self {
268        Self {
269            stack: [const { MaybeUninit::uninit() }; INLINE_CAPACITY],
270            len: 0,
271        }
272    }
273
274    #[inline]
275    pub fn push_unchecked_inline(&mut self, item: ScalarContinuationFar<A, S>) {
276        debug_assert!(self.len < INLINE_CAPACITY);
277        unsafe { self.stack.get_unchecked_mut(self.len) }.write(item);
278        self.len += 1;
279    }
280
281    #[inline]
282    pub fn pop(&mut self) -> Option<ScalarContinuationFar<A, S>> {
283        if self.len == 0 {
284            None
285        } else {
286            self.len -= 1;
287            Some(unsafe { self.stack.get_unchecked(self.len).assume_init_read() })
288        }
289    }
290
291    #[inline]
292    pub fn clear(&mut self) {
293        while self.len > 0 {
294            self.len -= 1;
295            unsafe { self.stack.get_unchecked_mut(self.len).assume_init_drop() };
296        }
297    }
298}
299
300#[allow(dead_code)]
301impl<A, const INLINE_CAPACITY: usize> Drop for ScalarContinuationRestoreStack<A, INLINE_CAPACITY> {
302    fn drop(&mut self) {
303        self.clear();
304    }
305}
306
307#[allow(dead_code)]
308impl<A, S, const INLINE_CAPACITY: usize> Drop
309    for ScalarContinuationFarStack<A, S, INLINE_CAPACITY>
310{
311    fn drop(&mut self) {
312        self.clear();
313    }
314}
315
316impl<A: Axis<Coord = A>, S> QueryStackContext<A, S> {
317    pub fn new(stem_state: S) -> Self {
318        Self {
319            stem_state,
320            restore_dim: None,
321            old_off: A::zero(),
322            rd: A::zero(),
323        }
324    }
325}
326
327impl<A, S> QueryStackContext<A, S> {
328    pub fn into_parts(self) -> (S, A, A) {
329        (self.stem_state, self.old_off, self.rd)
330    }
331}
332
333impl<A, S> ScalarStackContext<A, S> for QueryStackContext<A, S> {
334    #[inline(always)]
335    fn from_parts(stem_state: S, old_off: A, rd: A) -> Self {
336        Self {
337            stem_state,
338            restore_dim: None,
339            old_off,
340            rd,
341        }
342    }
343
344    #[inline(always)]
345    fn from_parts_with_restore_dim(stem_state: S, restore_dim: usize, old_off: A, rd: A) -> Self {
346        Self {
347            stem_state,
348            restore_dim: Some(restore_dim),
349            old_off,
350            rd,
351        }
352    }
353
354    #[inline(always)]
355    fn into_parts(self) -> (S, A, A) {
356        QueryStackContext::into_parts(self)
357    }
358
359    #[inline(always)]
360    fn into_parts_with_restore_dim(self) -> (S, Option<usize>, A, A) {
361        (self.stem_state, self.restore_dim, self.old_off, self.rd)
362    }
363}
kiddo/kd_tree/query_stack.rs

kiddo/kd_tree/
query_stack.rs
