]>
Commit | Line | Data |
---|---|---|
9fa01778 | 1 | use core::borrow::Borrow; |
9cc50fc6 | 2 | use core::cmp::Ordering; |
6a06907d | 3 | use core::ops::{Bound, RangeBounds}; |
9cc50fc6 | 4 | |
60c5eb7d | 5 | use super::node::{marker, ForceResult::*, Handle, NodeRef}; |
9cc50fc6 | 6 | |
6a06907d | 7 | use SearchBound::*; |
9fa01778 | 8 | use SearchResult::*; |
9cc50fc6 | 9 | |
6a06907d XL |
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) => Included(t), | |
25 | Bound::Excluded(t) => Excluded(t), | |
26 | Bound::Unbounded => AllIncluded, | |
27 | } | |
28 | } | |
29 | } | |
30 | ||
9cc50fc6 SL |
31 | pub enum SearchResult<BorrowType, K, V, FoundType, GoDownType> { |
32 | Found(Handle<NodeRef<BorrowType, K, V, FoundType>, marker::KV>), | |
60c5eb7d | 33 | GoDown(Handle<NodeRef<BorrowType, K, V, GoDownType>, marker::Edge>), |
9cc50fc6 SL |
34 | } |
35 | ||
5869c6ff XL |
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 | } | |
9cc50fc6 SL |
64 | } |
65 | } | |
6a06907d XL |
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 | /// | |
cdc7bbd5 XL |
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. | |
6a06907d XL |
75 | /// |
76 | /// If not found, returns an `Err` with the leaf edge matching the entire | |
77 | /// range. | |
78 | /// | |
cdc7bbd5 XL |
79 | /// As a diagnostic service, panics if the range specifies impossible bounds. |
80 | /// | |
6a06907d XL |
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 | { | |
923072b8 FG |
100 | // Determine if map or set is being searched |
101 | let is_set = <V as super::set_val::IsSetVal>::is_set_val(); | |
102 | ||
6a06907d XL |
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 => { | |
923072b8 FG |
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 | } | |
6a06907d XL |
113 | } |
114 | (Bound::Included(s) | Bound::Excluded(s), Bound::Included(e) | Bound::Excluded(e)) | |
115 | if s > e => | |
116 | { | |
923072b8 FG |
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 | } | |
6a06907d XL |
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); | |
cdc7bbd5 XL |
129 | let (upper_edge_idx, upper_child_bound) = |
130 | unsafe { self.find_upper_bound_index(upper_bound, lower_edge_idx) }; | |
6a06907d XL |
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 | } | |
cdc7bbd5 | 140 | debug_assert_eq!(lower_edge_idx, upper_edge_idx); |
6a06907d XL |
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 | { | |
cdc7bbd5 | 180 | let (edge_idx, bound) = unsafe { self.find_upper_bound_index(bound, 0) }; |
6a06907d XL |
181 | let edge = unsafe { Handle::new_edge(self, edge_idx) }; |
182 | (edge, bound) | |
183 | } | |
9cc50fc6 SL |
184 | } |
185 | ||
5869c6ff XL |
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 | { | |
cdc7bbd5 | 199 | match unsafe { self.find_key_index(key, 0) } { |
5869c6ff XL |
200 | IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, idx) }), |
201 | IndexResult::Edge(idx) => GoDown(unsafe { Handle::new_edge(self, idx) }), | |
202 | } | |
9cc50fc6 | 203 | } |
9cc50fc6 | 204 | |
5869c6ff | 205 | /// Returns either the KV index in the node at which the key (or an equivalent) |
cdc7bbd5 | 206 | /// exists, or the edge index where the key belongs, starting from a particular index. |
5869c6ff XL |
207 | /// |
208 | /// The result is meaningful only if the tree is ordered by key, like the tree | |
209 | /// in a `BTreeMap` is. | |
cdc7bbd5 XL |
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 | |
5869c6ff XL |
214 | where |
215 | Q: Ord, | |
216 | K: Borrow<Q>, | |
217 | { | |
218 | let node = self.reborrow(); | |
219 | let keys = node.keys(); | |
cdc7bbd5 XL |
220 | debug_assert!(start_index <= keys.len()); |
221 | for (offset, k) in unsafe { keys.get_unchecked(start_index..) }.iter().enumerate() { | |
5869c6ff XL |
222 | match key.cmp(k.borrow()) { |
223 | Ordering::Greater => {} | |
cdc7bbd5 XL |
224 | Ordering::Equal => return IndexResult::KV(start_index + offset), |
225 | Ordering::Less => return IndexResult::Edge(start_index + offset), | |
5869c6ff | 226 | } |
9cc50fc6 | 227 | } |
5869c6ff | 228 | IndexResult::Edge(keys.len()) |
9cc50fc6 | 229 | } |
6a06907d XL |
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 { | |
cdc7bbd5 | 245 | Included(key) => match unsafe { self.find_key_index(key, 0) } { |
6a06907d XL |
246 | IndexResult::KV(idx) => (idx, AllExcluded), |
247 | IndexResult::Edge(idx) => (idx, bound), | |
248 | }, | |
cdc7bbd5 | 249 | Excluded(key) => match unsafe { self.find_key_index(key, 0) } { |
6a06907d XL |
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 | ||
cdc7bbd5 XL |
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>( | |
6a06907d XL |
264 | &self, |
265 | bound: SearchBound<&'r Q>, | |
cdc7bbd5 | 266 | start_index: usize, |
6a06907d XL |
267 | ) -> (usize, SearchBound<&'r Q>) |
268 | where | |
269 | Q: ?Sized + Ord, | |
270 | K: Borrow<Q>, | |
271 | { | |
272 | match bound { | |
cdc7bbd5 | 273 | Included(key) => match unsafe { self.find_key_index(key, start_index) } { |
6a06907d XL |
274 | IndexResult::KV(idx) => (idx + 1, AllExcluded), |
275 | IndexResult::Edge(idx) => (idx, bound), | |
276 | }, | |
cdc7bbd5 | 277 | Excluded(key) => match unsafe { self.find_key_index(key, start_index) } { |
6a06907d XL |
278 | IndexResult::KV(idx) => (idx, AllIncluded), |
279 | IndexResult::Edge(idx) => (idx, bound), | |
280 | }, | |
281 | AllIncluded => (self.len(), AllIncluded), | |
cdc7bbd5 | 282 | AllExcluded => (start_index, AllExcluded), |
6a06907d XL |
283 | } |
284 | } | |
9cc50fc6 | 285 | } |