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 pair of edge indices in it
  72     /// delimiting the range, and the corresponding pair of bounds for
  73     /// continuing the search in the child nodes, in case the node is internal.
  74     ///
  75     /// If not found, returns an `Err` with the leaf edge matching the entire
  76     /// range.
  77     ///
  78     /// The result is meaningful only if the tree is ordered by key.
  79     pub fn search_tree_for_bifurcation<'r, Q: ?Sized, R>(
  80         mut self,
  81         range: &'r R,
  82     ) -> Result<
  83         (
  84             NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
  85             usize,
  86             usize,
  87             SearchBound<&'r Q>,
  88             SearchBound<&'r Q>,
  89         ),
  90         Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
  91     >
  92     where
  93         Q: Ord,
  94         K: Borrow<Q>,
  95         R: RangeBounds<Q>,
  96     {
  97         // WARNING: Inlining these variables would be unsound (#81138)
  98         // We assume the bounds reported by `range` remain the same, but
  99         // an adversarial implementation could change between calls
 100         let (start, end) = (range.start_bound(), range.end_bound());
 101         match (start, end) {
 102             (Bound::Excluded(s), Bound::Excluded(e)) if s == e => {
 103                 panic!("range start and end are equal and excluded in BTreeMap")
 104             }
 105             (Bound::Included(s) | Bound::Excluded(s), Bound::Included(e) | Bound::Excluded(e))
 106                 if s > e =>
 107             {
 108                 panic!("range start is greater than range end in BTreeMap")
 109             }
 110             _ => {}
 111         }
 112         let mut lower_bound = SearchBound::from_range(start);
 113         let mut upper_bound = SearchBound::from_range(end);
 114         loop {
 115             let (lower_edge_idx, lower_child_bound) = self.find_lower_bound_index(lower_bound);
 116             let (upper_edge_idx, upper_child_bound) = self.find_upper_bound_index(upper_bound);
 117             if lower_edge_idx > upper_edge_idx {
 118                 panic!("Ord is ill-defined in BTreeMap range")
 119             }
 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             let common_edge = unsafe { Handle::new_edge(self, lower_edge_idx) };
 130             match common_edge.force() {
 131                 Leaf(common_edge) => return Err(common_edge),
 132                 Internal(common_edge) => {
 133                     self = common_edge.descend();
 134                     lower_bound = lower_child_bound;
 135                     upper_bound = upper_child_bound;
 136                 }
 137             }
 138         }
 139     }
 140
 141     /// Finds an edge in the node delimiting the lower bound of a range.
 142     /// Also returns the lower bound to be used for continuing the search in
 143     /// the matching child node, if `self` is an internal node.
 144     ///
 145     /// The result is meaningful only if the tree is ordered by key.
 146     pub fn find_lower_bound_edge<'r, Q>(
 147         self,
 148         bound: SearchBound<&'r Q>,
 149     ) -> (Handle<Self, marker::Edge>, SearchBound<&'r Q>)
 150     where
 151         Q: ?Sized + Ord,
 152         K: Borrow<Q>,
 153     {
 154         let (edge_idx, bound) = self.find_lower_bound_index(bound);
 155         let edge = unsafe { Handle::new_edge(self, edge_idx) };
 156         (edge, bound)
 157     }
 158
 159     /// Clone of `find_lower_bound_edge` for the upper bound.
 160     pub fn find_upper_bound_edge<'r, Q>(
 161         self,
 162         bound: SearchBound<&'r Q>,
 163     ) -> (Handle<Self, marker::Edge>, SearchBound<&'r Q>)
 164     where
 165         Q: ?Sized + Ord,
 166         K: Borrow<Q>,
 167     {
 168         let (edge_idx, bound) = self.find_upper_bound_index(bound);
 169         let edge = unsafe { Handle::new_edge(self, edge_idx) };
 170         (edge, bound)
 171     }
 172 }
 173
 174 impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
 175     /// Looks up a given key in the node, without recursion.
 176     /// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
 177     /// returns a `GoDown` with the handle of the edge where the key might be found
 178     /// (if the node is internal) or where the key can be inserted.
 179     ///
 180     /// The result is meaningful only if the tree is ordered by key, like the tree
 181     /// in a `BTreeMap` is.
 182     pub fn search_node<Q: ?Sized>(self, key: &Q) -> SearchResult<BorrowType, K, V, Type, Type>
 183     where
 184         Q: Ord,
 185         K: Borrow<Q>,
 186     {
 187         match self.find_key_index(key) {
 188             IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, idx) }),
 189             IndexResult::Edge(idx) => GoDown(unsafe { Handle::new_edge(self, idx) }),
 190         }
 191     }
 192
 193     /// Returns either the KV index in the node at which the key (or an equivalent)
 194     /// exists, or the edge index where the key belongs.
 195     ///
 196     /// The result is meaningful only if the tree is ordered by key, like the tree
 197     /// in a `BTreeMap` is.
 198     fn find_key_index<Q: ?Sized>(&self, key: &Q) -> IndexResult
 199     where
 200         Q: Ord,
 201         K: Borrow<Q>,
 202     {
 203         let node = self.reborrow();
 204         let keys = node.keys();
 205         for (i, k) in keys.iter().enumerate() {
 206             match key.cmp(k.borrow()) {
 207                 Ordering::Greater => {}
 208                 Ordering::Equal => return IndexResult::KV(i),
 209                 Ordering::Less => return IndexResult::Edge(i),
 210             }
 211         }
 212         IndexResult::Edge(keys.len())
 213     }
 214
 215     /// Finds an edge index in the node delimiting the lower bound of a range.
 216     /// Also returns the lower bound to be used for continuing the search in
 217     /// the matching child node, if `self` is an internal node.
 218     ///
 219     /// The result is meaningful only if the tree is ordered by key.
 220     fn find_lower_bound_index<'r, Q>(
 221         &self,
 222         bound: SearchBound<&'r Q>,
 223     ) -> (usize, SearchBound<&'r Q>)
 224     where
 225         Q: ?Sized + Ord,
 226         K: Borrow<Q>,
 227     {
 228         match bound {
 229             Included(key) => match self.find_key_index(key) {
 230                 IndexResult::KV(idx) => (idx, AllExcluded),
 231                 IndexResult::Edge(idx) => (idx, bound),
 232             },
 233             Excluded(key) => match self.find_key_index(key) {
 234                 IndexResult::KV(idx) => (idx + 1, AllIncluded),
 235                 IndexResult::Edge(idx) => (idx, bound),
 236             },
 237             AllIncluded => (0, AllIncluded),
 238             AllExcluded => (self.len(), AllExcluded),
 239         }
 240     }
 241
 242     /// Clone of `find_lower_bound_index` for the upper bound.
 243     fn find_upper_bound_index<'r, Q>(
 244         &self,
 245         bound: SearchBound<&'r Q>,
 246     ) -> (usize, SearchBound<&'r Q>)
 247     where
 248         Q: ?Sized + Ord,
 249         K: Borrow<Q>,
 250     {
 251         match bound {
 252             Included(key) => match self.find_key_index(key) {
 253                 IndexResult::KV(idx) => (idx + 1, AllExcluded),
 254                 IndexResult::Edge(idx) => (idx, bound),
 255             },
 256             Excluded(key) => match self.find_key_index(key) {
 257                 IndexResult::KV(idx) => (idx, AllIncluded),
 258                 IndexResult::Edge(idx) => (idx, bound),
 259             },
 260             AllIncluded => (self.len(), AllIncluded),
 261             AllExcluded => (0, AllExcluded),
 262         }
 263     }
 264 }