light_batched_merkle_tree/
initialize_state_tree.rs

1use borsh::{BorshDeserialize, BorshSerialize};
2use light_account_checks::{checks::check_account_balance_is_rent_exempt, AccountInfoTrait};
3use light_compressed_account::{pubkey::Pubkey, QueueType, TreeType};
4use light_merkle_tree_metadata::{
5    access::AccessMetadata, fee::compute_rollover_fee, merkle_tree::MerkleTreeMetadata,
6    queue::QueueMetadata, rollover::RolloverMetadata,
7};
8
9use crate::{
10    constants::{
11        DEFAULT_BATCH_SIZE, DEFAULT_BATCH_STATE_TREE_HEIGHT, DEFAULT_CPI_CONTEXT_ACCOUNT_SIZE,
12        DEFAULT_ZKP_BATCH_SIZE,
13    },
14    errors::BatchedMerkleTreeError,
15    merkle_tree::{get_merkle_tree_account_size, BatchedMerkleTreeAccount},
16    queue::{get_output_queue_account_size, BatchedQueueAccount},
17};
18
19#[repr(C)]
20#[derive(Debug, Clone, Copy, BorshDeserialize, BorshSerialize, PartialEq)]
21pub struct InitStateTreeAccountsInstructionData {
22    /// Unchecked identifier of the state tree.
23    pub index: u64,
24    /// Program owning the tree, enforced in the system program.
25    pub program_owner: Option<Pubkey>,
26    /// Optional forester pubkey for trees not forested
27    /// by light foresters enforced in registry program.
28    pub forester: Option<Pubkey>,
29    pub additional_bytes: u64,
30    pub input_queue_batch_size: u64,
31    pub output_queue_batch_size: u64,
32    pub input_queue_zkp_batch_size: u64,
33    pub output_queue_zkp_batch_size: u64,
34    pub bloom_filter_num_iters: u64,
35    pub bloom_filter_capacity: u64,
36    pub root_history_capacity: u32,
37    pub network_fee: Option<u64>,
38    pub rollover_threshold: Option<u64>,
39    /// Placeholder unimplemented.
40    pub close_threshold: Option<u64>,
41    pub height: u32,
42}
43
44impl Default for InitStateTreeAccountsInstructionData {
45    fn default() -> Self {
46        Self {
47            index: 0,
48            program_owner: None,
49            forester: None,
50            additional_bytes: DEFAULT_CPI_CONTEXT_ACCOUNT_SIZE,
51            bloom_filter_num_iters: 3,
52            input_queue_batch_size: DEFAULT_BATCH_SIZE,
53            output_queue_batch_size: DEFAULT_BATCH_SIZE,
54            input_queue_zkp_batch_size: DEFAULT_ZKP_BATCH_SIZE,
55            output_queue_zkp_batch_size: DEFAULT_ZKP_BATCH_SIZE,
56            height: DEFAULT_BATCH_STATE_TREE_HEIGHT,
57            root_history_capacity: (DEFAULT_BATCH_SIZE / DEFAULT_ZKP_BATCH_SIZE * 2) as u32,
58            bloom_filter_capacity: DEFAULT_BATCH_SIZE * 8,
59            network_fee: Some(5000),
60            rollover_threshold: Some(95),
61            close_threshold: None,
62        }
63    }
64}
65
66/// Initializes the state Merkle tree and output queue accounts.
67/// 1. Check rent exemption and that accounts are initialized with the correct size.
68/// 2. Initialize the output queue and state Merkle tree accounts.
69pub fn init_batched_state_merkle_tree_from_account_info<A: AccountInfoTrait>(
70    params: InitStateTreeAccountsInstructionData,
71    owner: Pubkey,
72    merkle_tree_account_info: &A,
73    queue_account_info: &A,
74    additional_bytes_rent: u64,
75) -> Result<(), BatchedMerkleTreeError> {
76    // 1. Check rent exemption and that accounts are initialized with the correct size.
77    let queue_rent;
78    let merkle_tree_rent;
79    {
80        let queue_account_size = get_output_queue_account_size(
81            params.output_queue_batch_size,
82            params.output_queue_zkp_batch_size,
83        );
84        let mt_account_size = get_merkle_tree_account_size(
85            params.input_queue_batch_size,
86            params.bloom_filter_capacity,
87            params.input_queue_zkp_batch_size,
88            params.root_history_capacity,
89            params.height,
90        );
91
92        queue_rent = check_account_balance_is_rent_exempt(queue_account_info, queue_account_size)?;
93
94        merkle_tree_rent =
95            check_account_balance_is_rent_exempt(merkle_tree_account_info, mt_account_size)?;
96    }
97
98    // 2. Initialize the output queue and state Merkle tree accounts.
99    let queue_data = &mut queue_account_info.try_borrow_mut_data()?;
100    let mt_data = &mut merkle_tree_account_info.try_borrow_mut_data()?;
101
102    init_batched_state_merkle_tree_accounts(
103        owner,
104        params,
105        queue_data,
106        queue_account_info.key().into(),
107        queue_rent,
108        mt_data,
109        merkle_tree_account_info.key().into(),
110        merkle_tree_rent,
111        additional_bytes_rent,
112    )?;
113    Ok(())
114}
115
116#[allow(clippy::too_many_arguments)]
117pub fn init_batched_state_merkle_tree_accounts<'a>(
118    owner: Pubkey,
119    params: InitStateTreeAccountsInstructionData,
120    output_queue_account_data: &mut [u8],
121    output_queue_pubkey: Pubkey,
122    queue_rent: u64,
123    mt_account_data: &'a mut [u8],
124    mt_pubkey: Pubkey,
125    merkle_tree_rent: u64,
126    additional_bytes_rent: u64,
127) -> Result<BatchedMerkleTreeAccount<'a>, BatchedMerkleTreeError> {
128    let height = params.height;
129    // Output queue
130    {
131        let rollover_fee = match params.rollover_threshold {
132            Some(rollover_threshold) => {
133                let rent = merkle_tree_rent + additional_bytes_rent + queue_rent;
134                compute_rollover_fee(rollover_threshold, height, rent)?
135            }
136            None => 0,
137        };
138
139        #[cfg(feature = "solana")]
140        solana_msg::msg!(" Output queue rollover_fee: {}", rollover_fee);
141        let metadata = QueueMetadata {
142            next_queue: Pubkey::default(),
143            access_metadata: AccessMetadata::new(owner, params.program_owner, params.forester),
144            rollover_metadata: RolloverMetadata::new(
145                params.index,
146                rollover_fee,
147                params.rollover_threshold,
148                params.network_fee.unwrap_or_default(),
149                params.close_threshold,
150                Some(params.additional_bytes),
151            ),
152            queue_type: QueueType::OutputStateV2 as u64,
153            associated_merkle_tree: mt_pubkey,
154        };
155
156        BatchedQueueAccount::init(
157            output_queue_account_data,
158            metadata,
159            params.output_queue_batch_size,
160            params.output_queue_zkp_batch_size,
161            // Output queues have no bloom filter.
162            0,
163            0,
164            output_queue_pubkey,
165            2u64.pow(params.height),
166        )?;
167    }
168    let metadata = MerkleTreeMetadata {
169        next_merkle_tree: Pubkey::default(),
170        access_metadata: AccessMetadata::new(owner, params.program_owner, params.forester),
171        rollover_metadata: RolloverMetadata::new(
172            params.index,
173            // The complete rollover fee is charged when creating an output
174            // compressed account by inserting it into the output queue.
175            0,
176            params.rollover_threshold,
177            params.network_fee.unwrap_or_default(),
178            params.close_threshold,
179            None,
180        ),
181        associated_queue: output_queue_pubkey,
182    };
183
184    // The state Merkle tree account includes the input queue.
185    // A nullifier is inserted when compressed state is spent.
186    // Spending compressed state requires proving its inclusion,
187    // which needs a root from the tree account.
188    BatchedMerkleTreeAccount::init(
189        mt_account_data,
190        &mt_pubkey,
191        metadata,
192        params.root_history_capacity,
193        params.input_queue_batch_size,
194        params.input_queue_zkp_batch_size,
195        height,
196        params.bloom_filter_num_iters,
197        params.bloom_filter_capacity,
198        TreeType::StateV2,
199    )
200}
201
202/// Only used for testing. For production use the default config.
203pub fn validate_batched_tree_params(params: InitStateTreeAccountsInstructionData) {
204    assert!(params.input_queue_batch_size > 0);
205    assert!(params.output_queue_batch_size > 0);
206    assert_eq!(
207        params.input_queue_batch_size % params.input_queue_zkp_batch_size,
208        0,
209        "Input queue batch size must divisible by input_queue_zkp_batch_size."
210    );
211    assert_eq!(
212        params.output_queue_batch_size % params.output_queue_zkp_batch_size,
213        0,
214        "Output queue batch size must divisible by output_queue_zkp_batch_size."
215    );
216    assert!(
217        match_circuit_size(params.input_queue_zkp_batch_size),
218        "Zkp batch size not supported. Supported 10, 500"
219    );
220    assert!(
221        match_circuit_size(params.output_queue_zkp_batch_size),
222        "Zkp batch size not supported. Supported 10, 500"
223    );
224
225    assert!(params.bloom_filter_num_iters > 0);
226    assert!(params.bloom_filter_capacity >= params.input_queue_batch_size * 8);
227    assert_eq!(
228        params.bloom_filter_capacity % 8,
229        0,
230        "Bloom filter capacity must be divisible by 8."
231    );
232    assert!(params.bloom_filter_capacity > 0);
233    assert!(params.root_history_capacity > 0);
234    assert!(params.input_queue_batch_size > 0);
235
236    // Validate root_history_capacity is sufficient for both input and output operations
237    let required_capacity = (params.output_queue_batch_size / params.output_queue_zkp_batch_size)
238        + (params.input_queue_batch_size / params.input_queue_zkp_batch_size);
239    assert!(
240        params.root_history_capacity >= required_capacity as u32,
241        "root_history_capacity ({}) must be >= {} (output_queue_batch_size / output_queue_zkp_batch_size + input_queue_batch_size / input_queue_zkp_batch_size)",
242        params.root_history_capacity,
243        required_capacity
244    );
245
246    assert_eq!(params.close_threshold, None);
247    assert_eq!(params.height, DEFAULT_BATCH_STATE_TREE_HEIGHT);
248}
249
250/// Only 10 and 500 are supported.
251pub fn match_circuit_size(size: u64) -> bool {
252    matches!(size, 10 | 500)
253}
254
255#[test]
256fn test_validate_root_history_capacity_state_tree() {
257    // Test with valid params (default should pass)
258    let params = InitStateTreeAccountsInstructionData::default();
259    validate_batched_tree_params(params); // Should not panic
260}
261
262#[test]
263#[should_panic(expected = "root_history_capacity")]
264fn test_validate_root_history_capacity_insufficient_state_tree() {
265    let params = InitStateTreeAccountsInstructionData {
266        root_history_capacity: 1, // Much too small
267        input_queue_batch_size: 1000,
268        output_queue_batch_size: 1000,
269        input_queue_zkp_batch_size: 10,
270        output_queue_zkp_batch_size: 10,
271        // Required: (1000/10) + (1000/10) = 200, but we set only 1
272        ..Default::default()
273    };
274    validate_batched_tree_params(params); // Should panic
275}
276#[cfg(feature = "test-only")]
277pub mod test_utils {
278    use light_compressed_account::hash_to_bn254_field_size_be;
279
280    pub use super::InitStateTreeAccountsInstructionData;
281    use super::*;
282    use crate::{
283        constants::{
284            DEFAULT_BATCH_ROOT_HISTORY_LEN, STATE_BLOOM_FILTER_CAPACITY,
285            STATE_BLOOM_FILTER_NUM_HASHES, TEST_DEFAULT_BATCH_SIZE, TEST_DEFAULT_ZKP_BATCH_SIZE,
286        },
287        queue::{test_utils::assert_queue_inited, BatchedQueueMetadata},
288        queue_batch_metadata::QueueBatches,
289    };
290
291    impl InitStateTreeAccountsInstructionData {
292        pub fn test_default() -> Self {
293            Self {
294                index: 0,
295                program_owner: None,
296                forester: None,
297                additional_bytes: DEFAULT_CPI_CONTEXT_ACCOUNT_SIZE,
298                bloom_filter_num_iters: 3,
299                input_queue_batch_size: TEST_DEFAULT_BATCH_SIZE,
300                output_queue_batch_size: TEST_DEFAULT_BATCH_SIZE,
301                input_queue_zkp_batch_size: TEST_DEFAULT_ZKP_BATCH_SIZE,
302                output_queue_zkp_batch_size: TEST_DEFAULT_ZKP_BATCH_SIZE,
303                height: DEFAULT_BATCH_STATE_TREE_HEIGHT,
304                root_history_capacity: DEFAULT_BATCH_ROOT_HISTORY_LEN,
305                bloom_filter_capacity: 20_000 * 8,
306                network_fee: Some(5000),
307                rollover_threshold: Some(95),
308                close_threshold: None,
309            }
310        }
311
312        pub fn e2e_test_default() -> Self {
313            Self {
314                index: 0,
315                program_owner: None,
316                forester: None,
317                additional_bytes: DEFAULT_CPI_CONTEXT_ACCOUNT_SIZE,
318                bloom_filter_num_iters: 3,
319                input_queue_batch_size: 500,
320                output_queue_batch_size: 500,
321                input_queue_zkp_batch_size: TEST_DEFAULT_ZKP_BATCH_SIZE,
322                output_queue_zkp_batch_size: TEST_DEFAULT_ZKP_BATCH_SIZE,
323                height: DEFAULT_BATCH_STATE_TREE_HEIGHT,
324                root_history_capacity: DEFAULT_BATCH_ROOT_HISTORY_LEN,
325                bloom_filter_capacity: 20_000 * 8,
326                network_fee: Some(5000),
327                rollover_threshold: Some(95),
328                close_threshold: None,
329            }
330        }
331
332        pub fn testnet_default() -> Self {
333            Self {
334                index: 0,
335                program_owner: None,
336                forester: None,
337                additional_bytes: DEFAULT_CPI_CONTEXT_ACCOUNT_SIZE,
338                bloom_filter_num_iters: STATE_BLOOM_FILTER_NUM_HASHES,
339                input_queue_batch_size: 15000,
340                output_queue_batch_size: 15000,
341                input_queue_zkp_batch_size: DEFAULT_ZKP_BATCH_SIZE,
342                output_queue_zkp_batch_size: DEFAULT_ZKP_BATCH_SIZE,
343                height: DEFAULT_BATCH_STATE_TREE_HEIGHT,
344                root_history_capacity: DEFAULT_BATCH_ROOT_HISTORY_LEN,
345                bloom_filter_capacity: STATE_BLOOM_FILTER_CAPACITY,
346                network_fee: Some(5000),
347                rollover_threshold: Some(95),
348                close_threshold: None,
349            }
350        }
351    }
352
353    pub fn get_state_merkle_tree_account_size_from_params(
354        params: InitStateTreeAccountsInstructionData,
355    ) -> usize {
356        crate::merkle_tree::get_merkle_tree_account_size(
357            params.input_queue_batch_size,
358            params.bloom_filter_capacity,
359            params.input_queue_zkp_batch_size,
360            params.root_history_capacity,
361            params.height,
362        )
363    }
364
365    pub fn assert_state_mt_zero_copy_initialized(
366        account_data: &mut [u8],
367        ref_account: crate::merkle_tree_metadata::BatchedMerkleTreeMetadata,
368        pubkey: &Pubkey,
369    ) {
370        let account = BatchedMerkleTreeAccount::state_from_bytes(account_data, pubkey)
371            .expect("from_bytes_unchecked_mut failed");
372        _assert_mt_zero_copy_initialized::<{ light_compressed_account::STATE_MERKLE_TREE_TYPE_V2 }>(
373            account,
374            ref_account,
375            TreeType::StateV2 as u64,
376        );
377    }
378
379    pub fn assert_address_mt_zero_copy_initialized(
380        account_data: &mut [u8],
381        ref_account: crate::merkle_tree_metadata::BatchedMerkleTreeMetadata,
382        pubkey: &Pubkey,
383    ) {
384        use crate::merkle_tree::BatchedMerkleTreeAccount;
385
386        let account = BatchedMerkleTreeAccount::address_from_bytes(account_data, pubkey)
387            .expect("from_bytes_unchecked_mut failed");
388        _assert_mt_zero_copy_initialized::<{ light_compressed_account::STATE_MERKLE_TREE_TYPE_V2 }>(
389            account,
390            ref_account,
391            TreeType::AddressV1 as u64,
392        );
393    }
394
395    fn _assert_mt_zero_copy_initialized<const TREE_TYPE: u64>(
396        account: BatchedMerkleTreeAccount,
397        ref_account: crate::merkle_tree_metadata::BatchedMerkleTreeMetadata,
398        tree_type: u64,
399    ) {
400        use light_compressed_account::QueueType;
401        use light_hasher::Hasher;
402
403        let queue = account.queue_batches;
404        let ref_queue = ref_account.queue_batches;
405        assert_eq!(*account, ref_account, "metadata mismatch");
406
407        assert_eq!(
408            account.root_history.capacity(),
409            ref_account.root_history_capacity as usize,
410            "root_history_capacity mismatch"
411        );
412        if tree_type == TreeType::StateV2 as u64 {
413            assert_eq!(
414                *account.root_history.get(0).unwrap(),
415                light_hasher::Poseidon::zero_bytes()[ref_account.height as usize],
416                "root_history not initialized"
417            );
418        }
419        if tree_type == TreeType::AddressV2 as u64 {
420            assert_eq!(
421                *account.root_history.get(0).unwrap(),
422                crate::constants::ADDRESS_TREE_INIT_ROOT_40,
423                "root_history not initialized"
424            );
425        }
426        assert_eq!(
427            account.hash_chain_stores[0].capacity(),
428            ref_account.queue_batches.get_num_zkp_batches() as usize,
429            "hash_chain_store mismatch"
430        );
431
432        let queue_type = if tree_type == TreeType::StateV2 as u64 {
433            QueueType::InputStateV2 as u64
434        } else {
435            QueueType::AddressV2 as u64
436        };
437        assert_queue_inited(queue, ref_queue, queue_type, &mut []);
438    }
439
440    #[derive(Debug, Clone, Copy)]
441    #[repr(C)]
442    pub struct CreateOutputQueueParams {
443        pub owner: Pubkey,
444        pub program_owner: Option<Pubkey>,
445        pub forester: Option<Pubkey>,
446        pub rollover_threshold: Option<u64>,
447        pub index: u64,
448        pub batch_size: u64,
449        pub zkp_batch_size: u64,
450        pub additional_bytes: u64,
451        pub rent: u64,
452        pub associated_merkle_tree: Pubkey,
453        pub queue_pubkey: Pubkey,
454        pub height: u32,
455        pub network_fee: u64,
456    }
457
458    impl CreateOutputQueueParams {
459        pub fn from(
460            params: InitStateTreeAccountsInstructionData,
461            owner: Pubkey,
462            rent: u64,
463            associated_merkle_tree: Pubkey,
464            queue_pubkey: Pubkey,
465        ) -> Self {
466            Self {
467                owner,
468                program_owner: params.program_owner,
469                forester: params.forester,
470                rollover_threshold: params.rollover_threshold,
471                index: params.index,
472                batch_size: params.output_queue_batch_size,
473                zkp_batch_size: params.output_queue_zkp_batch_size,
474                additional_bytes: params.additional_bytes,
475                rent,
476                associated_merkle_tree,
477                height: params.height,
478                network_fee: params.network_fee.unwrap_or_default(),
479                queue_pubkey,
480            }
481        }
482    }
483
484    pub fn create_output_queue_account(params: CreateOutputQueueParams) -> BatchedQueueMetadata {
485        let rollover_fee: u64 = match params.rollover_threshold {
486            Some(rollover_threshold) => {
487                compute_rollover_fee(rollover_threshold, params.height, params.rent).unwrap()
488            }
489            None => 0,
490        };
491        let metadata = QueueMetadata {
492            next_queue: Pubkey::default(),
493            access_metadata: AccessMetadata {
494                owner: params.owner,
495                program_owner: params.program_owner.unwrap_or_default(),
496                forester: params.forester.unwrap_or_default(),
497            },
498            rollover_metadata: RolloverMetadata {
499                close_threshold: u64::MAX,
500                index: params.index,
501                rolledover_slot: u64::MAX,
502                rollover_threshold: params.rollover_threshold.unwrap_or(u64::MAX),
503                rollover_fee,
504                network_fee: params.network_fee,
505                additional_bytes: params.additional_bytes,
506            },
507            queue_type: QueueType::OutputStateV2 as u64,
508            associated_merkle_tree: params.associated_merkle_tree,
509        };
510        let batch_metadata =
511            QueueBatches::new_output_queue(params.batch_size, params.zkp_batch_size).unwrap();
512        BatchedQueueMetadata {
513            metadata,
514            batch_metadata,
515            tree_capacity: 2u64.pow(params.height),
516            hashed_merkle_tree_pubkey: hash_to_bn254_field_size_be(
517                &params.associated_merkle_tree.to_bytes(),
518            ),
519            hashed_queue_pubkey: hash_to_bn254_field_size_be(&params.queue_pubkey.to_bytes()),
520        }
521    }
522}
light_batched_merkle_tree/initialize_state_tree.rs

light_batched_merkle_tree/
initialize_state_tree.rs
