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

   1 use core::borrow::Borrow;
   2 use core::cmp::Ordering;
   3
   4 use super::node::{marker, ForceResult::*, Handle, NodeRef};
   5
   6 use SearchResult::*;
   7
   8 pub enum SearchResult<BorrowType, K, V, FoundType, GoDownType> {
   9     Found(Handle<NodeRef<BorrowType, K, V, FoundType>, marker::KV>),
  10     GoDown(Handle<NodeRef<BorrowType, K, V, GoDownType>, marker::Edge>),
  11 }
  12
  13 pub enum IndexResult {
  14     KV(usize),
  15     Edge(usize),
  16 }
  17
  18 impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
  19     /// Looks up a given key in a (sub)tree headed by the node, recursively.
  20     /// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
  21     /// returns a `GoDown` with the handle of the leaf edge where the key belongs.
  22     ///
  23     /// The result is meaningful only if the tree is ordered by key, like the tree
  24     /// in a `BTreeMap` is.
  25     pub fn search_tree<Q: ?Sized>(
  26         mut self,
  27         key: &Q,
  28     ) -> SearchResult<BorrowType, K, V, marker::LeafOrInternal, marker::Leaf>
  29     where
  30         Q: Ord,
  31         K: Borrow<Q>,
  32     {
  33         loop {
  34             self = match self.search_node(key) {
  35                 Found(handle) => return Found(handle),
  36                 GoDown(handle) => match handle.force() {
  37                     Leaf(leaf) => return GoDown(leaf),
  38                     Internal(internal) => internal.descend(),
  39                 },
  40             }
  41         }
  42     }
  43 }
  44
  45 impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
  46     /// Looks up a given key in the node, without recursion.
  47     /// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
  48     /// returns a `GoDown` with the handle of the edge where the key might be found
  49     /// (if the node is internal) or where the key can be inserted.
  50     ///
  51     /// The result is meaningful only if the tree is ordered by key, like the tree
  52     /// in a `BTreeMap` is.
  53     pub fn search_node<Q: ?Sized>(self, key: &Q) -> SearchResult<BorrowType, K, V, Type, Type>
  54     where
  55         Q: Ord,
  56         K: Borrow<Q>,
  57     {
  58         match self.find_index(key) {
  59             IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, idx) }),
  60             IndexResult::Edge(idx) => GoDown(unsafe { Handle::new_edge(self, idx) }),
  61         }
  62     }
  63
  64     /// Returns either the KV index in the node at which the key (or an equivalent)
  65     /// exists, or the edge index where the key belongs.
  66     ///
  67     /// The result is meaningful only if the tree is ordered by key, like the tree
  68     /// in a `BTreeMap` is.
  69     fn find_index<Q: ?Sized>(&self, key: &Q) -> IndexResult
  70     where
  71         Q: Ord,
  72         K: Borrow<Q>,
  73     {
  74         let node = self.reborrow();
  75         let keys = node.keys();
  76         for (i, k) in keys.iter().enumerate() {
  77             match key.cmp(k.borrow()) {
  78                 Ordering::Greater => {}
  79                 Ordering::Equal => return IndexResult::KV(i),
  80                 Ordering::Less => return IndexResult::Edge(i),
  81             }
  82         }
  83         IndexResult::Edge(keys.len())
  84     }
  85 }