feat: rollover state trees (#1371)

* feat: rollover batch state trees * test: batch state tree rollover
Lightprotocol · Dec 4, 2024 · 5ffede6 · 5ffede6
1 parent dfe7903
commit 5ffede6
Show file tree

Hide file tree

Showing 20 changed files with 2,052 additions and 101 deletions.
diff --git a/programs/account-compression/src/errors.rs b/programs/account-compression/src/errors.rs
@@ -76,4 +76,6 @@ pub enum AccountCompressionErrorCode {
     InvalidBatch,
     LeafIndexNotInBatch,
     UnsupportedParameters,
+    InvalidTreeType,
+    InvalidNetworkFee,
 }
diff --git a/programs/account-compression/src/instructions/intialize_batched_state_merkle_tree.rs b/programs/account-compression/src/instructions/intialize_batched_state_merkle_tree.rs
@@ -219,11 +219,13 @@ pub fn bytes_to_struct_checked<T: Clone + Copy + Discriminator, const INIT: bool
 
     if INIT {
         if bytes[0..8] != [0; 8] {
+            #[cfg(target_os = "solana")]
             msg!("Discriminator bytes must be zero for initialization.");
             return err!(AccountCompressionErrorCode::InvalidDiscriminator);
         }
         bytes[0..8].copy_from_slice(&T::DISCRIMINATOR);
     } else if T::DISCRIMINATOR != bytes[0..8] {
+        #[cfg(target_os = "solana")]
         msg!(
             "Expected discriminator: {:?}, actual {:?} ",
             T::DISCRIMINATOR,
@@ -377,10 +379,7 @@ pub fn assert_mt_zero_copy_inited(
         ref_account,
         "metadata mismatch"
     );
-    println!(
-        "zero_copy_account.root_history.capacity(): {}",
-        zero_copy_account.root_history.metadata().capacity()
-    );
+
     assert_eq!(
         zero_copy_account.root_history.capacity(),
         ref_account.root_history_capacity as usize,
@@ -421,6 +420,7 @@ pub fn get_output_queue_account_default(
     associated_merkle_tree: Pubkey,
     height: u32,
     num_batches: u64,
+    network_fee: u64,
 ) -> BatchedQueueAccount {
     let rollover_fee = match rollover_threshold {
         Some(rollover_threshold) => compute_rollover_fee(rollover_threshold, height, rent)
@@ -441,7 +441,7 @@ pub fn get_output_queue_account_default(
             rolledover_slot: u64::MAX,
             rollover_threshold: rollover_threshold.unwrap_or(u64::MAX),
             rollover_fee,
-            network_fee: 5000,
+            network_fee,
             additional_bytes,
         },
         queue_type: QueueType::Output as u64,
@@ -566,6 +566,7 @@ pub mod tests {
             mt_pubkey,
             params.height,
             params.output_queue_num_batches,
+            params.network_fee.unwrap_or_default(),
         );
         assert_queue_zero_copy_inited(
             output_queue_account_data.as_mut_slice(),

diff --git a/programs/account-compression/src/instructions/mod.rs b/programs/account-compression/src/instructions/mod.rs
@@ -42,3 +42,6 @@ pub use batch_nullify::*;
 
 pub mod batch_append;
 pub use batch_append::*;
+
+pub mod rollover_batch_state_merkle_tree;
+pub use rollover_batch_state_merkle_tree::*;
diff --git a/programs/account-compression/src/instructions/rollover_address_merkle_tree_and_queue.rs b/programs/account-compression/src/instructions/rollover_address_merkle_tree_and_queue.rs
@@ -48,7 +48,7 @@ impl<'info> GroupAccounts<'info> for RolloverAddressMerkleTreeAndQueue<'info> {
 ///
 /// Actions:
 /// 1. mark Merkle tree as rolled over in this slot
-/// 2. initialize new Merkle tree and nullifier queue with the same parameters
+/// 2. initialize new Merkle tree and queue with the same parameters
 pub fn process_rollover_address_merkle_tree_and_queue<'a, 'b, 'c: 'info, 'info>(
     ctx: Context<'a, 'b, 'c, 'info, RolloverAddressMerkleTreeAndQueue<'info>>,
 ) -> Result<()> {