diff --git a/program-libs/batched-merkle-tree/src/batch.rs b/program-libs/batched-merkle-tree/src/batch.rs index ddd968c5f7..319bf51d14 100644 --- a/program-libs/batched-merkle-tree/src/batch.rs +++ b/program-libs/batched-merkle-tree/src/batch.rs @@ -13,7 +13,7 @@ pub enum BatchState { Fill, /// Batch has been inserted into the tree. Inserted, - /// Batch is full, and insertion is in progress. + /// Batch is full. Full, } @@ -34,6 +34,14 @@ impl From for u64 { } } +/// Batch structure that holds +/// the metadata and state of a batch. +/// +/// A batch: +/// - has a size and a number of zkp batches. +/// - size must be divisible by zkp batch size. +/// - is part of a queue, by default a queue has two batches. +/// - is inserted into the tree by zkp batch. #[repr(C)] #[derive(Clone, Copy, Debug, PartialEq, Eq, KnownLayout, Immutable, IntoBytes, FromBytes)] pub struct Batch { @@ -203,15 +211,6 @@ impl Batch { .ok_or(BatchedMerkleTreeError::LeafIndexNotInBatch) } - pub fn store_value( - &mut self, - value: &[u8; 32], - value_store: &mut ZeroCopyVecU64<[u8; 32]>, - ) -> Result<(), BatchedMerkleTreeError> { - value_store.push(*value)?; - Ok(()) - } - /// Stores the value in a value store, /// and adds the value to the current hash chain. pub fn store_and_hash_value( @@ -221,12 +220,16 @@ impl Batch { hashchain_store: &mut ZeroCopyVecU64<[u8; 32]>, ) -> Result<(), BatchedMerkleTreeError> { self.add_to_hash_chain(value, hashchain_store)?; - self.store_value(value, value_store) + value_store.push(*value)?; + Ok(()) } - /// Inserts into the bloom filter and + /// Insert into the bloom filter and /// add value a the current hash chain. /// (used by nullifier & address queues) + /// 1. Add value to hash chain. + /// 2. Insert value into the bloom filter at bloom_filter_index. + /// 3. Check that value is not in any other bloom filter. pub fn insert( &mut self, bloom_filter_value: &[u8; 32], @@ -235,35 +238,52 @@ impl Batch { hashchain_store: &mut ZeroCopyVecU64<[u8; 32]>, bloom_filter_index: usize, ) -> Result<(), BatchedMerkleTreeError> { + // 1. add value to hash chain self.add_to_hash_chain(hashchain_value, hashchain_store)?; - - for (i, bloom_filter) in bloom_filter_stores.iter_mut().enumerate() { - if i == bloom_filter_index { - let mut bloom_filter = BloomFilter::new( - self.num_iters as usize, - self.bloom_filter_capacity, - bloom_filter.as_mut_slice(), - )?; - bloom_filter.insert(bloom_filter_value)?; - } else { - self.check_non_inclusion(bloom_filter_value, bloom_filter.as_mut_slice())?; + // insert into bloom filter & check non inclusion + { + let (before, after) = bloom_filter_stores.split_at_mut(bloom_filter_index); + let (bloom_filter, after) = after + .split_first_mut() + .ok_or(BatchedMerkleTreeError::InvalidIndex)?; + + // 2. Insert value into the bloom filter at bloom_filter_index. + BloomFilter::new( + self.num_iters as usize, + self.bloom_filter_capacity, + bloom_filter.as_mut_slice(), + )? + .insert(bloom_filter_value)?; + + // 3. Check that value is not in any other bloom filter. + for bf_store in before.iter_mut().chain(after.iter_mut()) { + self.check_non_inclusion(bloom_filter_value, bf_store.as_mut_slice())?; } } Ok(()) } + /// Add a value to the current hash chain, and advance batch state. + /// 1. Check that the batch is ready. + /// 2. If the zkp batch is empty, start a new hash chain. + /// 3. If the zkp batch is not empty, add value to last hash chain. + /// 4. If the zkp batch is full, increment the zkp batch index. + /// 5. If all zkp batches are full, set batch state to full. pub fn add_to_hash_chain( &mut self, value: &[u8; 32], hashchain_store: &mut ZeroCopyVecU64<[u8; 32]>, ) -> Result<(), BatchedMerkleTreeError> { + // 1. Check that the batch is ready. if self.get_state() != BatchState::Fill { return Err(BatchedMerkleTreeError::BatchNotReady); } let start_new_hash_chain = self.num_inserted == 0; if start_new_hash_chain { + // 2. Start a new hash chain. hashchain_store.push(*value)?; } else if let Some(last_hashchain) = hashchain_store.last() { + // 3. Add value to last hash chain. let hashchain = Poseidon::hashv(&[last_hashchain, value.as_slice()])?; *hashchain_store.last_mut().unwrap() = hashchain; } else { @@ -272,11 +292,13 @@ impl Batch { } self.num_inserted += 1; + // 4. If the zkp batch is full, increment the zkp batch index. let zkp_batch_is_full = self.num_inserted == self.zkp_batch_size; if zkp_batch_is_full { self.current_zkp_batch_index += 1; self.num_inserted = 0; + // 5. If all zkp batches are full, set batch state to full. let batch_is_full = self.current_zkp_batch_index == self.get_num_zkp_batches(); if batch_is_full { self.advance_state_to_full()?; @@ -295,9 +317,7 @@ impl Batch { let mut bloom_filter = BloomFilter::new(self.num_iters as usize, self.bloom_filter_capacity, store)?; if bloom_filter.contains(value) { - #[cfg(target_os = "solana")] - msg!("Value already exists in the bloom filter."); - return Err(BatchedMerkleTreeError::BatchInsertFailed); + return Err(BatchedMerkleTreeError::NonInclusionCheckFailed); } Ok(()) } diff --git a/program-libs/batched-merkle-tree/src/batch_metadata.rs b/program-libs/batched-merkle-tree/src/batch_metadata.rs index d86540fa06..f72e0e4462 100644 --- a/program-libs/batched-merkle-tree/src/batch_metadata.rs +++ b/program-libs/batched-merkle-tree/src/batch_metadata.rs @@ -1,7 +1,13 @@ use light_merkle_tree_metadata::{errors::MerkleTreeMetadataError, queue::QueueType}; +use light_zero_copy::{slice_mut::ZeroCopySliceMutU64, vec::ZeroCopyVecU64}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; -use crate::{batch::BatchState, errors::BatchedMerkleTreeError, BorshDeserialize, BorshSerialize}; +use crate::{ + batch::{Batch, BatchState}, + errors::BatchedMerkleTreeError, + queue::BatchedQueueMetadata, + BorshDeserialize, BorshSerialize, +}; #[repr(C)] #[derive( @@ -136,6 +142,33 @@ impl BatchMetadata { }; Ok((num_value_stores, num_stores, num_batches)) } + + pub fn queue_account_size(&self, queue_type: u64) -> Result { + let (num_value_vec, num_bloom_filter_stores, num_hashchain_store) = + self.get_size_parameters(queue_type)?; + let account_size = if queue_type != QueueType::BatchedOutput as u64 { + 0 + } else { + BatchedQueueMetadata::LEN + }; + let batches_size = + ZeroCopySliceMutU64::::required_size_for_capacity(self.num_batches); + let value_vecs_size = + ZeroCopyVecU64::<[u8; 32]>::required_size_for_capacity(self.batch_size) * num_value_vec; + // Bloomfilter capacity is in bits. + let bloom_filter_stores_size = + ZeroCopySliceMutU64::::required_size_for_capacity(self.bloom_filter_capacity / 8) + * num_bloom_filter_stores; + let hashchain_store_size = + ZeroCopyVecU64::<[u8; 32]>::required_size_for_capacity(self.get_num_zkp_batches()) + * num_hashchain_store; + let size = account_size + + batches_size + + value_vecs_size + + bloom_filter_stores_size + + hashchain_store_size; + Ok(size) + } } #[test] diff --git a/program-libs/batched-merkle-tree/src/errors.rs b/program-libs/batched-merkle-tree/src/errors.rs index 697771b473..09aa168012 100644 --- a/program-libs/batched-merkle-tree/src/errors.rs +++ b/program-libs/batched-merkle-tree/src/errors.rs @@ -45,6 +45,8 @@ pub enum BatchedMerkleTreeError { InvalidIndex, #[error("Batched Merkle tree is full.")] TreeIsFull, + #[error("Value already exists in bloom filter.")] + NonInclusionCheckFailed, } #[cfg(feature = "solana")] @@ -62,6 +64,7 @@ impl From for u32 { BatchedMerkleTreeError::InvalidBatchIndex => 14309, BatchedMerkleTreeError::InvalidIndex => 14310, BatchedMerkleTreeError::TreeIsFull => 14311, + BatchedMerkleTreeError::NonInclusionCheckFailed => 14312, BatchedMerkleTreeError::Hasher(e) => e.into(), BatchedMerkleTreeError::ZeroCopy(e) => e.into(), BatchedMerkleTreeError::MerkleTreeMetadata(e) => e.into(), diff --git a/program-libs/batched-merkle-tree/src/initialize_state_tree.rs b/program-libs/batched-merkle-tree/src/initialize_state_tree.rs index fd16160269..8f4accca38 100644 --- a/program-libs/batched-merkle-tree/src/initialize_state_tree.rs +++ b/program-libs/batched-merkle-tree/src/initialize_state_tree.rs @@ -311,7 +311,7 @@ pub fn get_state_merkle_tree_account_size_from_params( #[cfg(not(target_os = "solana"))] pub fn assert_state_mt_zero_copy_inited( account_data: &mut [u8], - ref_account: crate::merkle_tree::BatchedMerkleTreeMetadata, + ref_account: crate::merkle_tree_metadata::BatchedMerkleTreeMetadata, num_iters: u64, ) { let account = BatchedMerkleTreeAccount::state_from_bytes(account_data) @@ -327,7 +327,7 @@ pub fn assert_state_mt_zero_copy_inited( #[cfg(not(target_os = "solana"))] pub fn assert_address_mt_zero_copy_inited( account_data: &mut [u8], - ref_account: crate::merkle_tree::BatchedMerkleTreeMetadata, + ref_account: crate::merkle_tree_metadata::BatchedMerkleTreeMetadata, num_iters: u64, ) { use crate::{constants::BATCHED_ADDRESS_TREE_TYPE, merkle_tree::BatchedMerkleTreeAccount}; @@ -345,7 +345,7 @@ pub fn assert_address_mt_zero_copy_inited( #[cfg(not(target_os = "solana"))] fn _assert_mt_zero_copy_inited( mut account: BatchedMerkleTreeAccount, - ref_account: crate::merkle_tree::BatchedMerkleTreeMetadata, + ref_account: crate::merkle_tree_metadata::BatchedMerkleTreeMetadata, num_iters: u64, tree_type: u64, ) { diff --git a/program-libs/batched-merkle-tree/src/lib.rs b/program-libs/batched-merkle-tree/src/lib.rs index 1bbbf4393c..9da881e64f 100644 --- a/program-libs/batched-merkle-tree/src/lib.rs +++ b/program-libs/batched-merkle-tree/src/lib.rs @@ -7,6 +7,7 @@ pub mod event; pub mod initialize_address_tree; pub mod initialize_state_tree; pub mod merkle_tree; +pub mod merkle_tree_metadata; pub mod queue; pub mod rollover_address_tree; pub mod rollover_state_tree; diff --git a/program-libs/batched-merkle-tree/src/merkle_tree.rs b/program-libs/batched-merkle-tree/src/merkle_tree.rs index 948caec2de..3fb15e7405 100644 --- a/program-libs/batched-merkle-tree/src/merkle_tree.rs +++ b/program-libs/batched-merkle-tree/src/merkle_tree.rs @@ -1,17 +1,13 @@ use std::ops::{Deref, DerefMut}; -use aligned_sized::aligned_sized; use light_hasher::{Discriminator, Hasher, Poseidon}; use light_merkle_tree_metadata::{ - access::AccessMetadata, errors::MerkleTreeMetadataError, merkle_tree::{MerkleTreeMetadata, TreeType}, queue::QueueType, - rollover::RolloverMetadata, }; use light_utils::{ - account::{check_account_info_mut, check_discriminator, set_discriminator, DISCRIMINATOR_LEN}, - fee::compute_rollover_fee, + account::{check_account_info, check_discriminator, set_discriminator, DISCRIMINATOR_LEN}, hashchain::create_hash_chain_from_array, pubkey::Pubkey, }; @@ -24,242 +20,26 @@ use light_zero_copy::{ vec::ZeroCopyVecU64, }; use solana_program::{account_info::AccountInfo, msg}; -use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Ref}; +use zerocopy::Ref; use super::{ batch::Batch, - queue::{init_queue, input_queue_from_bytes, insert_into_current_batch, queue_account_size}, + queue::{init_queue, input_queue_from_bytes, insert_into_current_batch}, }; use crate::{ batch::BatchState, batch_metadata::BatchMetadata, constants::{ ACCOUNT_COMPRESSION_PROGRAM_ID, ADDRESS_TREE_INIT_ROOT_40, BATCHED_ADDRESS_TREE_TYPE, - BATCHED_STATE_TREE_TYPE, DEFAULT_BATCH_STATE_TREE_HEIGHT, TEST_DEFAULT_BATCH_SIZE, + BATCHED_STATE_TREE_TYPE, }, errors::BatchedMerkleTreeError, event::{BatchAddressAppendEvent, BatchAppendEvent, BatchNullifyEvent}, - initialize_address_tree::InitAddressTreeAccountsInstructionData, - initialize_state_tree::InitStateTreeAccountsInstructionData, + merkle_tree_metadata::BatchedMerkleTreeMetadata, queue::BatchedQueueAccount, BorshDeserialize, BorshSerialize, }; -#[repr(C)] -#[derive( - BorshSerialize, - BorshDeserialize, - Debug, - PartialEq, - Clone, - Copy, - FromBytes, - IntoBytes, - KnownLayout, - Immutable, -)] -#[aligned_sized(anchor)] -pub struct BatchedMerkleTreeMetadata { - pub metadata: MerkleTreeMetadata, - pub sequence_number: u64, - pub tree_type: u64, - pub next_index: u64, - pub height: u32, - pub root_history_capacity: u32, - pub capacity: u64, - pub queue_metadata: BatchMetadata, -} - -// TODO: make anchor consistent -impl Discriminator for BatchedMerkleTreeAccount<'_> { - const DISCRIMINATOR: [u8; 8] = *b"BatchMka"; -} - -impl Default for BatchedMerkleTreeMetadata { - fn default() -> Self { - BatchedMerkleTreeMetadata { - metadata: MerkleTreeMetadata::default(), - next_index: 0, - sequence_number: 0, - tree_type: TreeType::BatchedState as u64, - height: DEFAULT_BATCH_STATE_TREE_HEIGHT, - root_history_capacity: 20, - capacity: 2u64.pow(DEFAULT_BATCH_STATE_TREE_HEIGHT), - queue_metadata: BatchMetadata { - currently_processing_batch_index: 0, - num_batches: 2, - batch_size: TEST_DEFAULT_BATCH_SIZE, - bloom_filter_capacity: 20_000 * 8, - zkp_batch_size: 10, - ..Default::default() - }, - } - } -} - -#[repr(C)] -pub struct CreateTreeParams { - pub owner: Pubkey, - pub program_owner: Option, - pub forester: Option, - pub rollover_threshold: Option, - pub index: u64, - pub network_fee: u64, - pub batch_size: u64, - pub zkp_batch_size: u64, - pub bloom_filter_capacity: u64, - pub root_history_capacity: u32, - pub height: u32, - pub num_batches: u64, -} -impl CreateTreeParams { - pub fn from_state_ix_params(data: InitStateTreeAccountsInstructionData, owner: Pubkey) -> Self { - CreateTreeParams { - owner, - program_owner: data.program_owner, - forester: data.forester, - rollover_threshold: data.rollover_threshold, - index: data.index, - network_fee: data.network_fee.unwrap_or(0), - batch_size: data.input_queue_batch_size, - zkp_batch_size: data.input_queue_zkp_batch_size, - bloom_filter_capacity: data.bloom_filter_capacity, - root_history_capacity: data.root_history_capacity, - height: data.height, - num_batches: data.input_queue_num_batches, - } - } - - pub fn from_address_ix_params( - data: InitAddressTreeAccountsInstructionData, - owner: Pubkey, - ) -> Self { - CreateTreeParams { - owner, - program_owner: data.program_owner, - forester: data.forester, - rollover_threshold: data.rollover_threshold, - index: data.index, - network_fee: data.network_fee.unwrap_or(0), - batch_size: data.input_queue_batch_size, - zkp_batch_size: data.input_queue_zkp_batch_size, - bloom_filter_capacity: data.bloom_filter_capacity, - root_history_capacity: data.root_history_capacity, - height: data.height, - num_batches: data.input_queue_num_batches, - } - } -} - -impl BatchedMerkleTreeMetadata { - pub fn get_account_size(&self) -> Result { - let account_size = Self::LEN; - let root_history_size = ZeroCopyCyclicVecU64::<[u8; 32]>::required_size_for_capacity( - self.root_history_capacity as u64, - ); - let size = account_size - + root_history_size - + queue_account_size(&self.queue_metadata, QueueType::BatchedInput as u64)?; - Ok(size) - } - - pub fn new_state_tree(params: CreateTreeParams, associated_queue: Pubkey) -> Self { - Self::new_tree(TreeType::BatchedState, params, associated_queue, 0) - } - - pub fn new_address_tree(params: CreateTreeParams, rent: u64) -> Self { - let rollover_fee = match params.rollover_threshold { - Some(rollover_threshold) => { - compute_rollover_fee(rollover_threshold, params.height, rent).unwrap() - } - None => 0, - }; - let mut tree = Self::new_tree( - TreeType::BatchedAddress, - params, - Pubkey::default(), - rollover_fee, - ); - // inited address tree contains two elements. - tree.next_index = 2; - tree - } - - fn new_tree( - tree_type: TreeType, - params: CreateTreeParams, - associated_queue: Pubkey, - rollover_fee: u64, - ) -> Self { - let CreateTreeParams { - owner, - program_owner, - forester, - rollover_threshold, - index, - network_fee, - batch_size, - zkp_batch_size, - bloom_filter_capacity, - root_history_capacity, - height, - num_batches, - } = params; - Self { - metadata: MerkleTreeMetadata { - next_merkle_tree: Pubkey::default(), - access_metadata: AccessMetadata::new(owner, program_owner, forester), - rollover_metadata: RolloverMetadata::new( - index, - rollover_fee, - rollover_threshold, - network_fee, - None, - None, - ), - associated_queue, - }, - sequence_number: 0, - tree_type: tree_type as u64, - next_index: 0, - height, - root_history_capacity, - queue_metadata: BatchMetadata::new_input_queue( - batch_size, - bloom_filter_capacity, - zkp_batch_size, - num_batches, - ) - .unwrap(), - capacity: 2u64.pow(height), - } - } -} - -#[derive(Debug, PartialEq)] -pub struct BatchedMerkleTreeAccount<'a> { - metadata: Ref<&'a mut [u8], BatchedMerkleTreeMetadata>, - pub root_history: ZeroCopyCyclicVecU64<'a, [u8; 32]>, - pub batches: ZeroCopySliceMutU64<'a, Batch>, - pub value_vecs: Vec>, - pub bloom_filter_stores: Vec>, - pub hashchain_store: Vec>, -} - -impl Deref for BatchedMerkleTreeAccount<'_> { - type Target = BatchedMerkleTreeMetadata; - - fn deref(&self) -> &Self::Target { - &self.metadata - } -} - -impl DerefMut for BatchedMerkleTreeAccount<'_> { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.metadata - } -} - /// Public inputs: /// 1. old root (last root in root history) /// 2. new root (send to chain) @@ -290,16 +70,60 @@ pub struct InstructionDataBatchAppendInputs { pub compressed_proof: CompressedProof, } -impl<'a> BatchedMerkleTreeAccount<'a> { - pub fn get_metadata(&self) -> &BatchedMerkleTreeMetadata { - &self.metadata - } +/// Batched Merkle tree zero copy account. +/// The account is used to batched state +/// and address Merkle trees, plus the input and address queues, +/// in the Light Protocol account compression program. +/// +/// Tree roots can be used in zk proofs +/// outside of Light Protocol programs. +/// +/// To access a tree root by index use: +/// - get_state_root_by_index +/// - get_address_root_by_index +#[derive(Debug, PartialEq)] +pub struct BatchedMerkleTreeAccount<'a> { + metadata: Ref<&'a mut [u8], BatchedMerkleTreeMetadata>, + pub root_history: ZeroCopyCyclicVecU64<'a, [u8; 32]>, + pub batches: ZeroCopySliceMutU64<'a, Batch>, + pub value_vecs: Vec>, + pub bloom_filter_stores: Vec>, + pub hashchain_store: Vec>, +} - pub fn get_metadata_mut(&mut self) -> &mut BatchedMerkleTreeMetadata { - &mut self.metadata +impl Discriminator for BatchedMerkleTreeAccount<'_> { + const DISCRIMINATOR: [u8; 8] = *b"BatchMta"; +} + +impl<'a> BatchedMerkleTreeAccount<'a> { + /// Checks state Merkle tree account and returns the root. + pub fn get_state_root_by_index( + account_info: &AccountInfo<'a>, + index: usize, + ) -> Result<[u8; 32], BatchedMerkleTreeError> { + let tree = Self::state_from_account_info(account_info)?; + Ok(*tree + .get_root_by_index(index) + .ok_or(BatchedMerkleTreeError::InvalidIndex)?) } - // // TODO: add unit test + /// Checks address Merkle tree account and returns the root. + pub fn get_address_root_by_index( + account_info: &AccountInfo<'a>, + index: usize, + ) -> Result<[u8; 32], BatchedMerkleTreeError> { + let tree = Self::address_from_account_info(account_info)?; + Ok(*tree + .get_root_by_index(index) + .ok_or(BatchedMerkleTreeError::InvalidIndex)?) + } + + /// Deserialize a batched state Merkle tree from account info. + /// Should be used in solana programs. + /// Checks that: + /// 1. the program owner is the light account compression program, + /// 2. discriminator, + /// 3. tree type is batched state tree type. pub fn state_from_account_info( account_info: &AccountInfo<'a>, ) -> Result, BatchedMerkleTreeError> { @@ -309,7 +133,9 @@ impl<'a> BatchedMerkleTreeAccount<'a> { ) } - // TODO: add failing test + /// Deserialize a state BatchedMerkleTreeAccount from bytes. + /// Should only be used in client. + /// Checks the discriminator and tree type. #[cfg(not(target_os = "solana"))] pub fn state_from_bytes( account_data: &'a mut [u8], @@ -317,6 +143,12 @@ impl<'a> BatchedMerkleTreeAccount<'a> { Self::from_bytes::(account_data) } + /// Deserialize a batched address Merkle tree from account info. + /// Should be used in solana programs. + /// Checks that: + /// 1. the program owner is the light account compression program, + /// 2. discriminator, + /// 3. tree type is batched address tree type. pub fn address_from_account_info( account_info: &AccountInfo<'a>, ) -> Result, BatchedMerkleTreeError> { @@ -326,11 +158,11 @@ impl<'a> BatchedMerkleTreeAccount<'a> { ) } - pub fn from_account_info( + fn from_account_info( program_id: &solana_program::pubkey::Pubkey, account_info: &AccountInfo<'a>, ) -> Result, BatchedMerkleTreeError> { - check_account_info_mut::(program_id, account_info)?; + check_account_info::(program_id, account_info)?; let mut data = account_info.try_borrow_mut_data()?; // Necessary to convince the borrow checker. @@ -339,6 +171,9 @@ impl<'a> BatchedMerkleTreeAccount<'a> { Self::from_bytes::(data_slice) } + /// Deserialize a state BatchedMerkleTreeAccount from bytes. + /// Should only be used in client. + /// Checks the discriminator and tree type. #[cfg(not(target_os = "solana"))] pub fn address_from_bytes( account_data: &'a mut [u8], @@ -424,7 +259,11 @@ impl<'a> BatchedMerkleTreeAccount<'a> { account_metadata.root_history_capacity as u64, account_data, )?; - // fill root history with zero bytes + + // Initialize root history with zero bytes to enable + // unified logic to zero out roots. + // An unitialized root history vector + // would be an edge case. for _ in 0..root_history.capacity() { root_history.push([0u8; 32]); } @@ -838,7 +677,9 @@ impl<'a> BatchedMerkleTreeAccount<'a> { Ok(()) } - /// Zero out roots corresponding to sequence numbers < sequence_number. + /// Zero out roots corresponding to batch.sequence numbers > tree.sequence_number. + /// batch.sequence_number marks the sequence number all roots are overwritten + /// which can prove inclusion of a value inserted in the queue. /// 1. Check whether overlapping roots exist. /// 2. If yes: /// 2.1 Get, first safe root index. @@ -847,8 +688,6 @@ impl<'a> BatchedMerkleTreeAccount<'a> { // 1. Check whether overlapping roots exist. let overlapping_roots_exits = sequence_number > self.sequence_number; if overlapping_roots_exits { - // let root_index = root_index as usize; - let mut oldest_root_index = self.root_history.first_index(); // 2.1. Get, num of remaining roots. // Remaining roots have not been updated since @@ -899,7 +738,7 @@ impl<'a> BatchedMerkleTreeAccount<'a> { /// 3.1 zero out bloom filter /// 3.2 mark bloom filter as zeroed /// 3.3 zero out roots if needed - pub fn zero_out_previous_batch_bloom_filter(&mut self) -> Result<(), BatchedMerkleTreeError> { + fn zero_out_previous_batch_bloom_filter(&mut self) -> Result<(), BatchedMerkleTreeError> { let current_batch = self.queue_metadata.next_full_batch_index as usize; let batch_size = self.queue_metadata.batch_size; let previous_full_batch_index = current_batch.saturating_sub(1); @@ -968,10 +807,23 @@ impl<'a> BatchedMerkleTreeAccount<'a> { pub fn get_root_index(&self) -> u32 { self.root_history.last_index() as u32 } + pub fn get_root(&self) -> Option<[u8; 32]> { self.root_history.last().copied() } + pub fn get_root_by_index(&self, index: usize) -> Option<&[u8; 32]> { + self.root_history.get(index) + } + + pub fn get_metadata(&self) -> &BatchedMerkleTreeMetadata { + &self.metadata + } + + pub fn get_metadata_mut(&mut self) -> &mut BatchedMerkleTreeMetadata { + &mut self.metadata + } + // TODO: add unit test /// Checks non-inclusion in all bloom filters /// which are not zeroed. @@ -990,16 +842,16 @@ impl<'a> BatchedMerkleTreeAccount<'a> { Ok(()) } + pub fn tree_is_full(&self) -> bool { + self.next_index == self.capacity + } + pub fn check_tree_is_full(&self) -> Result<(), BatchedMerkleTreeError> { if self.tree_is_full() { return Err(BatchedMerkleTreeError::TreeIsFull); } Ok(()) } - - pub fn tree_is_full(&self) -> bool { - self.next_index == self.capacity - } } pub fn get_merkle_tree_account_size_default() -> usize { @@ -1007,6 +859,20 @@ pub fn get_merkle_tree_account_size_default() -> usize { mt_account.get_account_size().unwrap() } +impl Deref for BatchedMerkleTreeAccount<'_> { + type Target = BatchedMerkleTreeMetadata; + + fn deref(&self) -> &Self::Target { + &self.metadata + } +} + +impl DerefMut for BatchedMerkleTreeAccount<'_> { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.metadata + } +} + pub fn get_merkle_tree_account_size( batch_size: u64, bloom_filter_capacity: u64, diff --git a/program-libs/batched-merkle-tree/src/merkle_tree_metadata.rs b/program-libs/batched-merkle-tree/src/merkle_tree_metadata.rs new file mode 100644 index 0000000000..08429a2fb0 --- /dev/null +++ b/program-libs/batched-merkle-tree/src/merkle_tree_metadata.rs @@ -0,0 +1,207 @@ +use aligned_sized::aligned_sized; +use light_merkle_tree_metadata::{ + access::AccessMetadata, + merkle_tree::{MerkleTreeMetadata, TreeType}, + queue::QueueType, + rollover::RolloverMetadata, +}; +use light_utils::{fee::compute_rollover_fee, pubkey::Pubkey}; +use light_zero_copy::cyclic_vec::ZeroCopyCyclicVecU64; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +use crate::{ + batch_metadata::BatchMetadata, + constants::{DEFAULT_BATCH_STATE_TREE_HEIGHT, TEST_DEFAULT_BATCH_SIZE}, + errors::BatchedMerkleTreeError, + initialize_address_tree::InitAddressTreeAccountsInstructionData, + initialize_state_tree::InitStateTreeAccountsInstructionData, + BorshDeserialize, BorshSerialize, +}; + +#[repr(C)] +#[derive( + BorshSerialize, + BorshDeserialize, + Debug, + PartialEq, + Clone, + Copy, + FromBytes, + IntoBytes, + KnownLayout, + Immutable, +)] +#[aligned_sized(anchor)] +pub struct BatchedMerkleTreeMetadata { + pub metadata: MerkleTreeMetadata, + pub sequence_number: u64, + pub tree_type: u64, + pub next_index: u64, + pub height: u32, + pub root_history_capacity: u32, + pub capacity: u64, + pub queue_metadata: BatchMetadata, +} + +impl Default for BatchedMerkleTreeMetadata { + fn default() -> Self { + BatchedMerkleTreeMetadata { + metadata: MerkleTreeMetadata::default(), + next_index: 0, + sequence_number: 0, + tree_type: TreeType::BatchedState as u64, + height: DEFAULT_BATCH_STATE_TREE_HEIGHT, + root_history_capacity: 20, + capacity: 2u64.pow(DEFAULT_BATCH_STATE_TREE_HEIGHT), + queue_metadata: BatchMetadata { + currently_processing_batch_index: 0, + num_batches: 2, + batch_size: TEST_DEFAULT_BATCH_SIZE, + bloom_filter_capacity: 20_000 * 8, + zkp_batch_size: 10, + ..Default::default() + }, + } + } +} + +impl BatchedMerkleTreeMetadata { + pub fn get_account_size(&self) -> Result { + let account_size = Self::LEN; + let root_history_size = ZeroCopyCyclicVecU64::<[u8; 32]>::required_size_for_capacity( + self.root_history_capacity as u64, + ); + let size = account_size + + root_history_size + + self + .queue_metadata + .queue_account_size(QueueType::BatchedInput as u64)?; + Ok(size) + } + + pub fn new_state_tree(params: CreateTreeParams, associated_queue: Pubkey) -> Self { + Self::new_tree(TreeType::BatchedState, params, associated_queue, 0) + } + + pub fn new_address_tree(params: CreateTreeParams, rent: u64) -> Self { + let rollover_fee = match params.rollover_threshold { + Some(rollover_threshold) => { + compute_rollover_fee(rollover_threshold, params.height, rent).unwrap() + } + None => 0, + }; + let mut tree = Self::new_tree( + TreeType::BatchedAddress, + params, + Pubkey::default(), + rollover_fee, + ); + // inited address tree contains two elements. + tree.next_index = 2; + tree + } + + fn new_tree( + tree_type: TreeType, + params: CreateTreeParams, + associated_queue: Pubkey, + rollover_fee: u64, + ) -> Self { + let CreateTreeParams { + owner, + program_owner, + forester, + rollover_threshold, + index, + network_fee, + batch_size, + zkp_batch_size, + bloom_filter_capacity, + root_history_capacity, + height, + num_batches, + } = params; + Self { + metadata: MerkleTreeMetadata { + next_merkle_tree: Pubkey::default(), + access_metadata: AccessMetadata::new(owner, program_owner, forester), + rollover_metadata: RolloverMetadata::new( + index, + rollover_fee, + rollover_threshold, + network_fee, + None, + None, + ), + associated_queue, + }, + sequence_number: 0, + tree_type: tree_type as u64, + next_index: 0, + height, + root_history_capacity, + queue_metadata: BatchMetadata::new_input_queue( + batch_size, + bloom_filter_capacity, + zkp_batch_size, + num_batches, + ) + .unwrap(), + capacity: 2u64.pow(height), + } + } +} + +#[repr(C)] +pub struct CreateTreeParams { + pub owner: Pubkey, + pub program_owner: Option, + pub forester: Option, + pub rollover_threshold: Option, + pub index: u64, + pub network_fee: u64, + pub batch_size: u64, + pub zkp_batch_size: u64, + pub bloom_filter_capacity: u64, + pub root_history_capacity: u32, + pub height: u32, + pub num_batches: u64, +} +impl CreateTreeParams { + pub fn from_state_ix_params(data: InitStateTreeAccountsInstructionData, owner: Pubkey) -> Self { + CreateTreeParams { + owner, + program_owner: data.program_owner, + forester: data.forester, + rollover_threshold: data.rollover_threshold, + index: data.index, + network_fee: data.network_fee.unwrap_or(0), + batch_size: data.input_queue_batch_size, + zkp_batch_size: data.input_queue_zkp_batch_size, + bloom_filter_capacity: data.bloom_filter_capacity, + root_history_capacity: data.root_history_capacity, + height: data.height, + num_batches: data.input_queue_num_batches, + } + } + + pub fn from_address_ix_params( + data: InitAddressTreeAccountsInstructionData, + owner: Pubkey, + ) -> Self { + CreateTreeParams { + owner, + program_owner: data.program_owner, + forester: data.forester, + rollover_threshold: data.rollover_threshold, + index: data.index, + network_fee: data.network_fee.unwrap_or(0), + batch_size: data.input_queue_batch_size, + zkp_batch_size: data.input_queue_zkp_batch_size, + bloom_filter_capacity: data.bloom_filter_capacity, + root_history_capacity: data.root_history_capacity, + height: data.height, + num_batches: data.input_queue_num_batches, + } + } +} diff --git a/program-libs/batched-merkle-tree/src/queue.rs b/program-libs/batched-merkle-tree/src/queue.rs index 6db2405563..fcaea2091f 100644 --- a/program-libs/batched-merkle-tree/src/queue.rs +++ b/program-libs/batched-merkle-tree/src/queue.rs @@ -7,7 +7,7 @@ use light_merkle_tree_metadata::{ queue::{QueueMetadata, QueueType}, }; use light_utils::{ - account::{check_account_info_mut, check_discriminator, set_discriminator, DISCRIMINATOR_LEN}, + account::{check_account_info, check_discriminator, set_discriminator, DISCRIMINATOR_LEN}, pubkey::Pubkey, }; use light_zero_copy::{errors::ZeroCopyError, slice_mut::ZeroCopySliceMutU64, vec::ZeroCopyVecU64}; @@ -51,10 +51,6 @@ pub struct BatchedQueueMetadata { pub tree_capacity: u64, } -impl Discriminator for BatchedQueueAccount<'_> { - const DISCRIMINATOR: [u8; 8] = *b"queueacc"; -} - impl BatchedQueueMetadata { pub fn get_size_parameters(&self) -> Result<(usize, usize, usize), MerkleTreeMetadataError> { self.batch_metadata @@ -76,81 +72,82 @@ impl BatchedQueueMetadata { } } -pub fn queue_account_size( - batch_metadata: &BatchMetadata, - queue_type: u64, -) -> Result { - let (num_value_vec, num_bloom_filter_stores, num_hashchain_store) = - batch_metadata.get_size_parameters(queue_type)?; - let account_size = if queue_type != QueueType::BatchedOutput as u64 { - 0 - } else { - BatchedQueueMetadata::LEN - }; - let batches_size = - ZeroCopySliceMutU64::::required_size_for_capacity(batch_metadata.num_batches); - let value_vecs_size = - ZeroCopyVecU64::<[u8; 32]>::required_size_for_capacity(batch_metadata.batch_size) - * num_value_vec; - // Bloomfilter capacity is in bits. - let bloom_filter_stores_size = ZeroCopySliceMutU64::::required_size_for_capacity( - batch_metadata.bloom_filter_capacity / 8, - ) * num_bloom_filter_stores; - let hashchain_store_size = ZeroCopyVecU64::<[u8; 32]>::required_size_for_capacity( - batch_metadata.get_num_zkp_batches(), - ) * num_hashchain_store; - let size = account_size - + batches_size - + value_vecs_size - + bloom_filter_stores_size - + hashchain_store_size; - Ok(size) -} - +/// Batched queue zero copy account. +/// Used for output queues in light protocol. +/// Output queues store compressed account hashes, +/// to be appended to a batched Merkle tree +/// in batches with a zero-knowledge proof (zkp), +/// ie. it stores hashes and commits these to hash chains. +/// Each hash chain is used as public input for +/// a batch append zkp. +/// +/// An output queue is configured with: +/// 1. N batches +/// 2. N value vecs (one for each batch) +/// value vec length = batch size +/// 3. N hashchain stores (one for each batch) +/// hashchain store length = batch size /zkp batch size +/// +/// Default config: +/// - 2 batches +/// - 50,000 batch size +/// - 500 zkp batch size +/// +/// Initialization: +/// - an output queue is initialized +/// in combination with a state Merkle tree +/// - `init_batched_state_merkle_tree_from_account_info` +/// +/// For deserialization use: +/// - in program: `output_from_account_info` +/// - in client: `output_from_bytes` +/// +/// To insert a value use: +/// - `insert_into_current_batch` +/// +/// A compressed account can be spent or read +/// while in the output queue. +/// +/// To spend use: +/// - check_leaf_index_could_exist_in_batches in combination with +/// `prove_inclusion_by_index_and_zero_out_leaf` +/// +/// To read use: +/// - `prove_inclusion_by_index` #[derive(Debug, PartialEq)] pub struct BatchedQueueAccount<'a> { metadata: Ref<&'a mut [u8], BatchedQueueMetadata>, pub batches: ZeroCopySliceMutU64<'a, Batch>, pub value_vecs: Vec>, pub bloom_filter_stores: Vec>, - /// hashchain_store_capacity = batch_capacity / zkp_batch_size pub hashchain_store: Vec>, } -impl Deref for BatchedQueueAccount<'_> { - type Target = BatchedQueueMetadata; - - fn deref(&self) -> &Self::Target { - &self.metadata - } -} - -impl DerefMut for BatchedQueueAccount<'_> { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.metadata - } +impl Discriminator for BatchedQueueAccount<'_> { + const DISCRIMINATOR: [u8; 8] = *b"queueacc"; } impl<'a> BatchedQueueAccount<'a> { - pub fn get_metadata(&self) -> &BatchedQueueMetadata { - &self.metadata - } - - pub fn get_metadata_mut(&mut self) -> &mut BatchedQueueMetadata { - &mut self.metadata - } - + /// Deserialize an output BatchedQueueAccount from account info. + /// Should be used in solana programs. + /// Checks that: + /// 1. the program owner is the light account compression program, + /// 2. discriminator, + /// 3. queue type is output queue type. pub fn output_from_account_info( account_info: &AccountInfo<'a>, ) -> Result, BatchedMerkleTreeError> { Self::from_account_info::(&ACCOUNT_COMPRESSION_PROGRAM_ID, account_info) } - pub fn from_account_info( + /// Deserialize a BatchedQueueAccount from account info. + /// Should be used in solana programs. + /// Checks the program owner, discriminator and queue type. + fn from_account_info( program_id: &solana_program::pubkey::Pubkey, account_info: &AccountInfo<'a>, ) -> Result, BatchedMerkleTreeError> { - check_account_info_mut::(program_id, account_info)?; + check_account_info::(program_id, account_info)?; let account_data = &mut account_info.try_borrow_mut_data()?; // Necessary to convince the borrow checker. let account_data: &'a mut [u8] = unsafe { @@ -159,6 +156,9 @@ impl<'a> BatchedQueueAccount<'a> { Self::from_bytes::(account_data) } + /// Deserialize a BatchedQueueAccount from bytes. + /// Should only be used in client. + /// Checks the discriminator and queue type. #[cfg(not(target_os = "solana"))] pub fn output_from_bytes( account_data: &'a mut [u8], @@ -222,18 +222,16 @@ impl<'a> BatchedQueueAccount<'a> { bloom_filter_capacity, )?; if account_data_len - != queue_account_size( - &account_metadata.batch_metadata, - account_metadata.metadata.queue_type, - )? + != account_metadata + .batch_metadata + .queue_account_size(account_metadata.metadata.queue_type)? { msg!("account_data.len() {:?}", account_data_len); msg!( "queue_account_size {:?}", - queue_account_size( - &account_metadata.batch_metadata, - account_metadata.metadata.queue_type - )? + account_metadata + .batch_metadata + .queue_account_size(account_metadata.metadata.queue_type)? ); return Err(ZeroCopyError::InvalidAccountSize.into()); } @@ -255,6 +253,10 @@ impl<'a> BatchedQueueAccount<'a> { }) } + /// Insert a value into the current batch + /// of this output queue account. + /// 1. insert value into a value vec and hash chain store. + /// 2. Increment next_index. pub fn insert_into_current_batch( &mut self, hash_chain_value: &[u8; 32], @@ -277,6 +279,15 @@ impl<'a> BatchedQueueAccount<'a> { Ok(()) } + /// Proves inclusion of leaf index if it exists in one of the batches. + /// 1. Iterate over all batches + /// 2. Check if leaf index could exist in the batch. + /// 2.1 If yes, check whether value at index is equal to hash_chain_value. + /// Throw error if not. + /// 3. Return true if leaf index exists in one of the batches. + /// + /// Note, this method does not fail but returns `false` + /// if the leaf index is out of range for any batch. pub fn prove_inclusion_by_index( &mut self, leaf_index: u64, @@ -299,7 +310,12 @@ impl<'a> BatchedQueueAccount<'a> { Ok(false) } - pub fn leaf_index_could_exist_in_batches( + /// Check that leaf index could exist in one of the batches. + /// Returns Ok(()) if value of leaf index could exist in batch. + /// This doesn't mean that the value exists in the batch, + /// just that it is plausible. The value might already be spent + /// or never inserted in case an invalid index was provided. + pub fn check_leaf_index_could_exist_in_batches( &mut self, leaf_index: u64, ) -> Result<(), BatchedMerkleTreeError> { @@ -338,16 +354,26 @@ impl<'a> BatchedQueueAccount<'a> { Ok(()) } + pub fn get_metadata(&self) -> &BatchedQueueMetadata { + &self.metadata + } + + pub fn get_metadata_mut(&mut self) -> &mut BatchedQueueMetadata { + &mut self.metadata + } + + /// Returns the number of elements inserted in the current batch. pub fn get_num_inserted_in_current_batch(&self) -> u64 { - let next_full_batch = self.batch_metadata.currently_processing_batch_index; - let batch = self.batches.get(next_full_batch as usize).unwrap(); - batch.get_num_inserted() + batch.get_current_zkp_batch_index() * batch.zkp_batch_size + let next_full_batch = self.batch_metadata.currently_processing_batch_index as usize; + self.batches[next_full_batch].get_num_inserted_elements() } + /// Returns true if the pubkey is the associated Merkle tree of the queue. pub fn is_associated(&self, pubkey: &Pubkey) -> bool { self.metadata.metadata.associated_merkle_tree == *pubkey } + /// Check if the pubkey is the associated Merkle tree of the queue. pub fn check_is_associated(&self, pubkey: &Pubkey) -> Result<(), BatchedMerkleTreeError> { if !self.is_associated(pubkey) { return Err(MerkleTreeMetadataError::MerkleTreeAndQueueNotAssociated.into()); @@ -355,15 +381,31 @@ impl<'a> BatchedQueueAccount<'a> { Ok(()) } + /// Returns true if the tree is full. + pub fn tree_is_full(&self) -> bool { + self.tree_capacity == self.next_index + } + + /// Check if the tree is full. pub fn check_tree_is_full(&self) -> Result<(), BatchedMerkleTreeError> { if self.tree_is_full() { return Err(BatchedMerkleTreeError::TreeIsFull); } Ok(()) } +} - pub fn tree_is_full(&self) -> bool { - self.tree_capacity == self.next_index +impl Deref for BatchedQueueAccount<'_> { + type Target = BatchedQueueMetadata; + + fn deref(&self) -> &Self::Target { + &self.metadata + } +} + +impl DerefMut for BatchedQueueAccount<'_> { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.metadata } } @@ -378,7 +420,7 @@ impl<'a> BatchedQueueAccount<'a> { /// 3. If batch is full, increment currently_processing_batch_index. #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] -pub fn insert_into_current_batch( +pub(crate) fn insert_into_current_batch( queue_type: u64, batch_metadata: &mut BatchMetadata, batches: &mut ZeroCopySliceMutU64, @@ -467,7 +509,7 @@ pub fn insert_into_current_batch( } #[allow(clippy::type_complexity)] -pub fn output_queue_from_bytes( +pub(crate) fn output_queue_from_bytes( num_value_stores: usize, num_stores: usize, num_hashchain_stores: usize, @@ -492,7 +534,7 @@ pub fn output_queue_from_bytes( } #[allow(clippy::type_complexity)] -pub fn input_queue_from_bytes<'a>( +pub(crate) fn input_queue_from_bytes<'a>( batch_metadata: &BatchMetadata, account_data: &'a mut [u8], queue_type: u64, @@ -521,7 +563,7 @@ pub fn input_queue_from_bytes<'a>( } #[allow(clippy::type_complexity)] -pub fn init_queue<'a>( +pub(crate) fn init_queue<'a>( batch_metadata: &BatchMetadata, queue_type: u64, account_data: &'a mut [u8], @@ -582,11 +624,10 @@ pub fn get_output_queue_account_size_default() -> usize { }, ..Default::default() }; - queue_account_size( - &batch_metadata.batch_metadata, - QueueType::BatchedOutput as u64, - ) - .unwrap() + batch_metadata + .batch_metadata + .queue_account_size(QueueType::BatchedOutput as u64) + .unwrap() } pub fn get_output_queue_account_size_from_params( @@ -603,7 +644,10 @@ pub fn get_output_queue_account_size_from_params( }, ..Default::default() }; - queue_account_size(&metadata.batch_metadata, QueueType::BatchedOutput as u64).unwrap() + metadata + .batch_metadata + .queue_account_size(QueueType::BatchedOutput as u64) + .unwrap() } pub fn get_output_queue_account_size( @@ -622,9 +666,13 @@ pub fn get_output_queue_account_size( }, ..Default::default() }; - queue_account_size(&metadata.batch_metadata, QueueType::BatchedOutput as u64).unwrap() + metadata + .batch_metadata + .queue_account_size(QueueType::BatchedOutput as u64) + .unwrap() } +#[cfg(not(target_os = "solana"))] #[allow(clippy::too_many_arguments)] pub fn assert_queue_inited( batch_metadata: BatchMetadata, diff --git a/program-libs/batched-merkle-tree/src/rollover_address_tree.rs b/program-libs/batched-merkle-tree/src/rollover_address_tree.rs index 80eabac1a8..2ecf6e1e3c 100644 --- a/program-libs/batched-merkle-tree/src/rollover_address_tree.rs +++ b/program-libs/batched-merkle-tree/src/rollover_address_tree.rs @@ -10,14 +10,22 @@ use crate::{ rollover_state_tree::batched_tree_is_ready_for_rollover, }; -/// Checks: -/// 1. Merkle tree is ready to be rolled over -/// 2. Merkle tree is not already rolled over -/// 3. Rollover threshold is configured, if not tree cannot be rolled over +/// Rollover an almost full batched address tree, +/// ie create a new batched Merkle tree account +/// with the same parameters, and mark the old account as rolled over. +/// The old tree can be used until it is completely full. /// -/// Actions: -/// 1. mark Merkle tree as rolled over in this slot -/// 2. initialize new Merkle tree and nullifier queue with the same parameters +/// Steps: +/// 1. Check that Merkle tree is ready to be rolled over: +/// 1.1. rollover threshold is configured +/// 1.2. next index is greater than rollover threshold +/// 1.3. the network fee is not set if the current fee is zero +/// 2. Rollover the old Merkle tree and check: +/// 2.2. Rollover is configured. +/// 2.3. Tree is not already rolled over. +/// 2.4. Mark as rolled over in this slot. +/// 3. Initialize new batched address Merkle tree account +/// with the same parameters as the old account. pub fn rollover_batched_address_tree<'a>( old_merkle_tree: &mut BatchedMerkleTreeAccount<'a>, new_mt_data: &'a mut [u8], @@ -28,12 +36,12 @@ pub fn rollover_batched_address_tree<'a>( // 1. Check that old merkle tree is ready for rollover. batched_tree_is_ready_for_rollover(old_merkle_tree, &network_fee)?; - // Rollover the old merkle tree. + // 2. Rollover the old merkle tree. old_merkle_tree .metadata .rollover(Pubkey::default(), new_mt_pubkey)?; - // Initialize the new address merkle tree. + // 3. Initialize the new address merkle tree. let params = create_batched_address_tree_init_params(old_merkle_tree, network_fee); let owner = old_merkle_tree.metadata.access_metadata.owner; init_batched_address_merkle_tree_account(owner, params, new_mt_data, new_mt_rent) @@ -79,9 +87,9 @@ fn create_batched_address_tree_init_params( #[cfg(not(target_os = "solana"))] pub fn assert_address_mt_roll_over( mut old_mt_account_data: Vec, - mut old_ref_mt_account: crate::merkle_tree::BatchedMerkleTreeMetadata, + mut old_ref_mt_account: crate::merkle_tree_metadata::BatchedMerkleTreeMetadata, mut new_mt_account_data: Vec, - new_ref_mt_account: crate::merkle_tree::BatchedMerkleTreeMetadata, + new_ref_mt_account: crate::merkle_tree_metadata::BatchedMerkleTreeMetadata, new_mt_pubkey: Pubkey, bloom_filter_num_iters: u64, ) { diff --git a/program-libs/batched-merkle-tree/src/rollover_state_tree.rs b/program-libs/batched-merkle-tree/src/rollover_state_tree.rs index 2841f50ecc..dafb08d62e 100644 --- a/program-libs/batched-merkle-tree/src/rollover_state_tree.rs +++ b/program-libs/batched-merkle-tree/src/rollover_state_tree.rs @@ -7,7 +7,8 @@ use crate::{ initialize_state_tree::{ init_batched_state_merkle_tree_accounts, InitStateTreeAccountsInstructionData, }, - merkle_tree::{BatchedMerkleTreeAccount, BatchedMerkleTreeMetadata}, + merkle_tree::BatchedMerkleTreeAccount, + merkle_tree_metadata::BatchedMerkleTreeMetadata, queue::{BatchedQueueAccount, BatchedQueueMetadata}, }; @@ -28,29 +29,39 @@ pub struct RolloverBatchStateTreeParams<'a> { pub network_fee: Option, } -/// Checks: -/// 1. Merkle tree is ready to be rolled over -/// 2. Merkle tree is not already rolled over -/// 3. Rollover threshold is configured, if not tree cannot be rolled over +/// Rollover an almost full batched state tree, +/// ie create a new batched Merkle tree and output queue +/// with the same parameters, and mark the old accounts as rolled over. +/// The old tree and queue can be used until these are completely full. /// -/// Actions: -/// 1. mark Merkle tree as rolled over in this slot -/// 2. initialize new Merkle tree and output queue with the same parameters +/// Steps: +/// 1. Check that Merkle tree is ready to be rolled over: +/// 1.1. rollover threshold is configured +/// 1.2. next index is greater than rollover threshold +/// 1.3. the network fee is not set if the current fee is zero +/// 2. Rollover Merkle tree and check: +/// 2.1. Merkle tree and queue are associated. +/// 2.2. Rollover is configured. +/// 2.3. Tree is not already rolled over. +/// 2.4. Mark as rolled over in this slot. +/// 3. Rollover output queue and check: +/// 3.1. Merkle tree and queue are associated. +/// 3.2. Rollover is configured. +/// 3.3. Tree is not already rolled over. +/// 3.4. Mark as rolled over in this slot. +/// 4. Initialize new Merkle tree and output queue +/// with the same parameters as old accounts. pub fn rollover_batched_state_tree( params: RolloverBatchStateTreeParams, ) -> Result<(), BatchedMerkleTreeError> { - params - .old_output_queue - .check_is_associated(¶ms.old_mt_pubkey)?; - - // Check that old merkle tree is ready for rollover. + // 1. Check that old merkle tree is ready for rollover. batched_tree_is_ready_for_rollover(params.old_merkle_tree, ¶ms.network_fee)?; - // Rollover the old merkle tree. + // 2. Rollover the old merkle tree. params .old_merkle_tree .metadata .rollover(params.old_queue_pubkey, params.new_mt_pubkey)?; - // Rollover the old output queue. + // 3. Rollover the old output queue. params .old_output_queue .metadata @@ -58,7 +69,7 @@ pub fn rollover_batched_state_tree( let init_params = InitStateTreeAccountsInstructionData::from(¶ms); let owner = params.old_merkle_tree.metadata.access_metadata.owner; - // Initialize the new merkle tree and output queue. + // 4. Initialize the new merkle tree and output queue. init_batched_state_merkle_tree_accounts( owner, init_params, @@ -123,6 +134,10 @@ impl From<&RolloverBatchStateTreeParams<'_>> for InitStateTreeAccountsInstructio } // TODO: add unit test +/// Check that: +/// 1. rollover threshold is configured +/// 2. next index is greater than rollover threshold +/// 3. the network fee is not set if the current fee is zero pub fn batched_tree_is_ready_for_rollover( metadata: &BatchedMerkleTreeAccount<'_>, network_fee: &Option, diff --git a/program-libs/batched-merkle-tree/tests/zero_copy.rs b/program-libs/batched-merkle-tree/tests/account_access.rs similarity index 100% rename from program-libs/batched-merkle-tree/tests/zero_copy.rs rename to program-libs/batched-merkle-tree/tests/account_access.rs diff --git a/program-libs/batched-merkle-tree/tests/initialize_address_tree.rs b/program-libs/batched-merkle-tree/tests/initialize_address_tree.rs index 61afe03956..02e7a52a39 100644 --- a/program-libs/batched-merkle-tree/tests/initialize_address_tree.rs +++ b/program-libs/batched-merkle-tree/tests/initialize_address_tree.rs @@ -4,10 +4,8 @@ use light_batched_merkle_tree::{ init_batched_address_merkle_tree_account, InitAddressTreeAccountsInstructionData, }, initialize_state_tree::assert_address_mt_zero_copy_inited, - merkle_tree::{ - get_merkle_tree_account_size, get_merkle_tree_account_size_default, - BatchedMerkleTreeMetadata, CreateTreeParams, - }, + merkle_tree::{get_merkle_tree_account_size, get_merkle_tree_account_size_default}, + merkle_tree_metadata::{BatchedMerkleTreeMetadata, CreateTreeParams}, }; use light_utils::pubkey::Pubkey; use light_zero_copy::{ diff --git a/program-libs/batched-merkle-tree/tests/initialize_state_tree.rs b/program-libs/batched-merkle-tree/tests/initialize_state_tree.rs index 4145d3f628..b6ff2f86d3 100644 --- a/program-libs/batched-merkle-tree/tests/initialize_state_tree.rs +++ b/program-libs/batched-merkle-tree/tests/initialize_state_tree.rs @@ -5,10 +5,8 @@ use light_batched_merkle_tree::{ init_batched_state_merkle_tree_accounts, CreateOutputQueueParams, InitStateTreeAccountsInstructionData, }, - merkle_tree::{ - get_merkle_tree_account_size, get_merkle_tree_account_size_default, - BatchedMerkleTreeMetadata, CreateTreeParams, - }, + merkle_tree::{get_merkle_tree_account_size, get_merkle_tree_account_size_default}, + merkle_tree_metadata::{BatchedMerkleTreeMetadata, CreateTreeParams}, queue::{ assert_queue_zero_copy_inited, get_output_queue_account_size, get_output_queue_account_size_default, BatchedQueueMetadata, diff --git a/program-libs/batched-merkle-tree/tests/merkle_tree.rs b/program-libs/batched-merkle-tree/tests/merkle_tree.rs index 6cf6fe257e..c983a03472 100644 --- a/program-libs/batched-merkle-tree/tests/merkle_tree.rs +++ b/program-libs/batched-merkle-tree/tests/merkle_tree.rs @@ -18,9 +18,10 @@ use light_batched_merkle_tree::{ }, merkle_tree::{ assert_batch_append_event_event, assert_nullify_event, - get_merkle_tree_account_size_default, BatchedMerkleTreeAccount, BatchedMerkleTreeMetadata, + get_merkle_tree_account_size_default, BatchedMerkleTreeAccount, InstructionDataBatchAppendInputs, InstructionDataBatchNullifyInputs, }, + merkle_tree_metadata::BatchedMerkleTreeMetadata, queue::{ get_output_queue_account_size_default, get_output_queue_account_size_from_params, BatchedQueueAccount, BatchedQueueMetadata, diff --git a/program-libs/batched-merkle-tree/tests/queue.rs b/program-libs/batched-merkle-tree/tests/queue.rs index 844a1283fb..f78e97dd28 100644 --- a/program-libs/batched-merkle-tree/tests/queue.rs +++ b/program-libs/batched-merkle-tree/tests/queue.rs @@ -1,10 +1,7 @@ use light_batched_merkle_tree::{ batch_metadata::BatchMetadata, errors::BatchedMerkleTreeError, - queue::{ - assert_queue_zero_copy_inited, queue_account_size, BatchedQueueAccount, - BatchedQueueMetadata, - }, + queue::{assert_queue_zero_copy_inited, BatchedQueueAccount, BatchedQueueMetadata}, }; use light_merkle_tree_metadata::{ access::AccessMetadata, @@ -40,8 +37,13 @@ pub fn get_test_account_and_account_data( }, ..Default::default() }; - let account_data: Vec = - vec![0; queue_account_size(&account.batch_metadata, account.metadata.queue_type).unwrap()]; + let account_data: Vec = vec![ + 0; + account + .batch_metadata + .queue_account_size(account.metadata.queue_type) + .unwrap() + ]; (account, account_data) } diff --git a/program-libs/batched-merkle-tree/tests/rollover_address_tree.rs b/program-libs/batched-merkle-tree/tests/rollover_address_tree.rs index 2ee9ba237e..3770572304 100644 --- a/program-libs/batched-merkle-tree/tests/rollover_address_tree.rs +++ b/program-libs/batched-merkle-tree/tests/rollover_address_tree.rs @@ -6,8 +6,9 @@ use light_batched_merkle_tree::{ initialize_state_tree::assert_address_mt_zero_copy_inited, merkle_tree::{ get_merkle_tree_account_size, get_merkle_tree_account_size_default, - BatchedMerkleTreeAccount, BatchedMerkleTreeMetadata, CreateTreeParams, + BatchedMerkleTreeAccount, }, + merkle_tree_metadata::{BatchedMerkleTreeMetadata, CreateTreeParams}, rollover_address_tree::{assert_address_mt_roll_over, rollover_batched_address_tree}, }; use light_merkle_tree_metadata::errors::MerkleTreeMetadataError; diff --git a/program-libs/batched-merkle-tree/tests/rollover_state_tree.rs b/program-libs/batched-merkle-tree/tests/rollover_state_tree.rs index ffee6f568a..e9143512bc 100644 --- a/program-libs/batched-merkle-tree/tests/rollover_state_tree.rs +++ b/program-libs/batched-merkle-tree/tests/rollover_state_tree.rs @@ -7,8 +7,9 @@ use light_batched_merkle_tree::{ }, merkle_tree::{ get_merkle_tree_account_size, get_merkle_tree_account_size_default, - BatchedMerkleTreeAccount, BatchedMerkleTreeMetadata, CreateTreeParams, + BatchedMerkleTreeAccount, }, + merkle_tree_metadata::{BatchedMerkleTreeMetadata, CreateTreeParams}, queue::{ assert_queue_zero_copy_inited, get_output_queue_account_size, get_output_queue_account_size_default, BatchedQueueAccount, diff --git a/program-tests/account-compression-test/tests/batched_merkle_tree_test.rs b/program-tests/account-compression-test/tests/batched_merkle_tree_test.rs index d0b9caefd5..0fddfea1ec 100644 --- a/program-tests/account-compression-test/tests/batched_merkle_tree_test.rs +++ b/program-tests/account-compression-test/tests/batched_merkle_tree_test.rs @@ -13,9 +13,10 @@ use light_batched_merkle_tree::{ create_output_queue_account, CreateOutputQueueParams, InitStateTreeAccountsInstructionData, }, merkle_tree::{ - get_merkle_tree_account_size, BatchedMerkleTreeAccount, BatchedMerkleTreeMetadata, - CreateTreeParams, InstructionDataBatchAppendInputs, InstructionDataBatchNullifyInputs, + get_merkle_tree_account_size, BatchedMerkleTreeAccount, InstructionDataBatchAppendInputs, + InstructionDataBatchNullifyInputs, }, + merkle_tree_metadata::{BatchedMerkleTreeMetadata, CreateTreeParams}, queue::{ assert_queue_zero_copy_inited, get_output_queue_account_size, BatchedQueueAccount, BatchedQueueMetadata, diff --git a/program-tests/registry-test/tests/tests.rs b/program-tests/registry-test/tests/tests.rs index 0ed04caf74..6785088f51 100644 --- a/program-tests/registry-test/tests/tests.rs +++ b/program-tests/registry-test/tests/tests.rs @@ -15,7 +15,8 @@ use light_batched_merkle_tree::{ initialize_state_tree::{ assert_address_mt_zero_copy_inited, InitStateTreeAccountsInstructionData, }, - merkle_tree::{BatchedMerkleTreeAccount, BatchedMerkleTreeMetadata, CreateTreeParams}, + merkle_tree::BatchedMerkleTreeAccount, + merkle_tree_metadata::{BatchedMerkleTreeMetadata, CreateTreeParams}, queue::BatchedQueueAccount, }; use light_client::indexer::Indexer; diff --git a/program-tests/system-cpi-test/tests/test.rs b/program-tests/system-cpi-test/tests/test.rs index e8c9f3c0fa..49328bd54a 100644 --- a/program-tests/system-cpi-test/tests/test.rs +++ b/program-tests/system-cpi-test/tests/test.rs @@ -760,8 +760,8 @@ async fn only_test_create_pda() { assert_rpc_error( result, 0, - UtilsError::AccountNotMutable.into(), - // AccountCompressionErrorCode::AddressMerkleTreeAccountDiscriminatorMismatch.into(), + // UtilsError::AccountNotMutable.into(), + UtilsError::InvalidDiscriminator.into(), ) .unwrap(); diff --git a/programs/account-compression/src/instructions/migrate_state.rs b/programs/account-compression/src/instructions/migrate_state.rs index ef282433f8..bfbf70a569 100644 --- a/programs/account-compression/src/instructions/migrate_state.rs +++ b/programs/account-compression/src/instructions/migrate_state.rs @@ -137,7 +137,7 @@ fn migrate_state( mod migrate_state_test { use light_batched_merkle_tree::{ batch_metadata::BatchMetadata, - queue::{queue_account_size, BatchedQueueAccount, BatchedQueueMetadata}, + queue::{BatchedQueueAccount, BatchedQueueMetadata}, }; use light_concurrent_merkle_tree::ConcurrentMerkleTree; use light_hasher::Poseidon; @@ -180,11 +180,13 @@ mod migrate_state_test { }, tree_capacity: 2u64.pow(32), }; - let account_data: Vec = - vec![ - 0; - queue_account_size(&account.batch_metadata, account.metadata.queue_type).unwrap() - ]; + let account_data: Vec = vec![ + 0; + account + .batch_metadata + .queue_account_size(account.metadata.queue_type) + .unwrap() + ]; let mut mock_account = MockQueueAccount { account_data, account: None, diff --git a/programs/system/src/invoke/verify_state_proof.rs b/programs/system/src/invoke/verify_state_proof.rs index 60f8d3d063..a4d428b355 100644 --- a/programs/system/src/invoke/verify_state_proof.rs +++ b/programs/system/src/invoke/verify_state_proof.rs @@ -162,7 +162,7 @@ pub fn verify_input_accounts_proof_by_index( &mut BatchedQueueAccount::output_from_account_info(output_queue_account_info) .map_err(ProgramError::from)?; output_queue - .leaf_index_could_exist_in_batches(account.merkle_context.leaf_index as u64) + .check_leaf_index_could_exist_in_batches(account.merkle_context.leaf_index as u64) .map_err(ProgramError::from)?; } } diff --git a/sdk-libs/program-test/src/test_batch_forester.rs b/sdk-libs/program-test/src/test_batch_forester.rs index 5ef2469721..381a86f4ec 100644 --- a/sdk-libs/program-test/src/test_batch_forester.rs +++ b/sdk-libs/program-test/src/test_batch_forester.rs @@ -10,9 +10,10 @@ use light_batched_merkle_tree::{ create_output_queue_account, CreateOutputQueueParams, InitStateTreeAccountsInstructionData, }, merkle_tree::{ - get_merkle_tree_account_size, BatchedMerkleTreeAccount, BatchedMerkleTreeMetadata, - CreateTreeParams, InstructionDataBatchAppendInputs, InstructionDataBatchNullifyInputs, + get_merkle_tree_account_size, BatchedMerkleTreeAccount, InstructionDataBatchAppendInputs, + InstructionDataBatchNullifyInputs, }, + merkle_tree_metadata::{BatchedMerkleTreeMetadata, CreateTreeParams}, queue::{ assert_queue_zero_copy_inited, get_output_queue_account_size, BatchedQueueAccount, BatchedQueueMetadata,