library/alloc/src/collections/btree/search.rs

   1 use core::borrow::Borrow;
   2 use core::cmp::Ordering;
   3 use core::ops::{Bound, RangeBounds};
   4
   5 use super::node::{marker, ForceResult::*, Handle, NodeRef};
   6
   7 use SearchBound::*;
   8 use SearchResult::*;
   9
  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
  31 pub 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
  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             }
  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         // Inlining these variables should be avoided. We assume the bounds reported by `range`
 101         // remain the same, but an adversarial implementation could change between calls (#81138).
 102         let (start, end) = (range.start_bound(), range.end_bound());
 103         match (start, end) {
 104             (Bound::Excluded(s), Bound::Excluded(e)) if s == e => {
 105                 panic!("range start and end are equal and excluded in BTreeMap")
 106             }
 107             (Bound::Included(s) | Bound::Excluded(s), Bound::Included(e) | Bound::Excluded(e))
 108                 if s > e =>
 109             {
 110                 panic!("range start is greater than range end in BTreeMap")
 111             }
 112             _ => {}
 113         }
 114         let mut lower_bound = SearchBound::from_range(start);
 115         let mut upper_bound = SearchBound::from_range(end);
 116         loop {
 117             let (lower_edge_idx, lower_child_bound) = self.find_lower_bound_index(lower_bound);
 118             let (upper_edge_idx, upper_child_bound) =
 119                 unsafe { self.find_upper_bound_index(upper_bound, lower_edge_idx) };
 120             if lower_edge_idx < upper_edge_idx {
 121                 return Ok((
 122                     self,
 123                     lower_edge_idx,
 124                     upper_edge_idx,
 125                     lower_child_bound,
 126                     upper_child_bound,
 127                 ));
 128             }
 129             debug_assert_eq!(lower_edge_idx, upper_edge_idx);
 130             let common_edge = unsafe { Handle::new_edge(self, lower_edge_idx) };
 131             match common_edge.force() {
 132                 Leaf(common_edge) => return Err(common_edge),
 133                 Internal(common_edge) => {
 134                     self = common_edge.descend();
 135                     lower_bound = lower_child_bound;
 136                     upper_bound = upper_child_bound;
 137                 }
 138             }
 139         }
 140     }
 141
 142     /// Finds an edge in the node delimiting the lower bound of a range.
 143     /// Also returns the lower bound to be used for continuing the search in
 144     /// the matching child node, if `self` is an internal node.
 145     ///
 146     /// The result is meaningful only if the tree is ordered by key.
 147     pub fn find_lower_bound_edge<'r, Q>(
 148         self,
 149         bound: SearchBound<&'r Q>,
 150     ) -> (Handle<Self, marker::Edge>, SearchBound<&'r Q>)
 151     where
 152         Q: ?Sized + Ord,
 153         K: Borrow<Q>,
 154     {
 155         let (edge_idx, bound) = self.find_lower_bound_index(bound);
 156         let edge = unsafe { Handle::new_edge(self, edge_idx) };
 157         (edge, bound)
 158     }
 159
 160     /// Clone of `find_lower_bound_edge` for the upper bound.
 161     pub fn find_upper_bound_edge<'r, Q>(
 162         self,
 163         bound: SearchBound<&'r Q>,
 164     ) -> (Handle<Self, marker::Edge>, SearchBound<&'r Q>)
 165     where
 166         Q: ?Sized + Ord,
 167         K: Borrow<Q>,
 168     {
 169         let (edge_idx, bound) = unsafe { self.find_upper_bound_index(bound, 0) };
 170         let edge = unsafe { Handle::new_edge(self, edge_idx) };
 171         (edge, bound)
 172     }
 173 }
 174
 175 impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
 176     /// Looks up a given key in the node, without recursion.
 177     /// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
 178     /// returns a `GoDown` with the handle of the edge where the key might be found
 179     /// (if the node is internal) or where the key can be inserted.
 180     ///
 181     /// The result is meaningful only if the tree is ordered by key, like the tree
 182     /// in a `BTreeMap` is.
 183     pub fn search_node<Q: ?Sized>(self, key: &Q) -> SearchResult<BorrowType, K, V, Type, Type>
 184     where
 185         Q: Ord,
 186         K: Borrow<Q>,
 187     {
 188         match unsafe { self.find_key_index(key, 0) } {
 189             IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, idx) }),
 190             IndexResult::Edge(idx) => GoDown(unsafe { Handle::new_edge(self, idx) }),
 191         }
 192     }
 193
 194     /// Returns either the KV index in the node at which the key (or an equivalent)
 195     /// exists, or the edge index where the key belongs, starting from a particular index.
 196     ///
 197     /// The result is meaningful only if the tree is ordered by key, like the tree
 198     /// in a `BTreeMap` is.
 199     ///
 200     /// # Safety
 201     /// `start_index` must be a valid edge index for the node.
 202     unsafe fn find_key_index<Q: ?Sized>(&self, key: &Q, start_index: usize) -> IndexResult
 203     where
 204         Q: Ord,
 205         K: Borrow<Q>,
 206     {
 207         let node = self.reborrow();
 208         let keys = node.keys();
 209         debug_assert!(start_index <= keys.len());
 210         for (offset, k) in unsafe { keys.get_unchecked(start_index..) }.iter().enumerate() {
 211             match key.cmp(k.borrow()) {
 212                 Ordering::Greater => {}
 213                 Ordering::Equal => return IndexResult::KV(start_index + offset),
 214                 Ordering::Less => return IndexResult::Edge(start_index + offset),
 215             }
 216         }
 217         IndexResult::Edge(keys.len())
 218     }
 219
 220     /// Finds an edge index in the node delimiting the lower bound of a range.
 221     /// Also returns the lower bound to be used for continuing the search in
 222     /// the matching child node, if `self` is an internal node.
 223     ///
 224     /// The result is meaningful only if the tree is ordered by key.
 225     fn find_lower_bound_index<'r, Q>(
 226         &self,
 227         bound: SearchBound<&'r Q>,
 228     ) -> (usize, SearchBound<&'r Q>)
 229     where
 230         Q: ?Sized + Ord,
 231         K: Borrow<Q>,
 232     {
 233         match bound {
 234             Included(key) => match unsafe { self.find_key_index(key, 0) } {
 235                 IndexResult::KV(idx) => (idx, AllExcluded),
 236                 IndexResult::Edge(idx) => (idx, bound),
 237             },
 238             Excluded(key) => match unsafe { self.find_key_index(key, 0) } {
 239                 IndexResult::KV(idx) => (idx + 1, AllIncluded),
 240                 IndexResult::Edge(idx) => (idx, bound),
 241             },
 242             AllIncluded => (0, AllIncluded),
 243             AllExcluded => (self.len(), AllExcluded),
 244         }
 245     }
 246
 247     /// Mirror image of `find_lower_bound_index` for the upper bound,
 248     /// with an additional parameter to skip part of the key array.
 249     ///
 250     /// # Safety
 251     /// `start_index` must be a valid edge index for the node.
 252     unsafe fn find_upper_bound_index<'r, Q>(
 253         &self,
 254         bound: SearchBound<&'r Q>,
 255         start_index: usize,
 256     ) -> (usize, SearchBound<&'r Q>)
 257     where
 258         Q: ?Sized + Ord,
 259         K: Borrow<Q>,
 260     {
 261         match bound {
 262             Included(key) => match unsafe { self.find_key_index(key, start_index) } {
 263                 IndexResult::KV(idx) => (idx + 1, AllExcluded),
 264                 IndexResult::Edge(idx) => (idx, bound),
 265             },
 266             Excluded(key) => match unsafe { self.find_key_index(key, start_index) } {
 267                 IndexResult::KV(idx) => (idx, AllIncluded),
 268                 IndexResult::Edge(idx) => (idx, bound),
 269             },
 270             AllIncluded => (self.len(), AllIncluded),
 271             AllExcluded => (start_index, AllExcluded),
 272         }
 273     }
 274 }