quadtree_rs/
handle_iter.rs

1// Copyright 2019 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use {
16    crate::{area::Area, qtinner::QTInner, traversal::Traversal},
17    num::PrimInt,
18    std::{collections::HashSet, default::Default, iter::FusedIterator},
19};
20
21#[derive(Clone, Debug)]
22pub(crate) struct HandleIter<'a, U>
23where
24    U: PrimInt + Default,
25{
26    search_area: Area<U>,
27    handle_stack: Vec<u64>,
28    qt_stack: Vec<&'a QTInner<U>>,
29    visited: HashSet<u64>,
30}
31
32impl<'a, U> HandleIter<'a, U>
33where
34    U: PrimInt + Default,
35{
36    pub(crate) fn new(qt: &'a QTInner<U>, search_area: Area<U>) -> HandleIter<'a, U> {
37        HandleIter {
38            search_area,
39            handle_stack: vec![],
40            qt_stack: vec![qt],
41            visited: HashSet::new(),
42        }
43    }
44
45    // Descent is an optimization for queries. We don't want to traverse the entire tree searching
46    // for handles which (mostly) correspond to regions our @req doesn't intersect with.
47    //
48    // Instead, we can make a beeline for the lowest region which totally contains the @req (but no
49    // lower). We then have to actually evaluate every handle below that node.
50    //
51    // Along the way, if our query is meant to be of type Traversal::Overlapping, we collect the
52    // handles we meet along the way. They are guaranteed to intersect @req.
53    pub(crate) fn query_optimization(&mut self, req: Area<U>, traversal_method: Traversal) {
54        // This method expects to be called at a point in time when the HandleIter has just been
55        // created but has not yet been called.
56        assert!(self.qt_stack.len() == 1);
57        assert!(self.handle_stack.is_empty());
58        assert!(self.visited.is_empty());
59
60        self.descend_recurse_step(req, traversal_method);
61    }
62
63    fn descend_recurse_step(&mut self, req: Area<U>, traversal_method: Traversal) {
64        assert!(self.qt_stack.len() == 1);
65        // Peek into the stack. We have to peek rather than pop, because if we are about to go too
66        // far down we'd rather stop and return the HandleIter as-is.
67        if let Some(qt) = self.qt_stack.last() {
68            // If the region doesn't contain our @req, we're already too far down. Return here.
69            if !qt.region().contains(req) {
70                return;
71            }
72            assert!(qt.region().contains(req));
73
74            if let Some(subquadrants) = qt.subquadrants().as_ref() {
75                for subquadrant in subquadrants.iter() {
76                    // If we find a subquadrant which totally contains the @req, we want to make
77                    // that our new sole qt.
78                    if subquadrant.region().contains(req) {
79                        if traversal_method == Traversal::Overlapping {
80                            self.handle_stack.extend(qt.handles());
81                        }
82
83                        // TODO(ambuc): Could this be done with Vec::swap() or std::mem::replace()?
84                        assert!(self.qt_stack.len() == 1);
85                        self.qt_stack = vec![subquadrant];
86
87                        // Recurse on this step. It will naturally return, but we want to propogate
88                        // that return rather than continue to search the other subquadrants.
89                        return self.descend_recurse_step(req, traversal_method);
90                    }
91                }
92            }
93            // If there aren't any subquadrants, we're probably done.
94        }
95    }
96}
97
98impl<U> Iterator for HandleIter<'_, U>
99where
100    U: PrimInt + Default,
101{
102    type Item = u64;
103
104    #[inline]
105    fn next(&mut self) -> Option<Self::Item> {
106        loop {
107            while let Some(handle) = self.handle_stack.pop() {
108                if self.visited.insert(handle) {
109                    return Some(handle);
110                }
111            }
112
113            // Then check the qt_stack.
114            if let Some(qt) = self.qt_stack.pop() {
115                // Push my sub quadrants onto the qt_stack too.
116                if let Some(sub_quadrants) = qt.subquadrants().as_ref() {
117                    for sub_quadrant in sub_quadrants {
118                        if sub_quadrant.region().intersects(self.search_area) {
119                            self.qt_stack.push(sub_quadrant)
120                        }
121                    }
122                }
123
124                // Push my regions onto the region stack
125                match qt.handles().len() {
126                    0 => (),
127                    1 => {
128                        if self.visited.insert(qt.handles()[0]) {
129                            return Some(qt.handles()[0]);
130                        }
131                    }
132                    _ => self.handle_stack.extend(qt.handles()),
133                }
134
135                continue;
136            }
137
138            // Else there's nothing left to search.
139            return None;
140        }
141    }
142
143    #[inline]
144    fn size_hint(&self) -> (usize, Option<usize>) {
145        (0, None)
146    }
147}
148
149impl<U> FusedIterator for HandleIter<'_, U> where U: PrimInt + Default {}
quadtree_rs/handle_iter.rs

quadtree_rs/
handle_iter.rs
