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 /// Looks up a given key in a (sub)tree headed by the given node, recursively.
  14 /// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
  15 /// returns a `GoDown` with the handle of the possible leaf edge where the key
  16 /// belongs.
  17 pub fn search_tree<BorrowType, K, V, Q: ?Sized>(
  18     mut node: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
  19     key: &Q,
  20 ) -> SearchResult<BorrowType, K, V, marker::LeafOrInternal, marker::Leaf>
  21 where
  22     Q: Ord,
  23     K: Borrow<Q>,
  24 {
  25     loop {
  26         match search_node(node, key) {
  27             Found(handle) => return Found(handle),
  28             GoDown(handle) => match handle.force() {
  29                 Leaf(leaf) => return GoDown(leaf),
  30                 Internal(internal) => {
  31                     node = internal.descend();
  32                     continue;
  33                 }
  34             },
  35         }
  36     }
  37 }
  38
  39 /// Looks up a given key in a given node, without recursion.
  40 /// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
  41 /// returns a `GoDown` with the handle of the edge where the key might be found.
  42 /// If the node is a leaf, a `GoDown` edge is not an actual edge but a possible edge.
  43 pub fn search_node<BorrowType, K, V, Type, Q: ?Sized>(
  44     node: NodeRef<BorrowType, K, V, Type>,
  45     key: &Q,
  46 ) -> SearchResult<BorrowType, K, V, Type, Type>
  47 where
  48     Q: Ord,
  49     K: Borrow<Q>,
  50 {
  51     match search_linear(&node, key) {
  52         (idx, true) => Found(unsafe { Handle::new_kv(node, idx) }),
  53         (idx, false) => SearchResult::GoDown(unsafe { Handle::new_edge(node, idx) }),
  54     }
  55 }
  56
  57 /// Returns the index in the node at which the key (or an equivalent) exists
  58 /// or could exist, and whether it exists in the node itself. If it doesn't
  59 /// exist in the node itself, it may exist in the subtree with that index
  60 /// (if the node has subtrees). If the key doesn't exist in node or subtree,
  61 /// the returned index is the position or subtree where the key belongs.
  62 fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
  63     node: &NodeRef<BorrowType, K, V, Type>,
  64     key: &Q,
  65 ) -> (usize, bool)
  66 where
  67     Q: Ord,
  68     K: Borrow<Q>,
  69 {
  70     // This function is defined over all borrow types (immutable, mutable, owned).
  71     // Using `keys_at()` is fine here even if BorrowType is mutable, as all we return
  72     // is an index -- not a reference.
  73     let len = node.len();
  74     for i in 0..len {
  75         let k = unsafe { node.reborrow().key_at(i) };
  76         match key.cmp(k.borrow()) {
  77             Ordering::Greater => {}
  78             Ordering::Equal => return (i, true),
  79             Ordering::Less => return (i, false),
  80         }
  81     }
  82     (len, false)
  83 }