light_concurrent_merkle_tree/
copy.rs

1use std::ops::Deref;
2
3use light_bounded_vec::{BoundedVecMetadata, CyclicBoundedVecMetadata};
4use light_hasher::Hasher;
5use memoffset::{offset_of, span_of};
6
7use crate::{
8    errors::ConcurrentMerkleTreeError,
9    offset::copy::{read_bounded_vec_at, read_cyclic_bounded_vec_at, read_value_at},
10    ConcurrentMerkleTree,
11};
12
13#[derive(Debug)]
14pub struct ConcurrentMerkleTreeCopy<H, const HEIGHT: usize>(ConcurrentMerkleTree<H, HEIGHT>)
15where
16    H: Hasher;
17
18impl<H, const HEIGHT: usize> ConcurrentMerkleTreeCopy<H, HEIGHT>
19where
20    H: Hasher,
21{
22    pub fn struct_from_bytes_copy(
23        bytes: &[u8],
24    ) -> Result<(ConcurrentMerkleTree<H, HEIGHT>, usize), ConcurrentMerkleTreeError> {
25        let expected_size = ConcurrentMerkleTree::<H, HEIGHT>::non_dyn_fields_size();
26        if bytes.len() < expected_size {
27            return Err(ConcurrentMerkleTreeError::BufferSize(
28                expected_size,
29                bytes.len(),
30            ));
31        }
32
33        let height = usize::from_le_bytes(
34            bytes[span_of!(ConcurrentMerkleTree<H, HEIGHT>, height)]
35                .try_into()
36                .unwrap(),
37        );
38        let canopy_depth = usize::from_le_bytes(
39            bytes[span_of!(ConcurrentMerkleTree<H, HEIGHT>, canopy_depth)]
40                .try_into()
41                .unwrap(),
42        );
43
44        let mut offset = offset_of!(ConcurrentMerkleTree<H, HEIGHT>, next_index);
45
46        let next_index = unsafe { read_value_at(bytes, &mut offset) };
47        let sequence_number = unsafe { read_value_at(bytes, &mut offset) };
48        let rightmost_leaf = unsafe { read_value_at(bytes, &mut offset) };
49        let filled_subtrees_metadata: BoundedVecMetadata =
50            unsafe { read_value_at(bytes, &mut offset) };
51        let changelog_metadata: CyclicBoundedVecMetadata =
52            unsafe { read_value_at(bytes, &mut offset) };
53        let roots_metadata: CyclicBoundedVecMetadata = unsafe { read_value_at(bytes, &mut offset) };
54        let canopy_metadata: BoundedVecMetadata = unsafe { read_value_at(bytes, &mut offset) };
55
56        let expected_size = ConcurrentMerkleTree::<H, HEIGHT>::size_in_account(
57            height,
58            changelog_metadata.capacity(),
59            roots_metadata.capacity(),
60            canopy_depth,
61        );
62        if bytes.len() < expected_size {
63            return Err(ConcurrentMerkleTreeError::BufferSize(
64                expected_size,
65                bytes.len(),
66            ));
67        }
68
69        let filled_subtrees =
70            unsafe { read_bounded_vec_at(bytes, &mut offset, &filled_subtrees_metadata) };
71        let changelog =
72            unsafe { read_cyclic_bounded_vec_at(bytes, &mut offset, &changelog_metadata) };
73        let roots = unsafe { read_cyclic_bounded_vec_at(bytes, &mut offset, &roots_metadata) };
74        let canopy = unsafe { read_bounded_vec_at(bytes, &mut offset, &canopy_metadata) };
75
76        let mut merkle_tree = ConcurrentMerkleTree::new(
77            height,
78            changelog_metadata.capacity(),
79            roots_metadata.capacity(),
80            canopy_depth,
81        )?;
82        // SAFETY: Tree is initialized.
83        unsafe {
84            *merkle_tree.next_index = next_index;
85            *merkle_tree.sequence_number = sequence_number;
86            *merkle_tree.rightmost_leaf = rightmost_leaf;
87        }
88        merkle_tree.filled_subtrees = filled_subtrees;
89        merkle_tree.changelog = changelog;
90        merkle_tree.roots = roots;
91        merkle_tree.canopy = canopy;
92
93        Ok((merkle_tree, offset))
94    }
95
96    pub fn from_bytes_copy(bytes: &[u8]) -> Result<Self, ConcurrentMerkleTreeError> {
97        let (merkle_tree, _) = Self::struct_from_bytes_copy(bytes)?;
98        merkle_tree.check_size_constraints()?;
99        Ok(Self(merkle_tree))
100    }
101}
102
103impl<H, const HEIGHT: usize> Deref for ConcurrentMerkleTreeCopy<H, HEIGHT>
104where
105    H: Hasher,
106{
107    type Target = ConcurrentMerkleTree<H, HEIGHT>;
108
109    fn deref(&self) -> &Self::Target {
110        &self.0
111    }
112}
113
114#[cfg(test)]
115mod test {
116    use ark_bn254::Fr;
117    use ark_ff::{BigInteger, PrimeField, UniformRand};
118    use light_hasher::Poseidon;
119    use rand::{thread_rng, Rng};
120
121    use super::*;
122    use crate::zero_copy::ConcurrentMerkleTreeZeroCopyMut;
123
124    fn from_bytes_copy<
125        const HEIGHT: usize,
126        const CHANGELOG: usize,
127        const ROOTS: usize,
128        const CANOPY_DEPTH: usize,
129        const OPERATIONS: usize,
130    >() {
131        let mut mt_1 =
132            ConcurrentMerkleTree::<Poseidon, HEIGHT>::new(HEIGHT, CHANGELOG, ROOTS, CANOPY_DEPTH)
133                .unwrap();
134        mt_1.init().unwrap();
135
136        // Create a buffer with random bytes - the `*_init` method should
137        // initialize the buffer gracefully and the randomness shouldn't cause
138        // undefined behavior.
139        let mut bytes = vec![
140            0u8;
141            ConcurrentMerkleTree::<Poseidon, HEIGHT>::size_in_account(
142                HEIGHT,
143                CHANGELOG,
144                ROOTS,
145                CANOPY_DEPTH
146            )
147        ];
148        thread_rng().fill(bytes.as_mut_slice());
149
150        // Initialize a Merkle tree on top of a byte slice.
151        {
152            let mut mt =
153                ConcurrentMerkleTreeZeroCopyMut::<Poseidon, HEIGHT>::from_bytes_zero_copy_init(
154                    bytes.as_mut_slice(),
155                    HEIGHT,
156                    CANOPY_DEPTH,
157                    CHANGELOG,
158                    ROOTS,
159                )
160                .unwrap();
161            mt.init().unwrap();
162
163            // Ensure that it was properly initialized.
164            assert_eq!(mt.height, HEIGHT);
165            assert_eq!(mt.canopy_depth, CANOPY_DEPTH);
166            assert_eq!(mt.next_index(), 0);
167            assert_eq!(mt.sequence_number(), 0);
168            assert_eq!(mt.rightmost_leaf(), Poseidon::zero_bytes()[0]);
169
170            assert_eq!(mt.filled_subtrees.capacity(), HEIGHT);
171            assert_eq!(mt.filled_subtrees.len(), HEIGHT);
172
173            assert_eq!(mt.changelog.capacity(), CHANGELOG);
174            assert_eq!(mt.changelog.len(), 1);
175
176            assert_eq!(mt.roots.capacity(), ROOTS);
177            assert_eq!(mt.roots.len(), 1);
178
179            assert_eq!(
180                mt.canopy.capacity(),
181                ConcurrentMerkleTree::<Poseidon, HEIGHT>::canopy_size(CANOPY_DEPTH)
182            );
183
184            assert_eq!(mt.root(), Poseidon::zero_bytes()[HEIGHT]);
185        }
186
187        let mut rng = thread_rng();
188
189        for _ in 0..OPERATIONS {
190            // Reload the tree from bytes on each iteration.
191            let mut mt_2 =
192                ConcurrentMerkleTreeZeroCopyMut::<Poseidon, HEIGHT>::from_bytes_zero_copy_mut(
193                    &mut bytes,
194                )
195                .unwrap();
196
197            let leaf: [u8; 32] = Fr::rand(&mut rng)
198                .into_bigint()
199                .to_bytes_be()
200                .try_into()
201                .unwrap();
202
203            mt_1.append(&leaf).unwrap();
204            mt_2.append(&leaf).unwrap();
205
206            assert_eq!(mt_1, *mt_2);
207        }
208
209        // Read a copy of that Merkle tree.
210        let mt_2 = ConcurrentMerkleTreeCopy::<Poseidon, HEIGHT>::from_bytes_copy(&bytes).unwrap();
211
212        assert_eq!(mt_1.height, mt_2.height);
213        assert_eq!(mt_1.canopy_depth, mt_2.canopy_depth);
214        assert_eq!(mt_1.next_index(), mt_2.next_index());
215        assert_eq!(mt_1.sequence_number(), mt_2.sequence_number());
216        assert_eq!(mt_1.rightmost_leaf(), mt_2.rightmost_leaf());
217        assert_eq!(
218            mt_1.filled_subtrees.as_slice(),
219            mt_2.filled_subtrees.as_slice()
220        );
221    }
222
223    #[test]
224    fn test_from_bytes_copy_26_1400_2400_10_256_1024() {
225        const HEIGHT: usize = 26;
226        const CHANGELOG_SIZE: usize = 1400;
227        const ROOTS: usize = 2400;
228        const CANOPY_DEPTH: usize = 10;
229
230        const OPERATIONS: usize = 1024;
231
232        from_bytes_copy::<HEIGHT, CHANGELOG_SIZE, ROOTS, CANOPY_DEPTH, OPERATIONS>()
233    }
234}
light_concurrent_merkle_tree/copy.rs

light_concurrent_merkle_tree/
copy.rs
