fuel_merkle/sparse/
in_memory.rs

1use crate::{
2    common::{
3        Bytes32,
4        StorageMap,
5    },
6    sparse::{
7        self,
8        merkle_tree::MerkleTreeKey,
9        proof::Proof,
10        Primitive,
11    },
12    storage::{
13        Mappable,
14        StorageInspect,
15        StorageMutate,
16    },
17};
18use alloc::{
19    borrow::Cow,
20    vec::Vec,
21};
22
23/// The table of the Sparse Merkle tree's nodes. [`MerkleTree`] works with it as a sparse
24/// merkle tree, where the storage key is `Bytes32` and the value is the
25/// [`Buffer`](crate::sparse::Buffer) (raw presentation of the
26/// [`Node`](crate::sparse::Node)).
27#[derive(Debug)]
28pub struct NodesTable;
29
30impl Mappable for NodesTable {
31    type Key = Self::OwnedKey;
32    type OwnedKey = Bytes32;
33    type OwnedValue = Primitive;
34    type Value = Self::OwnedValue;
35}
36
37type Storage = StorageMap<NodesTable>;
38type SparseMerkleTree = sparse::MerkleTree<NodesTable, Storage>;
39
40#[derive(Debug)]
41pub struct MerkleTree {
42    tree: SparseMerkleTree,
43}
44
45impl MerkleTree {
46    pub fn new() -> Self {
47        Self {
48            tree: SparseMerkleTree::new(Storage::new()),
49        }
50    }
51
52    /// Build a sparse Merkle tree from a set of key-value pairs. This is
53    /// equivalent to creating an empty sparse Merkle tree and sequentially
54    /// calling [update](Self::update) for each key-value pair. This constructor
55    /// is more performant than calling individual sequential updates and is the
56    /// preferred approach when the key-values are known upfront. Leaves can be
57    /// appended to the returned tree using `update` to further accumulate leaf
58    /// data.
59    pub fn from_set<I, D>(set: I) -> Self
60    where
61        I: Iterator<Item = (MerkleTreeKey, D)>,
62        D: AsRef<[u8]>,
63    {
64        let tree = SparseMerkleTree::from_set(Storage::new(), set)
65            .expect("`Storage` can't return error");
66        Self { tree }
67    }
68
69    /// Calculate the sparse Merkle root from a set of key-value pairs. This is
70    /// similar to constructing a new tree from a set of key-value pairs using
71    /// [from_set](Self::from_set), except this method returns only the root; it
72    /// does not write to storage nor return a sparse Merkle tree instance. It
73    /// is equivalent to calling `from_set(..)`, followed by `root()`, but does
74    /// not incur the overhead of storage writes. This can be helpful when we
75    /// know all the key-values in the set upfront and we will not need to
76    /// update the set in the future.
77    pub fn root_from_set<I, D>(set: I) -> Bytes32
78    where
79        I: Iterator<Item = (MerkleTreeKey, D)>,
80        D: AsRef<[u8]>,
81    {
82        #[derive(Default)]
83        struct EmptyStorage;
84
85        impl StorageInspect<NodesTable> for EmptyStorage {
86            type Error = core::convert::Infallible;
87
88            fn get(&self, _: &Bytes32) -> Result<Option<Cow<Primitive>>, Self::Error> {
89                Ok(None)
90            }
91
92            fn contains_key(&self, _: &Bytes32) -> Result<bool, Self::Error> {
93                Ok(false)
94            }
95        }
96
97        impl StorageMutate<NodesTable> for EmptyStorage {
98            fn insert(&mut self, _: &Bytes32, _: &Primitive) -> Result<(), Self::Error> {
99                Ok(())
100            }
101
102            fn replace(
103                &mut self,
104                _: &Bytes32,
105                _: &Primitive,
106            ) -> Result<Option<Primitive>, Self::Error> {
107                Ok(None)
108            }
109
110            fn remove(&mut self, _: &Bytes32) -> Result<(), Self::Error> {
111                Ok(())
112            }
113
114            fn take(&mut self, _: &Bytes32) -> Result<Option<Primitive>, Self::Error> {
115                Ok(None)
116            }
117        }
118
119        let tree = sparse::MerkleTree::<NodesTable, _>::from_set(EmptyStorage, set)
120            .expect("`Storage` can't return error");
121        tree.root()
122    }
123
124    /// Calculate the sparse Merkle root as well as all nodes in the Merkle tree
125    /// from a set of key-value pairs. This is similar to constructing a new
126    /// tree from a set of key-value pairs using [from_set](Self::from_set),
127    /// except this method returns only the root and the list of leaves and
128    /// nodes in the tree; it does not return a sparse Merkle tree instance.
129    /// This can be helpful when we know all the key-values in the set upfront
130    /// and we need to defer storage writes, such as expensive database inserts,
131    /// for batch operations later in the process.
132    pub fn nodes_from_set<I, D>(set: I) -> (Bytes32, Vec<(Bytes32, Primitive)>)
133    where
134        I: Iterator<Item = (MerkleTreeKey, D)>,
135        D: AsRef<[u8]>,
136    {
137        #[derive(Default)]
138        struct VectorStorage {
139            storage: Vec<(Bytes32, Primitive)>,
140        }
141
142        impl StorageInspect<NodesTable> for VectorStorage {
143            type Error = core::convert::Infallible;
144
145            fn get(&self, _: &Bytes32) -> Result<Option<Cow<Primitive>>, Self::Error> {
146                unimplemented!("Read operation is not supported")
147            }
148
149            fn contains_key(&self, _: &Bytes32) -> Result<bool, Self::Error> {
150                unimplemented!("Read operation is not supported")
151            }
152        }
153
154        impl StorageMutate<NodesTable> for VectorStorage {
155            fn insert(
156                &mut self,
157                key: &Bytes32,
158                value: &Primitive,
159            ) -> Result<(), Self::Error> {
160                self.storage.push((*key, *value));
161                Ok(())
162            }
163
164            fn replace(
165                &mut self,
166                key: &Bytes32,
167                value: &Primitive,
168            ) -> Result<Option<Primitive>, Self::Error> {
169                self.storage.push((*key, *value));
170                Ok(None)
171            }
172
173            fn remove(&mut self, _: &Bytes32) -> Result<(), Self::Error> {
174                unimplemented!("Remove operation is not supported")
175            }
176
177            fn take(&mut self, _: &Bytes32) -> Result<Option<Primitive>, Self::Error> {
178                unimplemented!("Take operation is not supported")
179            }
180        }
181
182        let tree =
183            sparse::MerkleTree::<NodesTable, _>::from_set(VectorStorage::default(), set)
184                .expect("`Storage` can't return error");
185        let root = tree.root();
186        let nodes = tree.into_storage().storage;
187
188        (root, nodes)
189    }
190
191    pub fn update(&mut self, key: MerkleTreeKey, data: &[u8]) {
192        let _ = self.tree.update(key, data);
193    }
194
195    pub fn delete(&mut self, key: MerkleTreeKey) {
196        let _ = self.tree.delete(key);
197    }
198
199    pub fn root(&self) -> Bytes32 {
200        self.tree.root()
201    }
202
203    pub fn generate_proof(&self, key: &MerkleTreeKey) -> Option<Proof> {
204        self.tree.generate_proof(key).ok()
205    }
206}
207
208impl Default for MerkleTree {
209    fn default() -> Self {
210        Self::new()
211    }
212}
213
214#[cfg(test)]
215mod test {
216    use super::*;
217    use crate::common::sum;
218
219    fn key(data: &[u8]) -> MerkleTreeKey {
220        MerkleTreeKey::new_without_hash(sum(data))
221    }
222
223    #[test]
224    fn test_empty_root() {
225        let tree = MerkleTree::new();
226        let root = tree.root();
227        let expected_root =
228            "0000000000000000000000000000000000000000000000000000000000000000";
229        assert_eq!(hex::encode(root), expected_root);
230    }
231
232    #[test]
233    fn test_update_1() {
234        let mut tree = MerkleTree::new();
235
236        tree.update(key(b"\x00\x00\x00\x00"), b"DATA");
237
238        let root = tree.root();
239        let expected_root =
240            "39f36a7cb4dfb1b46f03d044265df6a491dffc1034121bc1071a34ddce9bb14b";
241        assert_eq!(hex::encode(root), expected_root);
242    }
243
244    #[test]
245    fn test_update_2() {
246        let mut tree = MerkleTree::new();
247
248        tree.update(key(b"\x00\x00\x00\x00"), b"DATA");
249        tree.update(key(b"\x00\x00\x00\x01"), b"DATA");
250
251        let root = tree.root();
252        let expected_root =
253            "8d0ae412ca9ca0afcb3217af8bcd5a673e798bd6fd1dfacad17711e883f494cb";
254        assert_eq!(hex::encode(root), expected_root);
255    }
256
257    #[test]
258    fn test_update_3() {
259        let mut tree = MerkleTree::new();
260
261        tree.update(key(b"\x00\x00\x00\x00"), b"DATA");
262        tree.update(key(b"\x00\x00\x00\x01"), b"DATA");
263        tree.update(key(b"\x00\x00\x00\x02"), b"DATA");
264
265        let root = tree.root();
266        let expected_root =
267            "52295e42d8de2505fdc0cc825ff9fead419cbcf540d8b30c7c4b9c9b94c268b7";
268        assert_eq!(hex::encode(root), expected_root);
269    }
270
271    #[test]
272    fn test_update_1_delete_1() {
273        let mut tree = MerkleTree::new();
274
275        tree.update(key(b"\x00\x00\x00\x00"), b"DATA");
276        tree.delete(key(b"\x00\x00\x00\x00"));
277
278        let root = tree.root();
279        let expected_root =
280            "0000000000000000000000000000000000000000000000000000000000000000";
281        assert_eq!(hex::encode(root), expected_root);
282    }
283
284    #[test]
285    fn test_update_2_delete_1() {
286        let mut tree = MerkleTree::new();
287
288        tree.update(key(b"\x00\x00\x00\x00"), b"DATA");
289        tree.update(key(b"\x00\x00\x00\x01"), b"DATA");
290        tree.delete(key(b"\x00\x00\x00\x01"));
291
292        let root = tree.root();
293        let expected_root =
294            "39f36a7cb4dfb1b46f03d044265df6a491dffc1034121bc1071a34ddce9bb14b";
295        assert_eq!(hex::encode(root), expected_root);
296    }
297}