1 | use core::borrow::Borrow; |
2 | use core::cmp::Ordering; |
3 | use core::ops::{Bound, RangeBounds}; |
4 | |
5 | use super::node::{marker, ForceResult::*, Handle, NodeRef}; |
6 | |
7 | use SearchBound::*; |
8 | use SearchResult::*; |
9 | |
10 | pub enum SearchBound<T> { |
11 | /// An inclusive bound to look for, just like `Bound::Included(T)`. |
12 | Included(T), |
13 | /// An exclusive bound to look for, just like `Bound::Excluded(T)`. |
14 | Excluded(T), |
15 | /// An unconditional inclusive bound, just like `Bound::Unbounded`. |
16 | AllIncluded, |
17 | /// An unconditional exclusive bound. |
18 | AllExcluded, |
19 | } |
20 | |
21 | impl<T> SearchBound<T> { |
22 | pub fn from_range(range_bound: Bound<T>) -> Self { |
23 | match range_bound { |
24 | Bound::Included(t: T) => Included(t), |
25 | Bound::Excluded(t: T) => Excluded(t), |
26 | Bound::Unbounded => AllIncluded, |
27 | } |
28 | } |
29 | } |
30 | |
31 | pub enum SearchResult<BorrowType, K, V, FoundType, GoDownType> { |
32 | Found(Handle<NodeRef<BorrowType, K, V, FoundType>, marker::KV>), |
33 | GoDown(Handle<NodeRef<BorrowType, K, V, GoDownType>, marker::Edge>), |
34 | } |
35 | |
36 | pub enum IndexResult { |
37 | KV(usize), |
38 | Edge(usize), |
39 | } |
40 | |
41 | impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> { |
42 | /// Looks up a given key in a (sub)tree headed by the node, recursively. |
43 | /// Returns a `Found` with the handle of the matching KV, if any. Otherwise, |
44 | /// returns a `GoDown` with the handle of the leaf edge where the key belongs. |
45 | /// |
46 | /// The result is meaningful only if the tree is ordered by key, like the tree |
47 | /// in a `BTreeMap` is. |
48 | pub fn search_tree<Q: ?Sized>( |
49 | mut self, |
50 | key: &Q, |
51 | ) -> SearchResult<BorrowType, K, V, marker::LeafOrInternal, marker::Leaf> |
52 | where |
53 | Q: Ord, |
54 | K: Borrow<Q>, |
55 | { |
56 | loop { |
57 | self = match self.search_node(key) { |
58 | Found(handle) => return Found(handle), |
59 | GoDown(handle) => match handle.force() { |
60 | Leaf(leaf) => return GoDown(leaf), |
61 | Internal(internal) => internal.descend(), |
62 | }, |
63 | } |
64 | } |
65 | } |
66 | |
67 | /// Descends to the nearest node where the edge matching the lower bound |
68 | /// of the range is different from the edge matching the upper bound, i.e., |
69 | /// the nearest node that has at least one key contained in the range. |
70 | /// |
71 | /// If found, returns an `Ok` with that node, the strictly ascending pair of |
72 | /// edge indices in the node delimiting the range, and the corresponding |
73 | /// pair of bounds for continuing the search in the child nodes, in case |
74 | /// the node is internal. |
75 | /// |
76 | /// If not found, returns an `Err` with the leaf edge matching the entire |
77 | /// range. |
78 | /// |
79 | /// As a diagnostic service, panics if the range specifies impossible bounds. |
80 | /// |
81 | /// The result is meaningful only if the tree is ordered by key. |
82 | pub fn search_tree_for_bifurcation<'r, Q: ?Sized, R>( |
83 | mut self, |
84 | range: &'r R, |
85 | ) -> Result< |
86 | ( |
87 | NodeRef<BorrowType, K, V, marker::LeafOrInternal>, |
88 | usize, |
89 | usize, |
90 | SearchBound<&'r Q>, |
91 | SearchBound<&'r Q>, |
92 | ), |
93 | Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>, |
94 | > |
95 | where |
96 | Q: Ord, |
97 | K: Borrow<Q>, |
98 | R: RangeBounds<Q>, |
99 | { |
100 | // Determine if map or set is being searched |
101 | let is_set = <V as super::set_val::IsSetVal>::is_set_val(); |
102 | |
103 | // Inlining these variables should be avoided. We assume the bounds reported by `range` |
104 | // remain the same, but an adversarial implementation could change between calls (#81138). |
105 | let (start, end) = (range.start_bound(), range.end_bound()); |
106 | match (start, end) { |
107 | (Bound::Excluded(s), Bound::Excluded(e)) if s == e => { |
108 | if is_set { |
109 | panic!("range start and end are equal and excluded in BTreeSet" ) |
110 | } else { |
111 | panic!("range start and end are equal and excluded in BTreeMap" ) |
112 | } |
113 | } |
114 | (Bound::Included(s) | Bound::Excluded(s), Bound::Included(e) | Bound::Excluded(e)) |
115 | if s > e => |
116 | { |
117 | if is_set { |
118 | panic!("range start is greater than range end in BTreeSet" ) |
119 | } else { |
120 | panic!("range start is greater than range end in BTreeMap" ) |
121 | } |
122 | } |
123 | _ => {} |
124 | } |
125 | let mut lower_bound = SearchBound::from_range(start); |
126 | let mut upper_bound = SearchBound::from_range(end); |
127 | loop { |
128 | let (lower_edge_idx, lower_child_bound) = self.find_lower_bound_index(lower_bound); |
129 | let (upper_edge_idx, upper_child_bound) = |
130 | unsafe { self.find_upper_bound_index(upper_bound, lower_edge_idx) }; |
131 | if lower_edge_idx < upper_edge_idx { |
132 | return Ok(( |
133 | self, |
134 | lower_edge_idx, |
135 | upper_edge_idx, |
136 | lower_child_bound, |
137 | upper_child_bound, |
138 | )); |
139 | } |
140 | debug_assert_eq!(lower_edge_idx, upper_edge_idx); |
141 | let common_edge = unsafe { Handle::new_edge(self, lower_edge_idx) }; |
142 | match common_edge.force() { |
143 | Leaf(common_edge) => return Err(common_edge), |
144 | Internal(common_edge) => { |
145 | self = common_edge.descend(); |
146 | lower_bound = lower_child_bound; |
147 | upper_bound = upper_child_bound; |
148 | } |
149 | } |
150 | } |
151 | } |
152 | |
153 | /// Finds an edge in the node delimiting the lower bound of a range. |
154 | /// Also returns the lower bound to be used for continuing the search in |
155 | /// the matching child node, if `self` is an internal node. |
156 | /// |
157 | /// The result is meaningful only if the tree is ordered by key. |
158 | pub fn find_lower_bound_edge<'r, Q>( |
159 | self, |
160 | bound: SearchBound<&'r Q>, |
161 | ) -> (Handle<Self, marker::Edge>, SearchBound<&'r Q>) |
162 | where |
163 | Q: ?Sized + Ord, |
164 | K: Borrow<Q>, |
165 | { |
166 | let (edge_idx, bound) = self.find_lower_bound_index(bound); |
167 | let edge = unsafe { Handle::new_edge(self, edge_idx) }; |
168 | (edge, bound) |
169 | } |
170 | |
171 | /// Clone of `find_lower_bound_edge` for the upper bound. |
172 | pub fn find_upper_bound_edge<'r, Q>( |
173 | self, |
174 | bound: SearchBound<&'r Q>, |
175 | ) -> (Handle<Self, marker::Edge>, SearchBound<&'r Q>) |
176 | where |
177 | Q: ?Sized + Ord, |
178 | K: Borrow<Q>, |
179 | { |
180 | let (edge_idx, bound) = unsafe { self.find_upper_bound_index(bound, 0) }; |
181 | let edge = unsafe { Handle::new_edge(self, edge_idx) }; |
182 | (edge, bound) |
183 | } |
184 | } |
185 | |
186 | impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> { |
187 | /// Looks up a given key in the node, without recursion. |
188 | /// Returns a `Found` with the handle of the matching KV, if any. Otherwise, |
189 | /// returns a `GoDown` with the handle of the edge where the key might be found |
190 | /// (if the node is internal) or where the key can be inserted. |
191 | /// |
192 | /// The result is meaningful only if the tree is ordered by key, like the tree |
193 | /// in a `BTreeMap` is. |
194 | pub fn search_node<Q: ?Sized>(self, key: &Q) -> SearchResult<BorrowType, K, V, Type, Type> |
195 | where |
196 | Q: Ord, |
197 | K: Borrow<Q>, |
198 | { |
199 | match unsafe { self.find_key_index(key, 0) } { |
200 | IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, idx) }), |
201 | IndexResult::Edge(idx) => GoDown(unsafe { Handle::new_edge(self, idx) }), |
202 | } |
203 | } |
204 | |
205 | /// Returns either the KV index in the node at which the key (or an equivalent) |
206 | /// exists, or the edge index where the key belongs, starting from a particular index. |
207 | /// |
208 | /// The result is meaningful only if the tree is ordered by key, like the tree |
209 | /// in a `BTreeMap` is. |
210 | /// |
211 | /// # Safety |
212 | /// `start_index` must be a valid edge index for the node. |
213 | unsafe fn find_key_index<Q: ?Sized>(&self, key: &Q, start_index: usize) -> IndexResult |
214 | where |
215 | Q: Ord, |
216 | K: Borrow<Q>, |
217 | { |
218 | let node = self.reborrow(); |
219 | let keys = node.keys(); |
220 | debug_assert!(start_index <= keys.len()); |
221 | for (offset, k) in unsafe { keys.get_unchecked(start_index..) }.iter().enumerate() { |
222 | match key.cmp(k.borrow()) { |
223 | Ordering::Greater => {} |
224 | Ordering::Equal => return IndexResult::KV(start_index + offset), |
225 | Ordering::Less => return IndexResult::Edge(start_index + offset), |
226 | } |
227 | } |
228 | IndexResult::Edge(keys.len()) |
229 | } |
230 | |
231 | /// Finds an edge index in the node delimiting the lower bound of a range. |
232 | /// Also returns the lower bound to be used for continuing the search in |
233 | /// the matching child node, if `self` is an internal node. |
234 | /// |
235 | /// The result is meaningful only if the tree is ordered by key. |
236 | fn find_lower_bound_index<'r, Q>( |
237 | &self, |
238 | bound: SearchBound<&'r Q>, |
239 | ) -> (usize, SearchBound<&'r Q>) |
240 | where |
241 | Q: ?Sized + Ord, |
242 | K: Borrow<Q>, |
243 | { |
244 | match bound { |
245 | Included(key) => match unsafe { self.find_key_index(key, 0) } { |
246 | IndexResult::KV(idx) => (idx, AllExcluded), |
247 | IndexResult::Edge(idx) => (idx, bound), |
248 | }, |
249 | Excluded(key) => match unsafe { self.find_key_index(key, 0) } { |
250 | IndexResult::KV(idx) => (idx + 1, AllIncluded), |
251 | IndexResult::Edge(idx) => (idx, bound), |
252 | }, |
253 | AllIncluded => (0, AllIncluded), |
254 | AllExcluded => (self.len(), AllExcluded), |
255 | } |
256 | } |
257 | |
258 | /// Mirror image of `find_lower_bound_index` for the upper bound, |
259 | /// with an additional parameter to skip part of the key array. |
260 | /// |
261 | /// # Safety |
262 | /// `start_index` must be a valid edge index for the node. |
263 | unsafe fn find_upper_bound_index<'r, Q>( |
264 | &self, |
265 | bound: SearchBound<&'r Q>, |
266 | start_index: usize, |
267 | ) -> (usize, SearchBound<&'r Q>) |
268 | where |
269 | Q: ?Sized + Ord, |
270 | K: Borrow<Q>, |
271 | { |
272 | match bound { |
273 | Included(key) => match unsafe { self.find_key_index(key, start_index) } { |
274 | IndexResult::KV(idx) => (idx + 1, AllExcluded), |
275 | IndexResult::Edge(idx) => (idx, bound), |
276 | }, |
277 | Excluded(key) => match unsafe { self.find_key_index(key, start_index) } { |
278 | IndexResult::KV(idx) => (idx, AllIncluded), |
279 | IndexResult::Edge(idx) => (idx, bound), |
280 | }, |
281 | AllIncluded => (self.len(), AllIncluded), |
282 | AllExcluded => (start_index, AllExcluded), |
283 | } |
284 | } |
285 | } |
286 | |