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#[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 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 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 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}