fix: wipe_previous_batch_bloom_filter

forester-utils/src/instructions.rs

@@ -1,8 +1,8 @@
 use light_batched_merkle_tree::{
     constants::{DEFAULT_BATCH_ADDRESS_TREE_HEIGHT, DEFAULT_BATCH_STATE_TREE_HEIGHT},
     merkle_tree::{
-        AppendBatchProofInputsIx, BatchProofInputsIx, BatchedMerkleTreeAccount,
-        InstructionDataBatchAppendInputs, InstructionDataBatchNullifyInputs,
+        BatchedMerkleTreeAccount, InstructionDataBatchAppendInputs,
+        InstructionDataBatchNullifyInputs,
     },
     queue::BatchedQueueAccount,
 };
@@ -58,12 +58,11 @@ where
         })?
         .unwrap();
 
-    let (old_root_index, leaves_hashchain, start_index, current_root, batch_size, full_batch_index) = {
+    let (leaves_hashchain, start_index, current_root, batch_size, full_batch_index) = {
         let merkle_tree =
             BatchedMerkleTreeAccount::address_from_bytes(merkle_tree_account.data.as_mut_slice())
                 .unwrap();
 
-        let old_root_index = merkle_tree.root_history.last_index();
         let full_batch_index = merkle_tree.queue_metadata.next_full_batch_index;
         let batch = &merkle_tree.batches[full_batch_index as usize];
         let zkp_batch_index = batch.get_num_inserted_zkps();
@@ -74,7 +73,6 @@ where
         let batch_size = batch.zkp_batch_size as usize;
 
         (
-            old_root_index,
             leaves_hashchain,
             start_index,
             current_root,
@@ -175,7 +173,7 @@ where
         })?;
 
     let client = Client::new();
-    let circuit_inputs_new_root = bigint_to_be_bytes_array::<32>(&inputs.new_root).unwrap();
+    let new_root = bigint_to_be_bytes_array::<32>(&inputs.new_root).unwrap();
     let inputs = to_json(&inputs);
 
     let response_result = client
@@ -192,10 +190,7 @@ where
         let (proof_a, proof_b, proof_c) = proof_from_json_struct(proof_json);
         let (proof_a, proof_b, proof_c) = compress_proof(&proof_a, &proof_b, &proof_c);
         let instruction_data = InstructionDataBatchNullifyInputs {
-            public_inputs: BatchProofInputsIx {
-                new_root: circuit_inputs_new_root,
-                old_root_index: old_root_index as u16,
-            },
+            new_root,
             compressed_proof: CompressedProof {
                 a: proof_a,
                 b: proof_b,
@@ -330,7 +325,7 @@ pub async fn create_append_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
     };
 
     Ok(InstructionDataBatchAppendInputs {
-        public_inputs: AppendBatchProofInputsIx { new_root },
+        new_root,
         compressed_proof: proof,
     })
 }
@@ -340,7 +335,7 @@ pub async fn create_nullify_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
     indexer: &mut I,
     merkle_tree_pubkey: Pubkey,
 ) -> Result<InstructionDataBatchNullifyInputs, ForesterUtilsError> {
-    let (zkp_batch_size, old_root, old_root_index, leaves_hashchain) = {
+    let (zkp_batch_size, old_root, leaves_hashchain) = {
         let mut account = rpc.get_account(merkle_tree_pubkey).await.unwrap().unwrap();
         let merkle_tree =
             BatchedMerkleTreeAccount::state_from_bytes(account.data.as_mut_slice()).unwrap();
@@ -349,9 +344,8 @@ pub async fn create_nullify_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
         let batch = &merkle_tree.batches[batch_idx];
         let zkp_idx = batch.get_num_inserted_zkps();
         let hashchain = merkle_tree.hashchain_store[batch_idx][zkp_idx as usize];
-        let root_idx = merkle_tree.root_history.last_index();
         let root = *merkle_tree.root_history.last().unwrap();
-        (zkp_size, root, root_idx, hashchain)
+        (zkp_size, root, hashchain)
     };
 
     let leaf_indices_tx_hashes =
@@ -434,10 +428,7 @@ pub async fn create_nullify_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
     };
 
     Ok(InstructionDataBatchNullifyInputs {
-        public_inputs: BatchProofInputsIx {
-            new_root,
-            old_root_index: old_root_index as u16,
-        },
+        new_root,
         compressed_proof: proof,
     })
 }

program-libs/batched-merkle-tree/src/batch.rs

@@ -10,7 +10,7 @@ use crate::errors::BatchedMerkleTreeError;
 #[repr(u64)]
 pub enum BatchState {
     /// Batch can be filled with values.
-    CanBeFilled,
+    Fill,
     /// Batch has been inserted into the tree.
     Inserted,
     /// Batch is full, and insertion is in progress.
@@ -20,7 +20,7 @@ pub enum BatchState {
 impl From<u64> for BatchState {
     fn from(value: u64) -> Self {
         match value {
-            0 => BatchState::CanBeFilled,
+            0 => BatchState::Fill,
             1 => BatchState::Inserted,
             2 => BatchState::Full,
             _ => panic!("Invalid BatchState value"),
@@ -73,7 +73,7 @@ impl Batch {
             bloom_filter_capacity,
             batch_size,
             num_inserted: 0,
-            state: BatchState::CanBeFilled.into(),
+            state: BatchState::Fill.into(),
             zkp_batch_size,
             current_zkp_batch_index: 0,
             num_inserted_zkps: 0,
@@ -102,9 +102,9 @@ impl Batch {
     }
 
     /// fill -> full -> inserted -> fill
-    pub fn advance_state_to_can_be_filled(&mut self) -> Result<(), BatchedMerkleTreeError> {
+    pub fn advance_state_to_fill(&mut self) -> Result<(), BatchedMerkleTreeError> {
         if self.get_state() == BatchState::Inserted {
-            self.state = BatchState::CanBeFilled.into();
+            self.state = BatchState::Fill.into();
         } else {
             msg!(
                 "Batch is in incorrect state {} expected Inserted 3",
@@ -131,7 +131,7 @@ impl Batch {
 
     /// fill -> full -> inserted -> fill
     pub fn advance_state_to_full(&mut self) -> Result<(), BatchedMerkleTreeError> {
-        if self.get_state() == BatchState::CanBeFilled {
+        if self.get_state() == BatchState::Fill {
             self.state = BatchState::Full.into();
         } else {
             msg!(
@@ -174,7 +174,7 @@ impl Batch {
         value: &[u8; 32],
         value_store: &mut ZeroCopyVecU64<[u8; 32]>,
     ) -> Result<(), BatchedMerkleTreeError> {
-        if self.get_state() != BatchState::CanBeFilled {
+        if self.get_state() != BatchState::Fill {
             return Err(BatchedMerkleTreeError::BatchNotReady);
         }
         value_store.push(*value)?;
@@ -253,31 +253,39 @@ impl Batch {
         self.batch_size / self.zkp_batch_size
     }
 
+    /// Marks the batch as inserted in the merkle tree.
+    /// 1. Checks that the batch is ready.
+    /// 2. increments the number of inserted zkps.
+    /// 3. If all zkps are inserted, sets the state to inserted.
+    /// 4. Returns the updated state of the batch.
     pub fn mark_as_inserted_in_merkle_tree(
         &mut self,
         sequence_number: u64,
         root_index: u32,
         root_history_length: u32,
-    ) -> Result<(), BatchedMerkleTreeError> {
-        // Check that batch is ready.
+    ) -> Result<BatchState, BatchedMerkleTreeError> {
+        // 1. Check that batch is ready.
         self.get_first_ready_zkp_batch()?;
 
         let num_zkp_batches = self.get_num_zkp_batches();
 
+        // 2. increments the number of inserted zkps.
         self.num_inserted_zkps += 1;
-        // Batch has been successfully inserted into the tree.
+        println!("num_inserted_zkps: {}", self.num_inserted_zkps);
+        // 3. If all zkps are inserted, sets the state to inserted.
         if self.num_inserted_zkps == num_zkp_batches {
-            self.current_zkp_batch_index = 0;
-            self.state = BatchState::Inserted.into();
+            println!("Setting state to inserted");
             self.num_inserted_zkps = 0;
+            self.current_zkp_batch_index = 0;
+            self.advance_state_to_inserted()?;
             // Saving sequence number and root index for the batch.
             // When the batch is cleared check that sequence number is greater or equal than self.sequence_number
             // if not advance current root index to root index
             self.sequence_number = sequence_number + root_history_length as u64;
             self.root_index = root_index;
         }
 
-        Ok(())
+        Ok(self.get_state())
     }
 
     pub fn get_hashchain_store_len(&self) -> usize {
@@ -513,7 +521,7 @@ mod tests {
         let mut batch = get_test_batch();
         assert_eq!(batch.get_num_zkp_batches(), 5);
         assert_eq!(batch.get_hashchain_store_len(), 5);
-        assert_eq!(batch.get_state(), BatchState::CanBeFilled);
+        assert_eq!(batch.get_state(), BatchState::Fill);
         assert_eq!(batch.get_num_inserted(), 0);
         assert_eq!(batch.get_current_zkp_batch_index(), 0);
         assert_eq!(batch.get_num_inserted_zkps(), 0);

program-libs/batched-merkle-tree/src/batch_metadata.rs

@@ -1,6 +1,7 @@
+use light_merkle_tree_metadata::{errors::MerkleTreeMetadataError, queue::QueueType};
 use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};
 
-use crate::{BorshDeserialize, BorshSerialize};
+use crate::{batch::BatchState, errors::BatchedMerkleTreeError, BorshDeserialize, BorshSerialize};
 
 #[repr(C)]
 #[derive(
@@ -17,43 +18,126 @@ use crate::{BorshDeserialize, BorshSerialize};
     Immutable,
 )]
 pub struct BatchMetadata {
+    /// Number of batches.
     pub num_batches: u64,
+    /// Number of elements in a batch.
     pub batch_size: u64,
+    /// Number of elements in a ZKP batch.
+    /// A batch has one or more ZKP batches.
     pub zkp_batch_size: u64,
+    /// Bloom filter capacity.
+    pub bloom_filter_capacity: u64,
+    /// Batch elements are currently inserted in.
     pub currently_processing_batch_index: u64,
+    /// Next batch to be inserted into the tree.
     pub next_full_batch_index: u64,
-    pub bloom_filter_capacity: u64,
 }
 
 impl BatchMetadata {
     pub fn get_num_zkp_batches(&self) -> u64 {
         self.batch_size / self.zkp_batch_size
     }
 
-    pub fn new_output_queue(batch_size: u64, zkp_batch_size: u64, num_batches: u64) -> Self {
-        BatchMetadata {
+    pub fn new_output_queue(
+        batch_size: u64,
+        zkp_batch_size: u64,
+        num_batches: u64,
+    ) -> Result<Self, BatchedMerkleTreeError> {
+        if batch_size % zkp_batch_size != 0 {
+            return Err(BatchedMerkleTreeError::BatchSizeNotDivisibleByZkpBatchSize);
+        }
+        Ok(BatchMetadata {
             num_batches,
             zkp_batch_size,
             batch_size,
             currently_processing_batch_index: 0,
             next_full_batch_index: 0,
             bloom_filter_capacity: 0,
-        }
+        })
     }
 
     pub fn new_input_queue(
         batch_size: u64,
         bloom_filter_capacity: u64,
         zkp_batch_size: u64,
         num_batches: u64,
-    ) -> Self {
-        BatchMetadata {
+    ) -> Result<Self, BatchedMerkleTreeError> {
+        if batch_size % zkp_batch_size != 0 {
+            return Err(BatchedMerkleTreeError::BatchSizeNotDivisibleByZkpBatchSize);
+        }
+        Ok(BatchMetadata {
             num_batches,
             zkp_batch_size,
             batch_size,
             currently_processing_batch_index: 0,
             next_full_batch_index: 0,
             bloom_filter_capacity,
+        })
+    }
+
+    /// Increment the next full batch index if current state is inserted.
+    pub fn increment_next_full_batch_index_if_inserted(&mut self, state: BatchState) {
+        if state == BatchState::Inserted {
+            self.next_full_batch_index += 1;
+            self.next_full_batch_index %= self.num_batches;
         }
     }
+
+    pub fn init(
+        &mut self,
+        num_batches: u64,
+        batch_size: u64,
+        zkp_batch_size: u64,
+    ) -> Result<(), BatchedMerkleTreeError> {
+        self.num_batches = num_batches;
+        self.batch_size = batch_size;
+        // Check that batch size is divisible by zkp_batch_size.
+        if batch_size % zkp_batch_size != 0 {
+            return Err(BatchedMerkleTreeError::BatchSizeNotDivisibleByZkpBatchSize);
+        }
+        self.zkp_batch_size = zkp_batch_size;
+        Ok(())
+    }
+
+    pub fn get_size_parameters(
+        &self,
+        queue_type: u64,
+    ) -> Result<(usize, usize, usize), MerkleTreeMetadataError> {
+        let num_batches = self.num_batches as usize;
+        // Input queues don't store values.
+        let num_value_stores = if queue_type == QueueType::BatchedOutput as u64 {
+            num_batches
+        } else if queue_type == QueueType::BatchedInput as u64 {
+            0
+        } else {
+            return Err(MerkleTreeMetadataError::InvalidQueueType);
+        };
+        // Output queues don't use bloom filters.
+        let num_stores = if queue_type == QueueType::BatchedInput as u64 {
+            num_batches
+        } else if queue_type == QueueType::BatchedOutput as u64 && self.bloom_filter_capacity == 0 {
+            0
+        } else {
+            return Err(MerkleTreeMetadataError::InvalidQueueType);
+        };
+        Ok((num_value_stores, num_stores, num_batches))
+    }
+}
+
+#[test]
+fn test_increment_next_full_batch_index_if_inserted() {
+    // create a new metadata struct
+    let mut metadata = BatchMetadata::new_input_queue(10, 10, 10, 2).unwrap();
+    assert_eq!(metadata.next_full_batch_index, 0);
+    // increment next full batch index
+    metadata.increment_next_full_batch_index_if_inserted(BatchState::Inserted);
+    assert_eq!(metadata.next_full_batch_index, 1);
+    // increment next full batch index
+    metadata.increment_next_full_batch_index_if_inserted(BatchState::Inserted);
+    assert_eq!(metadata.next_full_batch_index, 0);
+    // try incrementing next full batch index with state not inserted
+    metadata.increment_next_full_batch_index_if_inserted(BatchState::Fill);
+    assert_eq!(metadata.next_full_batch_index, 0);
+    metadata.increment_next_full_batch_index_if_inserted(BatchState::Full);
+    assert_eq!(metadata.next_full_batch_index, 0);
 }

program-libs/batched-merkle-tree/src/event.rs

@@ -25,3 +25,5 @@ pub struct BatchNullifyEvent {
     pub sequence_number: u64,
     pub batch_size: u64,
 }
+
+pub type BatchAddressAppendEvent = BatchNullifyEvent;

program-libs/batched-merkle-tree/src/initialize_state_tree.rs

@@ -477,7 +477,8 @@ pub fn create_output_queue_account(params: CreateOutputQueueParams) -> BatchedQu
         params.batch_size,
         params.zkp_batch_size,
         params.num_batches,
-    );
+    )
+    .unwrap();
     BatchedQueueMetadata {
         metadata,
         batch_metadata,