1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206

use std::{marker::PhantomData, mem, slice};

use light_bounded_vec::{BoundedVec, CyclicBoundedVec, Pod};
use light_concurrent_merkle_tree::{
    changelog::ChangelogEntry, errors::ConcurrentMerkleTreeError, ConcurrentMerkleTree,
};
use light_hasher::Hasher;
use num_traits::{CheckedAdd, CheckedSub, ToBytes, Unsigned};

use crate::{errors::IndexedMerkleTreeError, IndexedMerkleTree, RawIndexedElement};

#[derive(Debug)]
pub struct IndexedMerkleTreeCopy<'a, H, I, const HEIGHT: usize>(
    pub IndexedMerkleTree<'a, H, I, HEIGHT>,
)
where
    H: Hasher,
    I: CheckedAdd
        + CheckedSub
        + Copy
        + Clone
        + PartialOrd
        + ToBytes
        + TryFrom<usize>
        + Unsigned
        + Pod,
    usize: From<I>;

pub type IndexedMerkleTreeCopy22<'a, H, I> = IndexedMerkleTreeCopy<'a, H, I, 22>;
pub type IndexedMerkleTreeCopy26<'a, H, I> = IndexedMerkleTreeCopy<'a, H, I, 26>;
pub type IndexedMerkleTreeCopy32<'a, H, I> = IndexedMerkleTreeCopy<'a, H, I, 32>;
pub type IndexedMerkleTreeCopy40<'a, H, I> = IndexedMerkleTreeCopy<'a, H, I, 40>;

impl<'a, H, I, const HEIGHT: usize> IndexedMerkleTreeCopy<'a, H, I, HEIGHT>
where
    H: Hasher,
    I: CheckedAdd
        + CheckedSub
        + Copy
        + Clone
        + PartialOrd
        + ToBytes
        + TryFrom<usize>
        + Unsigned
        + Pod,
    usize: From<I>,
{
    /// Casts a byte slice into wrapped `IndexedMerkleTree` structure reference,
    /// including dynamic fields.
    ///
    /// # Purpose
    ///
    /// This method is meant to be used mostly in Solana programs, where memory
    /// constraints are tight and we want to make sure no data is copied.
    ///
    /// # Safety
    ///
    /// This is highly unsafe. Ensuring the alignment and that the slice
    /// provides actual data of the hash set is the caller's responsibility.
    ///
    /// Calling it in async context (or anyhwere where the underlying data can
    /// be moved in the memory) is certainly going to cause undefined behavior.
    pub unsafe fn copy_from_bytes(
        bytes_struct: &[u8],
        bytes_filled_subtrees: &[u8],
        bytes_changelog: &[u8],
        bytes_roots: &[u8],
        bytes_canopy: &[u8],
        bytes_indexed_changelog: &'a [u8],
    ) -> Result<Self, IndexedMerkleTreeError> {
        let expected_bytes_struct_size = mem::size_of::<IndexedMerkleTree<'a, H, I, HEIGHT>>();
        if bytes_struct.len() != expected_bytes_struct_size {
            return Err(IndexedMerkleTreeError::ConcurrentMerkleTree(
                ConcurrentMerkleTreeError::StructBufferSize(
                    expected_bytes_struct_size,
                    bytes_struct.len(),
                ),
            ));
        }
        let struct_ref: *mut IndexedMerkleTree<'a, H, I, HEIGHT> = bytes_struct.as_ptr() as _;

        let mut merkle_tree = unsafe {
            ConcurrentMerkleTree {
                height: (*struct_ref).merkle_tree.height,

                changelog_capacity: (*struct_ref).merkle_tree.changelog_capacity,
                changelog_length: (*struct_ref).merkle_tree.changelog_length,
                current_changelog_index: (*struct_ref).merkle_tree.current_changelog_index,

                roots_capacity: (*struct_ref).merkle_tree.roots_capacity,
                roots_length: (*struct_ref).merkle_tree.roots_length,
                current_root_index: (*struct_ref).merkle_tree.current_root_index,

                canopy_depth: (*struct_ref).merkle_tree.canopy_depth,

                next_index: (*struct_ref).merkle_tree.next_index,
                sequence_number: (*struct_ref).merkle_tree.sequence_number,
                rightmost_leaf: (*struct_ref).merkle_tree.rightmost_leaf,

                filled_subtrees: BoundedVec::with_capacity((*struct_ref).merkle_tree.height),
                changelog: CyclicBoundedVec::with_capacity(
                    (*struct_ref).merkle_tree.changelog_capacity,
                ),
                roots: CyclicBoundedVec::with_capacity((*struct_ref).merkle_tree.roots_capacity),
                canopy: BoundedVec::with_capacity(ConcurrentMerkleTree::<H, HEIGHT>::canopy_size(
                    (*struct_ref).merkle_tree.canopy_depth,
                )),

                _hasher: PhantomData,
            }
        };

        let expected_bytes_filled_subtrees_size =
            mem::size_of::<[u8; 32]>() * (*struct_ref).merkle_tree.height;
        if bytes_filled_subtrees.len() != expected_bytes_filled_subtrees_size {
            return Err(ConcurrentMerkleTreeError::FilledSubtreesBufferSize(
                expected_bytes_filled_subtrees_size,
                bytes_filled_subtrees.len(),
            )
            .into());
        }
        let filled_subtrees: &[[u8; 32]] = slice::from_raw_parts(
            bytes_filled_subtrees.as_ptr() as *const _,
            (*struct_ref).merkle_tree.height,
        );
        for subtree in filled_subtrees.iter() {
            merkle_tree.filled_subtrees.push(*subtree)?;
        }

        let expected_bytes_changelog_size =
            mem::size_of::<ChangelogEntry<HEIGHT>>() * (*struct_ref).merkle_tree.changelog_capacity;
        if bytes_changelog.len() != expected_bytes_changelog_size {
            return Err(ConcurrentMerkleTreeError::ChangelogBufferSize(
                expected_bytes_changelog_size,
                bytes_changelog.len(),
            )
            .into());
        }
        let changelog: &[ChangelogEntry<HEIGHT>] = slice::from_raw_parts(
            bytes_changelog.as_ptr() as *const _,
            (*struct_ref).merkle_tree.changelog_length,
        );
        for changelog_entry in changelog.iter() {
            merkle_tree.changelog.push(changelog_entry.clone());
        }

        let expected_bytes_roots_size =
            mem::size_of::<[u8; 32]>() * (*struct_ref).merkle_tree.roots_capacity;
        if bytes_roots.len() != expected_bytes_roots_size {
            return Err(ConcurrentMerkleTreeError::RootBufferSize(
                expected_bytes_roots_size,
                bytes_roots.len(),
            )
            .into());
        }
        let roots: &[[u8; 32]] = slice::from_raw_parts(
            bytes_roots.as_ptr() as *const _,
            (*struct_ref).merkle_tree.roots_length,
        );
        for root in roots.iter() {
            merkle_tree.roots.push(*root);
        }

        let canopy_size =
            ConcurrentMerkleTree::<H, HEIGHT>::canopy_size((*struct_ref).merkle_tree.canopy_depth);
        let expected_canopy_size = mem::size_of::<[u8; 32]>() * canopy_size;
        if bytes_canopy.len() != expected_canopy_size {
            return Err(ConcurrentMerkleTreeError::CanopyBufferSize(
                expected_canopy_size,
                bytes_canopy.len(),
            )
            .into());
        }
        let canopy: &[[u8; 32]] =
            slice::from_raw_parts(bytes_canopy.as_ptr() as *const _, canopy_size);
        for node in canopy.iter() {
            merkle_tree.canopy.push(*node)?;
        }

        let mut indexed_merkle_tree = unsafe {
            IndexedMerkleTree {
                merkle_tree,
                changelog: CyclicBoundedVec::with_capacity((*struct_ref).changelog.capacity()),
                _index: PhantomData,
            }
        };

        let expected_bytes_indexed_changelog_size =
            mem::size_of::<RawIndexedElement<I>>() * (*struct_ref).changelog.capacity();
        if bytes_indexed_changelog.len() != expected_bytes_indexed_changelog_size {
            return Err(IndexedMerkleTreeError::ChangelogBufferSize(
                expected_bytes_indexed_changelog_size,
                bytes_indexed_changelog.len(),
            ));
        }
        let indexed_changelog: &[RawIndexedElement<I>] = slice::from_raw_parts(
            bytes_indexed_changelog.as_ptr() as *const _,
            (*struct_ref).changelog.len(),
        );
        for changelog_entry in indexed_changelog.iter() {
            indexed_merkle_tree.changelog.push(*changelog_entry);
        }

        Ok(IndexedMerkleTreeCopy(indexed_merkle_tree))
    }
}