cw_merkle_tree/tree/
sparse_history_bounded.rs

1use std::fmt::Debug;
2
3use cosmwasm_std::{Empty, Order, Storage};
4use cw_storage_plus::{Bound, Item, Map, PrimaryKey};
5use serde::{de::DeserializeOwned, Serialize};
6
7use crate::{Hasher, MerkleTree, MerkleTreeError};
8
9use super::SparseMerkleTree;
10
11/// Like [SparseMerkleTree] but able to check valid root hash with previous root hashes upto specified history level.
12pub struct SparseMerkleTreeWithHistoryBounded<
13    'a,
14    L: Serialize + DeserializeOwned + Clone + Debug + PartialEq + PrimaryKey<'a>,
15    H: Hasher<L>,
16    const HISTORY_LEVEL: u32,
17> {
18    pub history_index: Item<'a, u32>,
19    pub root_history: Map<'a, L, Empty>,
20    pub root_index: Map<'a, u32, L>,
21    pub tree: SparseMerkleTree<'a, L, H>,
22}
23
24impl<
25        'a,
26        L: Serialize + DeserializeOwned + Clone + Debug + PartialEq + PrimaryKey<'a>,
27        H: Hasher<L>,
28        const HISTORY_LEVEL: u32,
29    > SparseMerkleTreeWithHistoryBounded<'a, L, H, HISTORY_LEVEL>
30{
31    pub const fn new(
32        hashes_ns: &'a str,
33        leafs_ns: &'a str,
34        level_ns: &'a str,
35        root_ns: &'a str,
36        root_history_ns: &'a str,
37        root_index_ns: &'a str,
38        history_index_ns: &'a str,
39    ) -> Self {
40        Self {
41            history_index: Item::new(history_index_ns),
42            root_history: Map::new(root_history_ns),
43            root_index: Map::new(root_index_ns),
44            tree: SparseMerkleTree::new(hashes_ns, leafs_ns, level_ns, root_ns),
45        }
46    }
47
48    /// Remove storage unused and out of range stored root.
49    /// The removed root might not be the earliest.
50    pub fn update_history_level(&self, storage: &mut dyn Storage) -> Result<(), MerkleTreeError> {
51        let updated_idx = self.history_index.may_load(storage)?.unwrap_or_default() % HISTORY_LEVEL;
52        self.history_index.save(storage, &updated_idx)?;
53
54        let mut root_range = self
55            .root_index
56            .range(
57                storage,
58                Some(Bound::inclusive(HISTORY_LEVEL)),
59                None,
60                Order::Ascending,
61            )
62            .collect::<Vec<_>>()
63            .into_iter();
64        while let Some(Ok((idx, root))) = root_range.next() {
65            self.root_index.remove(storage, idx);
66            self.root_history.remove(storage, root);
67        }
68
69        Ok(())
70    }
71}
72
73impl<
74        'a,
75        L: Serialize + DeserializeOwned + Clone + Debug + PartialEq + PrimaryKey<'a>,
76        H: Hasher<L>,
77        const HISTORY_LEVEL: u32,
78    > MerkleTree<L, H> for SparseMerkleTreeWithHistoryBounded<'a, L, H, HISTORY_LEVEL>
79{
80    fn init(
81        &self,
82        storage: &mut dyn Storage,
83        level: u8,
84        default_leaf: L,
85        hasher: &H,
86    ) -> Result<(), MerkleTreeError> {
87        self.tree.init(storage, level, default_leaf, hasher)
88    }
89
90    fn is_valid_root(&self, storage: &dyn Storage, root: &L) -> Result<bool, MerkleTreeError> {
91        Ok(self.root_history.has(storage, root.clone()))
92    }
93
94    fn insert(
95        &self,
96        storage: &mut dyn Storage,
97        leaf: L,
98        hasher: &H,
99    ) -> Result<(u64, L), MerkleTreeError> {
100        let (index, latest_root) = self.tree.insert(storage, leaf, hasher)?;
101        let cur_idx = self.history_index.may_load(storage)?.unwrap_or_default();
102        let next_idx = (cur_idx + 1) % HISTORY_LEVEL;
103
104        // Remove old root
105        if let Some(root) = self.root_index.may_load(storage, next_idx)? {
106            self.root_history.remove(storage, root);
107        }
108
109        // Insert new root
110        self.root_history
111            .save(storage, latest_root.clone(), &Empty {})?;
112        self.root_index.save(storage, next_idx, &latest_root)?;
113
114        // Update current index
115        self.history_index.save(storage, &next_idx)?;
116
117        Ok((index, latest_root))
118    }
119
120    fn get_latest_root(&self, storage: &dyn Storage) -> Result<L, MerkleTreeError> {
121        self.tree.get_latest_root(storage)
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use std::error::Error;
128
129    use cosmwasm_std::{testing::MockStorage, Uint256};
130
131    use crate::{test_utils::Blake2, Hasher, MerkleTree};
132
133    use super::SparseMerkleTreeWithHistoryBounded;
134
135    const TREE: SparseMerkleTreeWithHistoryBounded<Vec<u8>, Blake2, 5> =
136        SparseMerkleTreeWithHistoryBounded::new(
137            "hashes",
138            "leafs",
139            "level",
140            "zeros",
141            "root_history",
142            "root_index",
143            "history_index",
144        );
145    const ZERO: [u8; 32] = [
146        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,
147        0, 0,
148    ];
149
150    #[test]
151    fn init() -> Result<(), Box<dyn Error>> {
152        let mut storage = MockStorage::new();
153        let zero_vec = ZERO.to_vec();
154
155        TREE.init(
156            &mut storage,
157            20,
158            Blake2.hash_two(&zero_vec, &zero_vec)?,
159            &Blake2,
160        )?;
161
162        assert_eq!(
163            TREE.get_latest_root(&storage)?,
164            [
165                20, 114, 250, 18, 41, 94, 49, 107, 184, 78, 231, 47, 187, 225, 122, 14, 76, 178,
166                156, 226, 121, 99, 103, 48, 22, 79, 157, 174, 92, 246, 92, 50
167            ]
168        );
169
170        Ok(())
171    }
172
173    #[test]
174    fn insert() -> Result<(), Box<dyn Error>> {
175        let mut storage = MockStorage::new();
176        let zero_vec = ZERO.to_vec();
177        let one_vec = Uint256::one().to_be_bytes().to_vec();
178
179        TREE.init(
180            &mut storage,
181            20,
182            Blake2.hash_two(&zero_vec, &zero_vec)?,
183            &Blake2,
184        )?;
185
186        let leaf = Blake2.hash_two(&one_vec, &one_vec)?;
187
188        let (index, new_root) = TREE.insert(&mut storage, leaf.clone(), &Blake2)?;
189
190        assert_eq!(index, 0);
191        assert_eq!(
192            new_root,
193            [
194                144, 77, 181, 73, 235, 223, 13, 204, 30, 18, 199, 252, 182, 160, 89, 248, 240, 219,
195                173, 150, 189, 114, 165, 70, 40, 159, 110, 9, 165, 185, 17, 229
196            ]
197        );
198        assert_eq!(new_root, TREE.get_latest_root(&storage)?);
199        assert!(TREE.is_valid_root(&storage, &new_root)?);
200
201        let (index, new_root) = TREE.insert(&mut storage, leaf, &Blake2)?;
202
203        assert_eq!(index, 1);
204        assert_eq!(
205            new_root,
206            [
207                69, 102, 154, 15, 149, 187, 157, 26, 123, 248, 50, 67, 177, 207, 6, 143, 94, 80,
208                242, 17, 127, 26, 94, 197, 222, 220, 255, 245, 136, 20, 62, 132
209            ]
210        );
211        assert_eq!(new_root, TREE.get_latest_root(&storage)?);
212        assert!(TREE.is_valid_root(&storage, &new_root)?);
213
214        Ok(())
215    }
216
217    #[test]
218    fn root_history() -> Result<(), Box<dyn Error>> {
219        let mut storage = MockStorage::new();
220        let zero_vec = ZERO.to_vec();
221        let one_vec = Uint256::one().to_be_bytes().to_vec();
222
223        TREE.init(
224            &mut storage,
225            20,
226            Blake2.hash_two(&zero_vec, &zero_vec)?,
227            &Blake2,
228        )?;
229
230        let leaf = Blake2.hash_two(&one_vec, &one_vec)?;
231
232        let (_, very_old_root) = TREE.insert(&mut storage, leaf.clone(), &Blake2)?;
233        let (_, old_root) = TREE.insert(&mut storage, leaf.clone(), &Blake2)?;
234        let _ = TREE.insert(&mut storage, leaf.clone(), &Blake2)?;
235        let _ = TREE.insert(&mut storage, leaf.clone(), &Blake2)?;
236        let _ = TREE.insert(&mut storage, leaf.clone(), &Blake2)?;
237        let (_, new_root) = TREE.insert(&mut storage, leaf, &Blake2)?;
238
239        assert!(!TREE.is_valid_root(&storage, &very_old_root)?);
240        assert!(TREE.is_valid_root(&storage, &old_root)?);
241        assert!(TREE.is_valid_root(&storage, &new_root)?);
242
243        Ok(())
244    }
245}
cw_merkle_tree/tree/sparse_history_bounded.rs

cw_merkle_tree/tree/
sparse_history_bounded.rs
