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