light_indexed_merkle_tree/
copy.rs

1use std::{fmt, marker::PhantomData, ops::Deref};
2
3use light_bounded_vec::CyclicBoundedVecMetadata;
4use light_concurrent_merkle_tree::{
5    copy::ConcurrentMerkleTreeCopy,
6    errors::ConcurrentMerkleTreeError,
7    offset::copy::{read_cyclic_bounded_vec_at, read_value_at},
8};
9use light_hasher::Hasher;
10use num_traits::{CheckedAdd, CheckedSub, ToBytes, Unsigned};
11
12use crate::{errors::IndexedMerkleTreeError, IndexedMerkleTree};
13
14#[derive(Debug)]
15pub struct IndexedMerkleTreeCopy<H, I, const HEIGHT: usize, const NET_HEIGHT: usize>(
16    IndexedMerkleTree<H, I, HEIGHT, NET_HEIGHT>,
17)
18where
19    H: Hasher,
20    I: CheckedAdd
21        + CheckedSub
22        + Copy
23        + Clone
24        + fmt::Debug
25        + PartialOrd
26        + ToBytes
27        + TryFrom<usize>
28        + Unsigned,
29    usize: From<I>;
30
31impl<H, I, const HEIGHT: usize, const NET_HEIGHT: usize>
32    IndexedMerkleTreeCopy<H, I, HEIGHT, NET_HEIGHT>
33where
34    H: Hasher,
35    I: CheckedAdd
36        + CheckedSub
37        + Copy
38        + Clone
39        + fmt::Debug
40        + PartialOrd
41        + ToBytes
42        + TryFrom<usize>
43        + Unsigned,
44    usize: From<I>,
45{
46    /// Casts a byte slice into wrapped `IndexedMerkleTree` structure reference,
47    /// including dynamic fields.
48    ///
49    /// # Purpose
50    ///
51    /// This method is meant to be used mostly in Solana programs, where memory
52    /// constraints are tight and we want to make sure no data is copied.
53    pub fn from_bytes_copy(bytes: &[u8]) -> Result<Self, IndexedMerkleTreeError> {
54        let (merkle_tree, mut offset) =
55            ConcurrentMerkleTreeCopy::<H, HEIGHT>::struct_from_bytes_copy(bytes)?;
56
57        let indexed_changelog_metadata: CyclicBoundedVecMetadata =
58            unsafe { read_value_at(bytes, &mut offset) };
59
60        let expected_size = IndexedMerkleTree::<H, I, HEIGHT, NET_HEIGHT>::size_in_account(
61            merkle_tree.height,
62            merkle_tree.changelog.capacity(),
63            merkle_tree.roots.capacity(),
64            merkle_tree.canopy_depth,
65            indexed_changelog_metadata.capacity(),
66        );
67
68        if bytes.len() < expected_size {
69            return Err(IndexedMerkleTreeError::ConcurrentMerkleTree(
70                ConcurrentMerkleTreeError::BufferSize(expected_size, bytes.len()),
71            ));
72        }
73        let indexed_changelog =
74            unsafe { read_cyclic_bounded_vec_at(bytes, &mut offset, &indexed_changelog_metadata) };
75
76        Ok(Self(IndexedMerkleTree {
77            merkle_tree,
78            indexed_changelog,
79            _index: PhantomData,
80        }))
81    }
82}
83
84impl<H, I, const HEIGHT: usize, const NET_HEIGHT: usize> Deref
85    for IndexedMerkleTreeCopy<H, I, HEIGHT, NET_HEIGHT>
86where
87    H: Hasher,
88    I: CheckedAdd
89        + CheckedSub
90        + Copy
91        + Clone
92        + fmt::Debug
93        + PartialOrd
94        + ToBytes
95        + TryFrom<usize>
96        + Unsigned,
97    usize: From<I>,
98{
99    type Target = IndexedMerkleTree<H, I, HEIGHT, NET_HEIGHT>;
100
101    fn deref(&self) -> &Self::Target {
102        &self.0
103    }
104}
105
106#[cfg(test)]
107mod test {
108    use light_hasher::{bigint::bigint_to_be_bytes_array, Poseidon};
109    use num_bigint::RandBigInt;
110    use rand::thread_rng;
111
112    use super::*;
113    use crate::zero_copy::IndexedMerkleTreeZeroCopyMut;
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
