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