alloc/collections/btree/
search.rs

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) => Included(t),
26            Bound::Excluded(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}