1use core::borrow::Borrow;
2use core::cmp::Ordering;
3use core::ops::{Bound, RangeBounds};
4
5use super::node::{marker, ForceResult::*, Handle, NodeRef};
6
7use SearchBound::*;
8use SearchResult::*;
9
10pub 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
21impl<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
31pub 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
36pub enum IndexResult {
37 KV(usize),
38 Edge(usize),
39}
40
41impl<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
186impl<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