cw_merkle_tree/tree/
sparse_history.rs

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
11/// Like [SparseMerkleTree] but able to check valid root hash with all previous root hashes.
12pub 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}