1use std::fmt::Debug;
2
3use cosmwasm_std::Empty;
4use cw_storage_plus::{Map, PrimaryKey};
5use serde::{de::DeserializeOwned, Serialize};
6
7use crate::{Hasher, MerkleTree};
8
9use super::SparseMerkleTree;
10
11pub struct SparseMerkleTreeWithHistory<
13 'a,
14 L: Serialize + DeserializeOwned + Clone + Debug + PartialEq + PrimaryKey<'a>,
15 H: Hasher<L>,
16> {
17 pub tree: SparseMerkleTree<'a, L, H>,
18 pub root_history: Map<'a, L, Empty>,
19}
20
21impl<
22 'a,
23 L: Serialize + DeserializeOwned + Clone + Debug + PartialEq + PrimaryKey<'a>,
24 H: Hasher<L>,
25 > SparseMerkleTreeWithHistory<'a, L, H>
26{
27 pub const fn new(
28 hashes_ns: &'a str,
29 leafs_ns: &'a str,
30 level_ns: &'a str,
31 root_ns: &'a str,
32 root_history_ns: &'a str,
33 ) -> Self {
34 Self {
35 tree: SparseMerkleTree::new(hashes_ns, leafs_ns, level_ns, root_ns),
36 root_history: Map::new(root_history_ns),
37 }
38 }
39}
40
41impl<
42 'a,
43 L: Serialize + DeserializeOwned + Clone + Debug + PartialEq + PrimaryKey<'a>,
44 H: Hasher<L>,
45 > MerkleTree<L, H> for SparseMerkleTreeWithHistory<'a, L, H>
46{
47 fn init(
48 &self,
49 storage: &mut dyn cosmwasm_std::Storage,
50 level: u8,
51 default_leaf: L,
52 hasher: &H,
53 ) -> Result<(), crate::MerkleTreeError> {
54 self.tree.init(storage, level, default_leaf, hasher)
55 }
56
57 fn is_valid_root(
58 &self,
59 storage: &dyn cosmwasm_std::Storage,
60 root: &L,
61 ) -> Result<bool, crate::MerkleTreeError> {
62 Ok(self.root_history.has(storage, root.clone()))
63 }
64
65 fn insert(
66 &self,
67 storage: &mut dyn cosmwasm_std::Storage,
68 leaf: L,
69 hasher: &H,
70 ) -> Result<(u64, L), crate::MerkleTreeError> {
71 let (index, latest_root) = self.tree.insert(storage, leaf, hasher)?;
72
73 self.root_history
74 .save(storage, latest_root.clone(), &Empty {})?;
75
76 Ok((index, latest_root))
77 }
78
79 fn get_latest_root(
80 &self,
81 storage: &dyn cosmwasm_std::Storage,
82 ) -> Result<L, crate::MerkleTreeError> {
83 self.tree.get_latest_root(storage)
84 }
85}
86
87#[cfg(test)]
88mod tests {
89 use std::error::Error;
90
91 use cosmwasm_std::{testing::MockStorage, Uint256};
92
93 use crate::{test_utils::Blake2, Hasher, MerkleTree};
94
95 use super::SparseMerkleTreeWithHistory;
96
97 const TREE: SparseMerkleTreeWithHistory<Vec<u8>, Blake2> =
98 SparseMerkleTreeWithHistory::new("hashes", "leafs", "level", "zeros", "root_history");
99 const ZERO: [u8; 32] = [
100 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
101 0, 0,
102 ];
103
104 #[test]
105 fn init() -> Result<(), Box<dyn Error>> {
106 let mut storage = MockStorage::new();
107 let zero_vec = ZERO.to_vec();
108
109 TREE.init(
110 &mut storage,
111 20,
112 Blake2.hash_two(&zero_vec, &zero_vec)?,
113 &Blake2,
114 )?;
115
116 assert_eq!(
117 TREE.get_latest_root(&storage)?,
118 [
119 20, 114, 250, 18, 41, 94, 49, 107, 184, 78, 231, 47, 187, 225, 122, 14, 76, 178,
120 156, 226, 121, 99, 103, 48, 22, 79, 157, 174, 92, 246, 92, 50
121 ]
122 );
123
124 Ok(())
125 }
126
127 #[test]
128 fn insert() -> Result<(), Box<dyn Error>> {
129 let mut storage = MockStorage::new();
130 let zero_vec = ZERO.to_vec();
131 let one_vec = Uint256::one().to_be_bytes().to_vec();
132
133 TREE.init(
134 &mut storage,
135 20,
136 Blake2.hash_two(&zero_vec, &zero_vec)?,
137 &Blake2,
138 )?;
139
140 let leaf = Blake2.hash_two(&one_vec, &one_vec)?;
141
142 let (index, new_root) = TREE.insert(&mut storage, leaf.clone(), &Blake2)?;
143
144 assert_eq!(index, 0);
145 assert_eq!(
146 new_root,
147 [
148 144, 77, 181, 73, 235, 223, 13, 204, 30, 18, 199, 252, 182, 160, 89, 248, 240, 219,
149 173, 150, 189, 114, 165, 70, 40, 159, 110, 9, 165, 185, 17, 229
150 ]
151 );
152 assert_eq!(new_root, TREE.get_latest_root(&storage)?);
153 assert!(TREE.is_valid_root(&storage, &new_root)?);
154
155 let (index, new_root) = TREE.insert(&mut storage, leaf, &Blake2)?;
156
157 assert_eq!(index, 1);
158 assert_eq!(
159 new_root,
160 [
161 69, 102, 154, 15, 149, 187, 157, 26, 123, 248, 50, 67, 177, 207, 6, 143, 94, 80,
162 242, 17, 127, 26, 94, 197, 222, 220, 255, 245, 136, 20, 62, 132
163 ]
164 );
165 assert_eq!(new_root, TREE.get_latest_root(&storage)?);
166 assert!(TREE.is_valid_root(&storage, &new_root)?);
167
168 Ok(())
169 }
170
171 #[test]
172 fn root_history() -> Result<(), Box<dyn Error>> {
173 let mut storage = MockStorage::new();
174 let zero_vec = ZERO.to_vec();
175 let one_vec = Uint256::one().to_be_bytes().to_vec();
176
177 TREE.init(
178 &mut storage,
179 20,
180 Blake2.hash_two(&zero_vec, &zero_vec)?,
181 &Blake2,
182 )?;
183
184 let leaf = Blake2.hash_two(&one_vec, &one_vec)?;
185
186 let (_, old_root) = TREE.insert(&mut storage, leaf.clone(), &Blake2)?;
187 let (_, new_root) = TREE.insert(&mut storage, leaf, &Blake2)?;
188
189 assert!(TREE.is_valid_root(&storage, &old_root)?);
190 assert!(TREE.is_valid_root(&storage, &new_root)?);
191
192 Ok(())
193 }
194}