     1
     2
     3
     4
     5
     6
     7
     8
     9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85

    use core::borrow::Borrow;
    use core::cmp::Ordering;

    use super::node::{marker, ForceResult::*, Handle, NodeRef};

    use SearchResult::*;

    pub enum SearchResult<BorrowType, K, V, FoundType, GoDownType> {
        Found(Handle<NodeRef<BorrowType, K, V, FoundType>, marker::KV>),
        GoDown(Handle<NodeRef<BorrowType, K, V, GoDownType>, marker::Edge>),
    }

    pub enum IndexResult {
        KV(usize),
        Edge(usize),
    }

    impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
        /// Looks up a given key in a (sub)tree headed by the node, recursively.
        /// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
        /// returns a `GoDown` with the handle of the leaf edge where the key belongs.
        ///
        /// The result is meaningful only if the tree is ordered by key, like the tree
        /// in a `BTreeMap` is.
        pub fn search_tree<Q: ?Sized>(
            mut self,
            key: &Q,
        ) -> SearchResult<BorrowType, K, V, marker::LeafOrInternal, marker::Leaf>
        where
            Q: Ord,
            K: Borrow<Q>,
        {
            loop {
                self = match self.search_node(key) {
                    Found(handle) => return Found(handle),
                    GoDown(handle) => match handle.force() {
                        Leaf(leaf) => return GoDown(leaf),
                        Internal(internal) => internal.descend(),
                    },
                }
            }
        }
    }

    impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
        /// Looks up a given key in the node, without recursion.
        /// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
        /// returns a `GoDown` with the handle of the edge where the key might be found
        /// (if the node is internal) or where the key can be inserted.
        ///
        /// The result is meaningful only if the tree is ordered by key, like the tree
        /// in a `BTreeMap` is.
        pub fn search_node<Q: ?Sized>(self, key: &Q) -> SearchResult<BorrowType, K, V, Type, Type>
        where
            Q: Ord,
            K: Borrow<Q>,
        {
            match self.find_index(key) {
                IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, idx) }),
                IndexResult::Edge(idx) => GoDown(unsafe { Handle::new_edge(self, idx) }),
            }
        }

        /// Returns either the KV index in the node at which the key (or an equivalent)
        /// exists, or the edge index where the key belongs.
        ///
        /// The result is meaningful only if the tree is ordered by key, like the tree
        /// in a `BTreeMap` is.
        fn find_index<Q: ?Sized>(&self, key: &Q) -> IndexResult
        where
            Q: Ord,
            K: Borrow<Q>,
        {
            let node = self.reborrow();
            let keys = node.keys();
            for (i, k) in keys.iter().enumerate() {
                match key.cmp(k.borrow()) {
                    Ordering::Greater => {}
                    Ordering::Equal => return IndexResult::KV(i),
                    Ordering::Less => return IndexResult::Edge(i),
                }
            }
            IndexResult::Edge(keys.len())
        }
    }