hpl_account_compression/state/
concurrent_merkle_tree_header.rs

1use anchor_lang::prelude::*;
2use borsh::{BorshDeserialize, BorshSerialize};
3
4use spl_concurrent_merkle_tree::concurrent_merkle_tree::ConcurrentMerkleTree;
5use std::mem::size_of;
6
7use crate::error::AccountCompressionError;
8
9pub const CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1: usize = 2 + 54;
10
11#[derive(Debug, Copy, Clone, PartialEq, BorshDeserialize, BorshSerialize)]
12#[repr(u8)]
13pub enum CompressionAccountType {
14    /// Uninitialized
15    Uninitialized,
16
17    /// SPL ConcurrentMerkleTree data structure, may include a Canopy
18    ConcurrentMerkleTree,
19}
20
21impl std::fmt::Display for CompressionAccountType {
22    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
23        write!(f, "{:?}", &self)
24    }
25}
26
27#[cfg(feature = "idl-build")]
28impl anchor_lang::IdlBuild for CompressionAccountType {}
29
30/// Initialization parameters for an SPL ConcurrentMerkleTree.
31///
32/// Only the following permutations are valid:
33///
34/// | max_depth | max_buffer_size       |
35/// | --------- | --------------------- |
36/// | 14        | (64, 256, 1024, 2048) |           
37/// | 20        | (64, 256, 1024, 2048) |           
38/// | 24        | (64, 256, 512, 1024, 2048) |           
39/// | 26        | (64, 256, 512, 1024, 2048) |           
40/// | 30        | (512, 1024, 2048) |           
41///
42#[repr(C)]
43#[derive(AnchorDeserialize, AnchorSerialize)]
44pub struct ConcurrentMerkleTreeHeader {
45    /// Account type
46    pub account_type: CompressionAccountType,
47    /// Versioned header
48    pub header: ConcurrentMerkleTreeHeaderData,
49}
50
51#[repr(C)]
52#[derive(AnchorDeserialize, AnchorSerialize)]
53pub struct ConcurrentMerkleTreeHeaderDataV1 {
54    /// Buffer of changelogs stored on-chain.
55    /// Must be a power of 2; see above table for valid combinations.
56    max_buffer_size: u32,
57
58    /// Depth of the SPL ConcurrentMerkleTree to store.
59    /// Tree capacity can be calculated as power(2, max_depth).
60    /// See above table for valid options.
61    max_depth: u32,
62
63    /// Authority that validates the content of the trees.
64    /// Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs.
65    authority: Pubkey,
66
67    /// Slot corresponding to when the Merkle tree was created.
68    /// Provides a lower-bound on what slot to start (re-)building a tree from.
69    creation_slot: u64,
70
71    /// A flag indicating whether the tree has been initialized with a root.
72    /// This field was added together with the `finalize_tree_with_root` instruction.
73    /// It takes 1 byte of space taken from the previous padding for existing accounts.
74    is_batch_initialized: bool,
75
76    /// Needs padding for the account to be 8-byte aligned
77    /// 8-byte alignment is necessary to zero-copy the SPL ConcurrentMerkleTree
78    _padding: [u8; 5],
79}
80
81#[repr(C)]
82#[derive(AnchorDeserialize, AnchorSerialize)]
83pub enum ConcurrentMerkleTreeHeaderData {
84    V1(ConcurrentMerkleTreeHeaderDataV1),
85}
86
87impl ConcurrentMerkleTreeHeader {
88    pub fn initialize(
89        &mut self,
90        max_depth: u32,
91        max_buffer_size: u32,
92        authority: &Pubkey,
93        creation_slot: u64,
94    ) {
95        self.account_type = CompressionAccountType::ConcurrentMerkleTree;
96
97        match self.header {
98            ConcurrentMerkleTreeHeaderData::V1(ref mut header) => {
99                // Double check header is empty after deserialization from zero'd bytes
100                assert_eq!(header.max_buffer_size, 0);
101                assert_eq!(header.max_depth, 0);
102                header.max_buffer_size = max_buffer_size;
103                header.max_depth = max_depth;
104                header.authority = *authority;
105                header.creation_slot = creation_slot;
106                // is_batch_initialized is left false by default
107            }
108        }
109    }
110
111    /// Initializes the header with the given parameters and sets the `is_batch_initialized` flag to
112    /// true.
113    pub fn initialize_batched(
114        &mut self,
115        max_depth: u32,
116        max_buffer_size: u32,
117        authority: &Pubkey,
118        creation_slot: u64,
119    ) {
120        self.initialize(max_depth, max_buffer_size, authority, creation_slot);
121        match self.header {
122            ConcurrentMerkleTreeHeaderData::V1(ref mut header) => {
123                header.is_batch_initialized = true;
124            }
125        }
126    }
127
128    pub fn get_max_depth(&self) -> u32 {
129        match &self.header {
130            ConcurrentMerkleTreeHeaderData::V1(header) => header.max_depth,
131        }
132    }
133
134    pub fn get_max_buffer_size(&self) -> u32 {
135        match &self.header {
136            ConcurrentMerkleTreeHeaderData::V1(header) => header.max_buffer_size,
137        }
138    }
139
140    pub fn get_creation_slot(&self) -> u64 {
141        match &self.header {
142            ConcurrentMerkleTreeHeaderData::V1(header) => header.creation_slot,
143        }
144    }
145
146    pub fn get_is_batch_initialized(&self) -> bool {
147        match &self.header {
148            ConcurrentMerkleTreeHeaderData::V1(header) => header.is_batch_initialized,
149        }
150    }
151
152    pub fn set_new_authority(&mut self, new_authority: &Pubkey) {
153        match self.header {
154            ConcurrentMerkleTreeHeaderData::V1(ref mut header) => {
155                header.authority = new_authority.clone();
156                msg!("Authority transferred to: {:?}", header.authority);
157            }
158        }
159    }
160
161    pub fn assert_valid(&self) -> Result<()> {
162        require_eq!(
163            self.account_type,
164            CompressionAccountType::ConcurrentMerkleTree,
165            AccountCompressionError::IncorrectAccountType,
166        );
167        Ok(())
168    }
169
170    pub fn assert_valid_authority(&self, expected_authority: &Pubkey) -> Result<()> {
171        self.assert_valid()?;
172        match &self.header {
173            ConcurrentMerkleTreeHeaderData::V1(header) => {
174                require_eq!(
175                    header.authority,
176                    *expected_authority,
177                    AccountCompressionError::IncorrectAuthority,
178                );
179            }
180        }
181        Ok(())
182    }
183
184    pub fn assert_valid_leaf_index(&self, leaf_index: u32) -> Result<()> {
185        if leaf_index >= (1 << self.get_max_depth()) {
186            return Err(AccountCompressionError::LeafIndexOutOfBounds.into());
187        }
188        Ok(())
189    }
190
191    pub fn assert_is_batch_initialized(&self) -> Result<()> {
192        match &self.header {
193            ConcurrentMerkleTreeHeaderData::V1(header) => {
194                require!(
195                    header.is_batch_initialized,
196                    AccountCompressionError::BatchNotInitialized
197                );
198            }
199        }
200        Ok(())
201    }
202}
203
204pub fn merkle_tree_get_size(header: &ConcurrentMerkleTreeHeader) -> Result<usize> {
205    // Note: max_buffer_size MUST be a power of 2
206    match (header.get_max_depth(), header.get_max_buffer_size()) {
207        (3, 8) => Ok(size_of::<ConcurrentMerkleTree<3, 8>>()),
208        (5, 8) => Ok(size_of::<ConcurrentMerkleTree<5, 8>>()),
209        (6, 16) => Ok(size_of::<ConcurrentMerkleTree<6, 16>>()),
210        (7, 16) => Ok(size_of::<ConcurrentMerkleTree<7, 16>>()),
211        (8, 16) => Ok(size_of::<ConcurrentMerkleTree<8, 16>>()),
212        (9, 16) => Ok(size_of::<ConcurrentMerkleTree<9, 16>>()),
213        (10, 32) => Ok(size_of::<ConcurrentMerkleTree<10, 32>>()),
214        (11, 32) => Ok(size_of::<ConcurrentMerkleTree<11, 32>>()),
215        (12, 32) => Ok(size_of::<ConcurrentMerkleTree<12, 32>>()),
216        (13, 32) => Ok(size_of::<ConcurrentMerkleTree<13, 32>>()),
217        (14, 64) => Ok(size_of::<ConcurrentMerkleTree<14, 64>>()),
218        (14, 256) => Ok(size_of::<ConcurrentMerkleTree<14, 256>>()),
219        (14, 1024) => Ok(size_of::<ConcurrentMerkleTree<14, 1024>>()),
220        (14, 2048) => Ok(size_of::<ConcurrentMerkleTree<14, 2048>>()),
221        (15, 64) => Ok(size_of::<ConcurrentMerkleTree<15, 64>>()),
222        (16, 64) => Ok(size_of::<ConcurrentMerkleTree<16, 64>>()),
223        (17, 64) => Ok(size_of::<ConcurrentMerkleTree<17, 64>>()),
224        (18, 64) => Ok(size_of::<ConcurrentMerkleTree<18, 64>>()),
225        (19, 64) => Ok(size_of::<ConcurrentMerkleTree<19, 64>>()),
226        (20, 64) => Ok(size_of::<ConcurrentMerkleTree<20, 64>>()),
227        (20, 256) => Ok(size_of::<ConcurrentMerkleTree<20, 256>>()),
228        (20, 1024) => Ok(size_of::<ConcurrentMerkleTree<20, 1024>>()),
229        (20, 2048) => Ok(size_of::<ConcurrentMerkleTree<20, 2048>>()),
230        (24, 64) => Ok(size_of::<ConcurrentMerkleTree<24, 64>>()),
231        (24, 256) => Ok(size_of::<ConcurrentMerkleTree<24, 256>>()),
232        (24, 512) => Ok(size_of::<ConcurrentMerkleTree<24, 512>>()),
233        (24, 1024) => Ok(size_of::<ConcurrentMerkleTree<24, 1024>>()),
234        (24, 2048) => Ok(size_of::<ConcurrentMerkleTree<24, 2048>>()),
235        (26, 512) => Ok(size_of::<ConcurrentMerkleTree<26, 512>>()),
236        (26, 1024) => Ok(size_of::<ConcurrentMerkleTree<26, 1024>>()),
237        (26, 2048) => Ok(size_of::<ConcurrentMerkleTree<26, 2048>>()),
238        (30, 512) => Ok(size_of::<ConcurrentMerkleTree<30, 512>>()),
239        (30, 1024) => Ok(size_of::<ConcurrentMerkleTree<30, 1024>>()),
240        (30, 2048) => Ok(size_of::<ConcurrentMerkleTree<30, 2048>>()),
241        _ => {
242            msg!(
243                "Failed to get size of max depth {} and max buffer size {}",
244                header.get_max_depth(),
245                header.get_max_buffer_size()
246            );
247            err!(AccountCompressionError::ConcurrentMerkleTreeConstantsError)
248        }
249    }
250}