light_indexed_merkle_tree/
copy.rs

1use std::{fmt, marker::PhantomData, ops::Deref};
2
3use crate::{errors::IndexedMerkleTreeError, IndexedMerkleTree};
4use light_bounded_vec::CyclicBoundedVecMetadata;
5use light_concurrent_merkle_tree::{
6    copy::ConcurrentMerkleTreeCopy, errors::ConcurrentMerkleTreeError,
7};
8use light_hasher::Hasher;
9use light_utils::offset::copy::{read_cyclic_bounded_vec_at, read_value_at};
10use num_traits::{CheckedAdd, CheckedSub, ToBytes, Unsigned};
11
12#[derive(Debug)]
13pub struct IndexedMerkleTreeCopy<H, I, const HEIGHT: usize, const NET_HEIGHT: usize>(
14    IndexedMerkleTree<H, I, HEIGHT, NET_HEIGHT>,
15)
16where
17    H: Hasher,
18    I: CheckedAdd
19        + CheckedSub
20        + Copy
21        + Clone
22        + fmt::Debug
23        + PartialOrd
24        + ToBytes
25        + TryFrom<usize>
26        + Unsigned,
27    usize: From<I>;
28
29impl<H, I, const HEIGHT: usize, const NET_HEIGHT: usize>
30    IndexedMerkleTreeCopy<H, I, HEIGHT, NET_HEIGHT>
31where
32    H: Hasher,
33    I: CheckedAdd
34        + CheckedSub
35        + Copy
36        + Clone
37        + fmt::Debug
38        + PartialOrd
39        + ToBytes
40        + TryFrom<usize>
41        + Unsigned,
42    usize: From<I>,
43{
44    /// Casts a byte slice into wrapped `IndexedMerkleTree` structure reference,
45    /// including dynamic fields.
46    ///
47    /// # Purpose
48    ///
49    /// This method is meant to be used mostly in Solana programs, where memory
50    /// constraints are tight and we want to make sure no data is copied.
51    pub fn from_bytes_copy(bytes: &[u8]) -> Result<Self, IndexedMerkleTreeError> {
52        let (merkle_tree, mut offset) =
53            ConcurrentMerkleTreeCopy::<H, HEIGHT>::struct_from_bytes_copy(bytes)?;
54
55        let indexed_changelog_metadata: CyclicBoundedVecMetadata =
56            unsafe { read_value_at(bytes, &mut offset) };
57
58        let expected_size = IndexedMerkleTree::<H, I, HEIGHT, NET_HEIGHT>::size_in_account(
59            merkle_tree.height,
60            merkle_tree.changelog.capacity(),
61            merkle_tree.roots.capacity(),
62            merkle_tree.canopy_depth,
63            indexed_changelog_metadata.capacity(),
64        );
65
66        if bytes.len() < expected_size {
67            return Err(IndexedMerkleTreeError::ConcurrentMerkleTree(
68                ConcurrentMerkleTreeError::BufferSize(expected_size, bytes.len()),
69            ));
70        }
71        let indexed_changelog =
72            unsafe { read_cyclic_bounded_vec_at(bytes, &mut offset, &indexed_changelog_metadata) };
73
74        Ok(Self(IndexedMerkleTree {
75            merkle_tree,
76            indexed_changelog,
77            _index: PhantomData,
78        }))
79    }
80}
81
82impl<H, I, const HEIGHT: usize, const NET_HEIGHT: usize> Deref
83    for IndexedMerkleTreeCopy<H, I, HEIGHT, NET_HEIGHT>
84where
85    H: Hasher,
86    I: CheckedAdd
87        + CheckedSub
88        + Copy
89        + Clone
90        + fmt::Debug
91        + PartialOrd
92        + ToBytes
93        + TryFrom<usize>
94        + Unsigned,
95    usize: From<I>,
96{
97    type Target = IndexedMerkleTree<H, I, HEIGHT, NET_HEIGHT>;
98
99    fn deref(&self) -> &Self::Target {
100        &self.0
101    }
102}
103
104#[cfg(test)]
105mod test {
106    use light_hasher::Poseidon;
107    use light_utils::bigint::bigint_to_be_bytes_array;
108    use num_bigint::RandBigInt;
109    use rand::thread_rng;
110
111    use crate::zero_copy::IndexedMerkleTreeZeroCopyMut;
112
113    use super::*;
114
115    fn from_bytes_copy<
116        const HEIGHT: usize,
117        const CHANGELOG_SIZE: usize,
118        const ROOTS: usize,
119        const CANOPY_DEPTH: usize,
120        const INDEXED_CHANGELOG_SIZE: usize,
121        const OPERATIONS: usize,
122        const NET_HEIGHT: usize,
123    >() {
124        let mut mt_1 = IndexedMerkleTree::<Poseidon, usize, HEIGHT, NET_HEIGHT>::new(
125            HEIGHT,
126            CHANGELOG_SIZE,
127            ROOTS,
128            CANOPY_DEPTH,
129            INDEXED_CHANGELOG_SIZE,
130        )
131        .unwrap();
132        mt_1.init().unwrap();
133
134        let mut bytes = vec![
135            0u8;
136            IndexedMerkleTree::<Poseidon, usize, HEIGHT, NET_HEIGHT>::size_in_account(
137                HEIGHT,
138                CHANGELOG_SIZE,
139                ROOTS,
140                CANOPY_DEPTH,
141                INDEXED_CHANGELOG_SIZE
142            )
143        ];
144
145        {
146            let mut mt_2 =
147                IndexedMerkleTreeZeroCopyMut::<Poseidon, usize, HEIGHT, NET_HEIGHT>::from_bytes_zero_copy_init(
148                    &mut bytes,
149                    HEIGHT,
150                    CANOPY_DEPTH,
151                    CHANGELOG_SIZE,
152                    ROOTS,
153                    INDEXED_CHANGELOG_SIZE,
154                )
155                .unwrap();
156            mt_2.init().unwrap();
157
158            assert_eq!(mt_1, *mt_2);
159        }
160
161        let mut rng = thread_rng();
162
163        for _ in 0..OPERATIONS {
164            // Reload the tree from bytes on each iteration.
165            let mut mt_2 =
166                IndexedMerkleTreeZeroCopyMut::<Poseidon, usize, HEIGHT,NET_HEIGHT>::from_bytes_zero_copy_mut(
167                    &mut bytes,
168                )
169                .unwrap();
170
171            let leaf: [u8; 32] = bigint_to_be_bytes_array::<32>(&rng.gen_biguint(248)).unwrap();
172            mt_1.append(&leaf).unwrap();
173            mt_2.append(&leaf).unwrap();
174
175            assert_eq!(mt_1, *mt_2);
176        }
177
178        // Read a copy of that Merkle tree.
179        let mt_2 =
180            IndexedMerkleTreeCopy::<Poseidon, usize, HEIGHT, NET_HEIGHT>::from_bytes_copy(&bytes)
181                .unwrap();
182
183        assert_eq!(mt_1, *mt_2);
184    }
185
186    #[test]
187    fn test_from_bytes_copy_26_1400_2400_10_256_1024() {
188        const HEIGHT: usize = 26;
189        const CHANGELOG_SIZE: usize = 1400;
190        const ROOTS: usize = 2400;
191        const CANOPY_DEPTH: usize = 10;
192        const INDEXED_CHANGELOG_SIZE: usize = 256;
193        const NET_HEIGHT: usize = 16;
194        const OPERATIONS: usize = 1024;
195
196        from_bytes_copy::<
197            HEIGHT,
198            CHANGELOG_SIZE,
199            ROOTS,
200            CANOPY_DEPTH,
201            INDEXED_CHANGELOG_SIZE,
202            OPERATIONS,
203            NET_HEIGHT,
204        >()
205    }
206}
light_indexed_merkle_tree/copy.rs

light_indexed_merkle_tree/
copy.rs
