merkle_tree_db/
treedb.rs

1use hash_db::{HashDBRef, EMPTY_PREFIX};
2
3use super::{
4    null_nodes, ChildSelector, DBValue, DataError, HashMap, Hasher, Key, KeyedTree, Node, NodeHash,
5    TreeError, TreeRecorder,
6};
7
8// TreeDBBuilder
9// ================================================================================================
10
11/// Used to construct a TreeDB
12pub struct TreeDBBuilder<'db, const D: usize, H: Hasher> {
13    db: &'db dyn HashDBRef<H, DBValue>,
14    root: &'db H::Out,
15    recorder: Option<&'db mut dyn TreeRecorder<H>>,
16}
17
18impl<'db, const D: usize, H: Hasher> TreeDBBuilder<'db, D, H> {
19    /// Construct a new TreeDBBuilder
20    pub fn new(db: &'db dyn HashDBRef<H, DBValue>, root: &'db H::Out) -> Result<Self, TreeError> {
21        //TODO: warm user if default root provided
22        if D > usize::MAX / 8 {
23            return Err(TreeError::DepthTooLarge(D, usize::MAX / 8));
24        }
25        Ok(Self {
26            db,
27            root,
28            recorder: None,
29        })
30    }
31
32    /// Add a recorder to the TreeDBBuilder
33    pub fn with_recorder(mut self, recorder: &'db mut dyn TreeRecorder<H>) -> Self {
34        self.recorder = Some(recorder);
35        self
36    }
37
38    /// Add an optional recorder to the TreeDBBuilder
39    pub fn with_optional_recorder<'recorder: 'db>(
40        mut self,
41        recorder: Option<&'recorder mut dyn TreeRecorder<H>>,
42    ) -> Self {
43        self.recorder = recorder.map(|r| r as _);
44        self
45    }
46
47    /// build a TreeDB
48    pub fn build(self) -> TreeDB<'db, D, H> {
49        let (null_nodes, default_root) = null_nodes::<H>(D * 8);
50        let root = if self.root == &H::Out::default() || self.root == &default_root {
51            NodeHash::Default(default_root)
52        } else {
53            NodeHash::Database(*self.root)
54        };
55        TreeDB {
56            db: self.db,
57            root,
58            recorder: self.recorder.map(core::cell::RefCell::new),
59            null_nodes,
60        }
61    }
62}
63
64// TreeDB
65// ================================================================================================
66
67/// An immutable merkle tree db that uses a byte slice key to specify the leaves in the tree.
68pub struct TreeDB<'db, const D: usize, H: Hasher> {
69    db: &'db dyn HashDBRef<H, DBValue>,
70    root: NodeHash<H>,
71    null_nodes: HashMap<H::Out, Node<H>>,
72    recorder: Option<core::cell::RefCell<&'db mut dyn TreeRecorder<H>>>,
73}
74
75impl<'db, const D: usize, H: Hasher> TreeDB<'db, D, H> {
76    /// Return the underlying db of a TreeDB
77    pub fn db(&self) -> &dyn HashDBRef<H, DBValue> {
78        self.db
79    }
80
81    /// Return the node associated with the provided hash. Retrieves the node from either the database
82    /// or the null node map if it is a default node.
83    fn lookup(&self, node_hash: &NodeHash<H>) -> Result<Node<H>, TreeError> {
84        let node = match node_hash {
85            NodeHash::InMemory(_) => {
86                return Err(TreeError::DataError(DataError::InMemoryNotSupported))
87            }
88            NodeHash::Database(hash) => {
89                let data = self.db.get(hash, EMPTY_PREFIX).ok_or(TreeError::DataError(
90                    DataError::DatabaseDataNotFound(hash.as_ref().to_vec()),
91                ))?;
92                let node: Node<H> = data.try_into().map_err(TreeError::NodeError)?;
93
94                if let Some(recorder) = self.recorder.as_ref() {
95                    recorder.borrow_mut().record(&node);
96                }
97
98                Ok(node)
99            }
100            NodeHash::Default(hash) => {
101                self.null_nodes
102                    .get(hash)
103                    .cloned()
104                    .ok_or(TreeError::DataError(DataError::NullNodeDataNotFound(
105                        hash.as_ref().to_vec(),
106                    )))
107            }
108        }?;
109
110        Ok(node)
111    }
112
113    /// Returns a leaf node for the provided key. If the leaf node does not exist, returns None.
114    /// If a proof is provided, the sibling hashes along the lookup path are stored in the proof.
115    fn lookup_leaf_node(
116        &self,
117        key: &Key<D>,
118        proof: &mut Option<Vec<DBValue>>,
119    ) -> Result<Option<Node<H>>, TreeError> {
120        let mut current_node = self.lookup(&self.root)?;
121
122        for bit in key.iter() {
123            let child_selector = ChildSelector::new(bit);
124            let child_hash = current_node
125                .child_hash(&child_selector)
126                .map_err(TreeError::NodeError)?;
127            if child_hash.is_default() && proof.is_none() {
128                return Ok(None);
129            }
130
131            // store the sibling hash in the proof
132            if let Some(proof) = proof.as_mut() {
133                let sibling_hash: H::Out = **current_node
134                    .child_hash(&child_selector.sibling())
135                    .map_err(TreeError::NodeError)?;
136                proof.push(sibling_hash.as_ref().to_vec());
137            }
138
139            current_node = self.lookup(child_hash)?;
140        }
141
142        Ok(Some(current_node))
143    }
144}
145
146impl<'db, H: Hasher, const D: usize> KeyedTree<H, D> for TreeDB<'db, D, H> {
147    /// Returns the root of the tree
148    fn root(&self) -> &H::Out {
149        &self.root
150    }
151
152    /// Returns the value associated with the given key
153    fn value(&self, key: &[u8]) -> Result<Option<DBValue>, TreeError> {
154        let key = Key::<D>::new(key).map_err(TreeError::KeyError)?;
155        let node = self.lookup_leaf_node(&key, &mut None)?;
156        match node {
157            Some(node) => Ok(Some(node.value().map_err(TreeError::NodeError)?.clone())),
158            None => Ok(None),
159        }
160    }
161
162    /// Returns the leaf associated with the given key
163    fn leaf(&self, key: &[u8]) -> Result<Option<H::Out>, TreeError> {
164        let key = Key::<D>::new(key).map_err(TreeError::KeyError)?;
165        let node = self.lookup_leaf_node(&key, &mut None)?;
166        match node {
167            Some(node) => Ok(Some(*node.hash())),
168            None => Ok(None),
169        }
170    }
171
172    /// Returns an inclusion proof of a value a the specified key.
173    /// Returns a tuple of form: (value, root, proof)  
174    fn proof(&self, key: &[u8]) -> Result<(Option<DBValue>, H::Out, Vec<DBValue>), TreeError> {
175        let key = Key::<D>::new(key).map_err(TreeError::KeyError)?;
176        let mut proof = Some(Vec::new());
177        let node = self.lookup_leaf_node(&key, &mut proof)?;
178        let root = *self.root.hash();
179        let mut proof = proof.unwrap();
180        proof.reverse();
181
182        match node {
183            Some(node) => {
184                let value = node.value().map_err(TreeError::NodeError)?.clone();
185                Ok((Some(value), root, proof))
186            }
187            None => Ok((None, root, proof)),
188        }
189    }
190
191    /// Verifies that the given value is in the tree with the given root at the given index
192    fn verify(
193        key: &[u8],
194        value: &[u8],
195        proof: &[DBValue],
196        root: &H::Out,
197    ) -> Result<bool, TreeError> {
198        let key = Key::<D>::new(key).map_err(TreeError::KeyError)?;
199        let mut hash = H::hash(value);
200        // iterate over the bits in the key in reverse order
201        for (bit, sibling) in (0..D * 8).rev().zip(proof.iter()) {
202            let bit = key.bit(bit).map_err(TreeError::KeyError)?;
203            let child_selector = ChildSelector::new(bit);
204            match child_selector {
205                ChildSelector::Left => {
206                    hash = H::hash(&[hash.as_ref(), sibling].concat());
207                }
208                ChildSelector::Right => {
209                    hash = H::hash(&[sibling, hash.as_ref()].concat());
210                }
211            }
212        }
213        Ok(hash == *root)
214    }
215}
merkle_tree_db/treedb.rs

merkle_tree_db/
treedb.rs
