Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions anchor-programs/system/src/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -88,4 +88,6 @@ pub enum SystemProgramError {
TooManyOutputAccounts,
#[msg("Failed to borrow account data")]
BorrowingDataFailed,
#[msg("DuplicateAccountInInputsAndReadOnly")]
DuplicateAccountInInputsAndReadOnly,
}
3 changes: 1 addition & 2 deletions program-libs/batched-merkle-tree/src/batch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -989,10 +989,9 @@ mod tests {
num_iters,
bloom_filter_capacity,
Pubkey::new_unique(),
16, // Tree height 4 -> capacity 16
)
.unwrap();
// Tree height 4 -> capacity 16
account.tree_capacity = 16;
assert_eq!(account.get_num_inserted_in_current_batch(), 0);
// Fill first batch
for i in 1..=batch_size {
Expand Down
40 changes: 37 additions & 3 deletions program-libs/batched-merkle-tree/src/initialize_address_tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ use crate::{
DEFAULT_BATCH_ROOT_HISTORY_LEN, DEFAULT_BATCH_SIZE,
},
errors::BatchedMerkleTreeError,
initialize_state_tree::match_circuit_size,
merkle_tree::{get_merkle_tree_account_size, BatchedMerkleTreeAccount},
BorshDeserialize, BorshSerialize,
};
Expand Down Expand Up @@ -129,6 +128,7 @@ pub fn init_batched_address_merkle_tree_account(
)
}

/// Only used for testing. For production use the default config.
pub fn validate_batched_address_tree_params(params: InitAddressTreeAccountsInstructionData) {
assert!(params.input_queue_batch_size > 0);
assert_eq!(
Expand All @@ -138,7 +138,7 @@ pub fn validate_batched_address_tree_params(params: InitAddressTreeAccountsInstr
);
assert!(
match_circuit_size(params.input_queue_zkp_batch_size),
"Zkp batch size not supported. Supported 1, 10, 100, 500, 1000"
"Zkp batch size not supported. Supported: 10, 250"
);

assert!(params.bloom_filter_num_iters > 0);
Expand All @@ -151,10 +151,24 @@ pub fn validate_batched_address_tree_params(params: InitAddressTreeAccountsInstr
assert!(params.bloom_filter_capacity > 0);
assert!(params.root_history_capacity > 0);
assert!(params.input_queue_batch_size > 0);

// Validate root_history_capacity is sufficient for input operations
// (address trees only have input queues, no output queues)
let required_capacity = params.input_queue_batch_size / params.input_queue_zkp_batch_size;
assert!(
params.root_history_capacity >= required_capacity as u32,
"root_history_capacity ({}) must be >= {} (input_queue_batch_size / input_queue_zkp_batch_size)",
params.root_history_capacity,
required_capacity
);

assert_eq!(params.close_threshold, None);
assert_eq!(params.height, DEFAULT_BATCH_ADDRESS_TREE_HEIGHT);
}

/// Only 10 and 250 are supported.
pub fn match_circuit_size(size: u64) -> bool {
matches!(size, 10 | 250)
}
pub fn get_address_merkle_tree_account_size_from_params(
params: InitAddressTreeAccountsInstructionData,
) -> usize {
Expand Down Expand Up @@ -330,3 +344,23 @@ fn test_height_not_40() {
};
validate_batched_address_tree_params(params);
}

#[test]
fn test_validate_root_history_capacity_address_tree() {
// Test with valid params (default should pass)
let params = InitAddressTreeAccountsInstructionData::default();
validate_batched_address_tree_params(params); // Should not panic
}

#[test]
#[should_panic(expected = "root_history_capacity")]
fn test_validate_root_history_capacity_insufficient_address_tree() {
let params = InitAddressTreeAccountsInstructionData {
root_history_capacity: 1, // Much too small
input_queue_batch_size: 1000,
input_queue_zkp_batch_size: 10,
// Required: 1000/10 = 100, but we set only 1
..Default::default()
};
validate_batched_address_tree_params(params); // Should panic
}
42 changes: 39 additions & 3 deletions program-libs/batched-merkle-tree/src/initialize_state_tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,7 @@ pub fn init_batched_state_merkle_tree_accounts<'a>(
0,
0,
output_queue_pubkey,
2u64.pow(params.height),
)?;
}
let metadata = MerkleTreeMetadata {
Expand Down Expand Up @@ -198,6 +199,7 @@ pub fn init_batched_state_merkle_tree_accounts<'a>(
)
}

/// Only used for testing. For production use the default config.
pub fn validate_batched_tree_params(params: InitStateTreeAccountsInstructionData) {
assert!(params.input_queue_batch_size > 0);
assert!(params.output_queue_batch_size > 0);
Expand All @@ -213,11 +215,11 @@ pub fn validate_batched_tree_params(params: InitStateTreeAccountsInstructionData
);
assert!(
match_circuit_size(params.input_queue_zkp_batch_size),
"Zkp batch size not supported. Supported 1, 10, 100, 500, 1000"
"Zkp batch size not supported. Supported 10, 500"
);
assert!(
match_circuit_size(params.output_queue_zkp_batch_size),
"Zkp batch size not supported. Supported 1, 10, 100, 500, 1000"
"Zkp batch size not supported. Supported 10, 500"
);

assert!(params.bloom_filter_num_iters > 0);
Expand All @@ -230,12 +232,46 @@ pub fn validate_batched_tree_params(params: InitStateTreeAccountsInstructionData
assert!(params.bloom_filter_capacity > 0);
assert!(params.root_history_capacity > 0);
assert!(params.input_queue_batch_size > 0);

// Validate root_history_capacity is sufficient for both input and output operations
let required_capacity = (params.output_queue_batch_size / params.output_queue_zkp_batch_size)
+ (params.input_queue_batch_size / params.input_queue_zkp_batch_size);
assert!(
params.root_history_capacity >= required_capacity as u32,
"root_history_capacity ({}) must be >= {} (output_queue_batch_size / output_queue_zkp_batch_size + input_queue_batch_size / input_queue_zkp_batch_size)",
params.root_history_capacity,
required_capacity
);

assert_eq!(params.close_threshold, None);
assert_eq!(params.height, DEFAULT_BATCH_STATE_TREE_HEIGHT);
}

/// Only 10 and 500 are supported.
pub fn match_circuit_size(size: u64) -> bool {
matches!(size, 10 | 100 | 250 | 500 | 1000)
matches!(size, 10 | 500)
}

#[test]
fn test_validate_root_history_capacity_state_tree() {
// Test with valid params (default should pass)
let params = InitStateTreeAccountsInstructionData::default();
validate_batched_tree_params(params); // Should not panic
}

#[test]
#[should_panic(expected = "root_history_capacity")]
fn test_validate_root_history_capacity_insufficient_state_tree() {
let params = InitStateTreeAccountsInstructionData {
root_history_capacity: 1, // Much too small
input_queue_batch_size: 1000,
output_queue_batch_size: 1000,
input_queue_zkp_batch_size: 10,
output_queue_zkp_batch_size: 10,
// Required: (1000/10) + (1000/10) = 200, but we set only 1
..Default::default()
};
validate_batched_tree_params(params); // Should panic
}
#[cfg(feature = "test-only")]
pub mod test_utils {
Expand Down
50 changes: 35 additions & 15 deletions program-libs/batched-merkle-tree/src/merkle_tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -376,7 +376,7 @@ impl<'a> BatchedMerkleTreeAccount<'a> {
queue_account: &mut BatchedQueueAccount,
instruction_data: InstructionDataBatchAppendInputs,
) -> Result<MerkleTreeEvent, BatchedMerkleTreeError> {
self.check_tree_is_full()?;
self.check_tree_is_full(Some(queue_account.batch_metadata.zkp_batch_size))?;
let pending_batch_index = queue_account.batch_metadata.pending_batch_index as usize;
let new_root = instruction_data.new_root;
let circuit_batch_size = queue_account.batch_metadata.zkp_batch_size;
Expand Down Expand Up @@ -465,7 +465,7 @@ impl<'a> BatchedMerkleTreeAccount<'a> {
if self.tree_type != TreeType::AddressV2 as u64 {
return Err(MerkleTreeMetadataError::InvalidTreeType.into());
}
self.check_tree_is_full()?;
self.check_tree_is_full(Some(self.metadata.queue_batches.zkp_batch_size))?;
Ok(MerkleTreeEvent::BatchAddressAppend(
self.update_input_queue::<ADDRESS_QUEUE_TYPE_V2>(instruction_data)?,
))
Expand Down Expand Up @@ -886,8 +886,10 @@ impl<'a> BatchedMerkleTreeAccount<'a> {
Ok(())
}

pub fn tree_is_full(&self) -> bool {
self.next_index >= self.capacity
/// Checks if the tree is full, optionally for a batch size.
/// If batch_size is provided, checks if there is enough space for the batch.
pub fn tree_is_full(&self, batch_size: Option<u64>) -> bool {
self.next_index + batch_size.unwrap_or_default() >= self.capacity
}

pub fn check_queue_next_index_reached_tree_capacity(
Expand All @@ -899,8 +901,13 @@ impl<'a> BatchedMerkleTreeAccount<'a> {
Ok(())
}

pub fn check_tree_is_full(&self) -> Result<(), BatchedMerkleTreeError> {
if self.tree_is_full() {
/// Checks if the tree is full, optionally for a batch size.
/// If batch_size is provided, checks if there is enough space for the batch.
pub fn check_tree_is_full(
&self,
batch_size: Option<u64>,
) -> Result<(), BatchedMerkleTreeError> {
if self.tree_is_full(batch_size) {
return Err(BatchedMerkleTreeError::TreeIsFull);
}
Ok(())
Expand Down Expand Up @@ -1578,7 +1585,7 @@ mod test {
)
.unwrap();
// 1. empty tree is not full
assert!(!account.tree_is_full());
assert!(!account.tree_is_full(None));

let mut inserted_elements = vec![];
let rng = &mut rand::rngs::StdRng::from_seed([0u8; 32]);
Expand Down Expand Up @@ -1677,17 +1684,30 @@ mod test {
)
.unwrap();
// 1. empty tree is not full
assert!(!account.tree_is_full());
assert!(account.check_tree_is_full().is_ok());
assert!(!account.tree_is_full(None));
assert!(account.check_tree_is_full(None).is_ok());
assert!(!account.tree_is_full(Some(1)));
assert!(account.check_tree_is_full(Some(1)).is_ok());
account.next_index = account.capacity - 2;
assert!(!account.tree_is_full(None));
assert!(account.check_tree_is_full(None).is_ok());
assert!(!account.tree_is_full(Some(1)));
assert!(account.check_tree_is_full(Some(1)).is_ok());
account.next_index = account.capacity - 1;
assert!(!account.tree_is_full());
assert!(account.check_tree_is_full().is_ok());
assert!(!account.tree_is_full(None));
assert!(account.check_tree_is_full(None).is_ok());
assert!(account.tree_is_full(Some(1)));
assert!(account.check_tree_is_full(Some(1)).is_err());
account.next_index = account.capacity;
assert!(account.tree_is_full());
assert!(account.check_tree_is_full().is_err());
assert!(account.tree_is_full(None));
assert!(account.check_tree_is_full(None).is_err());
assert!(account.tree_is_full(Some(1)));
assert!(account.check_tree_is_full(Some(1)).is_err());
account.next_index = account.capacity + 1;
assert!(account.tree_is_full());
assert!(account.check_tree_is_full().is_err());
assert!(account.tree_is_full(None));
assert!(account.check_tree_is_full(None).is_err());
assert!(account.tree_is_full(Some(1)));
assert!(account.check_tree_is_full(Some(1)).is_err());
}
#[test]
fn test_increment_next_index() {
Expand Down
Loading