oasis_core_runtime/storage/mkvs/tree/
lookup.rs

1use anyhow::Result;
2
3use crate::storage::mkvs::{
4    cache::{Cache, ReadSyncFetcher},
5    sync::{GetRequest, Proof, ProofBuilder, ReadSync, TreeID},
6    tree::{Depth, Key, KeyTrait, NodeBox, NodeKind, NodePtrRef, Root, Tree, Value},
7};
8
9pub(super) struct FetcherSyncGet<'a> {
10    key: &'a Key,
11    include_siblings: bool,
12}
13
14impl<'a> FetcherSyncGet<'a> {
15    pub(super) fn new(key: &'a Key, include_siblings: bool) -> Self {
16        Self {
17            key,
18            include_siblings,
19        }
20    }
21}
22
23impl<'a> ReadSyncFetcher for FetcherSyncGet<'a> {
24    fn fetch(&self, root: Root, ptr: NodePtrRef, rs: &mut Box<dyn ReadSync>) -> Result<Proof> {
25        let rsp = rs.sync_get(GetRequest {
26            tree: TreeID {
27                root,
28                position: ptr.borrow().hash,
29            },
30            key: self.key.clone(),
31            include_siblings: self.include_siblings,
32        })?;
33        Ok(rsp.proof)
34    }
35}
36
37impl Tree {
38    /// Get an existing key.
39    pub fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
40        self._get_top(key, false)
41    }
42
43    pub fn get_proof(&self, key: &[u8]) -> Result<Option<Proof>> {
44        let boxed_key = key.to_vec();
45        let pending_root = self.cache.borrow().get_pending_root();
46
47        // Remember where the path from root to target node ends (will end).
48        self.cache.borrow_mut().mark_position();
49
50        let mut proof_builder = ProofBuilder::new(pending_root.as_ref().borrow().hash);
51
52        let result = self._get(pending_root, 0, &boxed_key, false, Some(&mut proof_builder))?;
53        match result {
54            Some(_) => Ok(Some(proof_builder.build())),
55            None => Ok(None),
56        }
57    }
58
59    /// Check if the key exists in the local cache.
60    pub fn cache_contains_key(&self, key: &[u8]) -> bool {
61        match self._get_top(key, true) {
62            Ok(Some(_)) => true,
63            Ok(None) => false,
64            Err(_) => false,
65        }
66    }
67
68    fn _get_top(&self, key: &[u8], check_only: bool) -> Result<Option<Vec<u8>>> {
69        let boxed_key = key.to_vec();
70        let pending_root = self.cache.borrow().get_pending_root();
71
72        // Remember where the path from root to target node ends (will end).
73        self.cache.borrow_mut().mark_position();
74
75        self._get(pending_root, 0, &boxed_key, check_only, None)
76    }
77
78    fn _get(
79        &self,
80        ptr: NodePtrRef,
81        bit_depth: Depth,
82        key: &Key,
83        check_only: bool,
84        mut proof_builder: Option<&mut ProofBuilder>,
85    ) -> Result<Option<Value>> {
86        let node_ref = self.cache.borrow_mut().deref_node_ptr(
87            ptr,
88            if check_only {
89                None
90            } else {
91                Some(FetcherSyncGet::new(key, false))
92            },
93        )?;
94
95        // Include nodes in proof if we have a proof builder.
96        if let (Some(pb), Some(node_ref)) = (proof_builder.as_mut(), &node_ref) {
97            pb.include(&node_ref.borrow());
98        }
99
100        match classify_noderef!(?node_ref) {
101            NodeKind::None => {
102                // Reached a nil node, there is nothing here.
103                Ok(None)
104            }
105            NodeKind::Internal => {
106                let node_ref = node_ref.unwrap();
107                if let NodeBox::Internal(ref mut n) = *node_ref.borrow_mut() {
108                    // Internal node.
109                    // Does lookup key end here? Look into LeafNode.
110                    if key.bit_length() == bit_depth + n.label_bit_length {
111                        return self._get(
112                            n.leaf_node.clone(),
113                            bit_depth + n.label_bit_length,
114                            key,
115                            check_only,
116                            proof_builder,
117                        );
118                    }
119
120                    // Lookup key is too short for the current n.Label. It's not stored.
121                    if key.bit_length() < bit_depth + n.label_bit_length {
122                        return Ok(None);
123                    }
124
125                    // Continue recursively based on a bit value.
126                    if key.get_bit(bit_depth + n.label_bit_length) {
127                        return self._get(
128                            n.right.clone(),
129                            bit_depth + n.label_bit_length,
130                            key,
131                            check_only,
132                            proof_builder,
133                        );
134                    } else {
135                        return self._get(
136                            n.left.clone(),
137                            bit_depth + n.label_bit_length,
138                            key,
139                            check_only,
140                            proof_builder,
141                        );
142                    }
143                }
144
145                unreachable!("node kind is internal node");
146            }
147            NodeKind::Leaf => {
148                // Reached a leaf node, check if key matches.
149                let node_ref = node_ref.unwrap();
150                if noderef_as!(node_ref, Leaf).key == *key {
151                    Ok(Some(noderef_as!(node_ref, Leaf).value.clone()))
152                } else {
153                    Ok(None)
154                }
155            }
156        }
157    }
158}