refactor: get_queue_elements (#1595)

* refactor: get_queue_elements

* chore: fix test

* fix: batch address test

* fix: account compression tests

Author: ananas-block

forester-utils/src/instructions.rs

@@ -59,7 +59,7 @@ where
         })?
         .unwrap();
 
-    let (leaves_hash_chain, start_index, current_root, batch_size, full_batch_index) = {
+    let (leaves_hash_chain, start_index, current_root, batch_size) = {
         let merkle_tree = BatchedMerkleTreeAccount::address_from_bytes(
             merkle_tree_account.data.as_mut_slice(),
             &merkle_tree_pubkey.into(),
@@ -75,13 +75,7 @@ where
         let current_root = *merkle_tree.root_history.last().unwrap();
         let batch_size = batch.zkp_batch_size as usize;
 
-        (
-            leaves_hash_chain,
-            start_index,
-            current_root,
-            batch_size,
-            full_batch_index,
-        )
+        (leaves_hash_chain, start_index, current_root, batch_size)
     };
 
     let batch_start_index = indexer
@@ -96,9 +90,8 @@ where
     let addresses = indexer
         .get_queue_elements(
             merkle_tree_pubkey.to_bytes(),
-            full_batch_index,
-            0,
             batch_size as u64,
+            None
         )
         .await
         .map_err(|e| {
@@ -115,7 +108,7 @@ where
     let non_inclusion_proofs = indexer
         .get_multiple_new_address_proofs_h40(
             merkle_tree_pubkey.to_bytes(),
-            addresses.clone(),
+            addresses.iter().map(|x|x.account_hash).collect(),
         )
         .await
         .map_err(|e| {
@@ -151,6 +144,10 @@ where
         })?
         .try_into()
         .unwrap();
+    let addresses = addresses
+        .iter()
+        .map(|x| x.account_hash)
+        .collect::<Vec<[u8; 32]>>();
 
     let inputs =
         get_batch_address_append_circuit_inputs::<{ DEFAULT_BATCH_ADDRESS_TREE_HEIGHT as usize }>(
@@ -227,7 +224,7 @@ pub async fn create_append_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
         )
     };
 
-    let (zkp_batch_size, full_batch_index, num_inserted_zkps, leaves_hash_chain) = {
+    let (zkp_batch_size, leaves_hash_chain) = {
         let mut output_queue_account = rpc.get_account(output_queue_pubkey).await.unwrap().unwrap();
         let output_queue =
             BatchedQueueAccount::output_from_bytes(output_queue_account.data.as_mut_slice())
@@ -242,41 +239,27 @@ pub async fn create_append_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
         let leaves_hash_chain =
             output_queue.hash_chain_stores[full_batch_index as usize][num_inserted_zkps as usize];
 
-        (
-            zkp_batch_size,
-            full_batch_index,
-            num_inserted_zkps,
-            leaves_hash_chain,
-        )
+        (zkp_batch_size, leaves_hash_chain)
     };
-    let start = num_inserted_zkps as usize * zkp_batch_size as usize;
-    let end = start + zkp_batch_size as usize;
 
-    let leaves = indexer
-        .get_queue_elements(
-            merkle_tree_pubkey.to_bytes(),
-            full_batch_index,
-            start as u64,
-            end as u64,
-        )
+    let indexer_response = indexer
+        .get_queue_elements(output_queue_pubkey.to_bytes(), zkp_batch_size, None)
         .await
         .unwrap();
 
-    info!("Leaves: {:?}", leaves);
-
-    let (old_leaves, merkle_proofs) = {
-        let mut old_leaves = vec![];
-        let mut merkle_proofs = vec![];
-        let indices =
-            (merkle_tree_next_index..merkle_tree_next_index + zkp_batch_size).collect::<Vec<_>>();
-        let proofs = indexer.get_proofs_by_indices(merkle_tree_pubkey, &indices);
-        proofs.iter().for_each(|proof| {
-            old_leaves.push(proof.leaf);
-            merkle_proofs.push(proof.proof.clone());
-        });
-
-        (old_leaves, merkle_proofs)
-    };
+    println!("indexer_response: {:?}", indexer_response);
+    let old_leaves = indexer_response
+        .iter()
+        .map(|x| x.leaf)
+        .collect::<Vec<[u8; 32]>>();
+    let leaves = indexer_response
+        .iter()
+        .map(|x| x.account_hash)
+        .collect::<Vec<[u8; 32]>>();
+    let merkle_proofs = indexer_response
+        .iter()
+        .map(|x| x.proof.clone())
+        .collect::<Vec<Vec<[u8; 32]>>>();
 
     info!("Old leaves: {:?}", old_leaves);
 
@@ -356,8 +339,11 @@ pub async fn create_nullify_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
         (zkp_size, root, hash_chain)
     };
 
-    let leaf_indices_tx_hashes =
-        indexer.get_leaf_indices_tx_hashes(merkle_tree_pubkey, zkp_batch_size as usize);
+    let leaf_indices_tx_hashes = indexer
+        .get_queue_elements(merkle_tree_pubkey.to_bytes(), zkp_batch_size, None)
+        .await
+        .unwrap();
+    println!("leaf_indices_tx_hashes {:?}", leaf_indices_tx_hashes);
 
     let mut leaves = Vec::new();
     let mut tx_hashes = Vec::new();
@@ -366,23 +352,19 @@ pub async fn create_nullify_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
     let mut merkle_proofs = Vec::new();
     let mut nullifiers = Vec::new();
 
-    let proofs = indexer.get_proofs_by_indices(
-        merkle_tree_pubkey,
-        &leaf_indices_tx_hashes
-            .iter()
-            .map(|leaf_info| leaf_info.leaf_index as u64)
-            .collect::<Vec<_>>(),
-    );
-
-    for (leaf_info, proof) in leaf_indices_tx_hashes.iter().zip(proofs.iter()) {
-        path_indices.push(leaf_info.leaf_index);
-        leaves.push(leaf_info.leaf);
-        old_leaves.push(proof.leaf);
-        merkle_proofs.push(proof.proof.clone());
-        tx_hashes.push(leaf_info.tx_hash);
+    for leaf_info in leaf_indices_tx_hashes.iter() {
+        path_indices.push(leaf_info.leaf_index as u32);
+        leaves.push(leaf_info.account_hash);
+        old_leaves.push(leaf_info.leaf);
+        merkle_proofs.push(leaf_info.proof.clone());
+        tx_hashes.push(leaf_info.tx_hash.unwrap());
         let index_bytes = leaf_info.leaf_index.to_be_bytes();
-        let nullifier =
-            Poseidon::hashv(&[&leaf_info.leaf, &index_bytes, &leaf_info.tx_hash]).unwrap();
+        let nullifier = Poseidon::hashv(&[
+            &leaf_info.account_hash,
+            &index_bytes,
+            &leaf_info.tx_hash.unwrap(),
+        ])
+        .unwrap();
         nullifiers.push(nullifier);
     }
 

forester/src/batch_processor/common.rs

@@ -158,8 +158,8 @@ impl<R: RpcConnection, I: Indexer<R> + IndexerType<R>> BatchProcessor<R, I> {
 
     async fn process_state_append(&self) -> Result<usize> {
         let mut rpc = self.context.rpc_pool.get_connection().await?;
-        let (num_inserted_zkps, zkp_batch_size) = self.get_num_inserted_zkps(&mut rpc).await?;
-        state::perform_append(&self.context, &mut rpc, num_inserted_zkps).await?;
+        let (_, zkp_batch_size) = self.get_num_inserted_zkps(&mut rpc).await?;
+        state::perform_append(&self.context, &mut rpc).await?;
         Ok(zkp_batch_size)
     }
 

forester/src/batch_processor/state.rs

@@ -18,7 +18,6 @@ use crate::{
 pub(crate) async fn perform_append<R: RpcConnection, I: Indexer<R> + IndexerType<R>>(
     context: &BatchContext<R, I>,
     rpc: &mut R,
-    num_inserted_zkps: u64,
 ) -> Result<()> {
     let instruction_data = create_append_batch_ix_data(
         rpc,
@@ -53,7 +52,6 @@ pub(crate) async fn perform_append<R: RpcConnection, I: Indexer<R> + IndexerType
         context.indexer.clone(),
         context.merkle_tree,
         context.output_queue,
-        num_inserted_zkps,
     )
     .await
     .expect("Failed to update test indexer after append");

forester/src/indexer_type.rs

@@ -69,7 +69,6 @@ pub trait IndexerType<R: RpcConnection>: sealed::Sealed {
         indexer: &mut impl Indexer<R>,
         merkle_tree_pubkey: Pubkey,
         output_queue: Pubkey,
-        num_inserted_zkps: u64,
     ) where
         Self: Sized;
 }
@@ -150,18 +149,12 @@ impl<R: RpcConnection + light_client::rpc::merkle_tree::MerkleTreeExt> IndexerTy
         indexer: &mut impl Indexer<R>,
         merkle_tree_pubkey: Pubkey,
         output_queue: Pubkey,
-        num_inserted_zkps: u64,
     ) where
         Self: Sized,
     {
         if let Some(test_indexer) = (indexer as &mut dyn Any).downcast_mut::<TestIndexer<R>>() {
             test_indexer
-                .update_test_indexer_after_append(
-                    rpc,
-                    merkle_tree_pubkey,
-                    output_queue,
-                    num_inserted_zkps,
-                )
+                .update_test_indexer_after_append(rpc, merkle_tree_pubkey, output_queue)
                 .await;
         }
     }
@@ -209,7 +202,6 @@ impl<R: RpcConnection> IndexerType<R> for PhotonIndexer<R> {
         _indexer: &mut impl Indexer<R>,
         _merkle_tree_pubkey: Pubkey,
         _output_queue: Pubkey,
-        _num_inserted_zkps: u64,
     ) {
         // No-op for production indexer
     }
@@ -328,14 +320,12 @@ pub async fn update_test_indexer_after_append<R: RpcConnection, I: Indexer<R> +
     indexer: Arc<Mutex<I>>,
     merkle_tree_pubkey: Pubkey,
     output_queue: Pubkey,
-    num_inserted_zkps: u64,
 ) -> Result<(), ForesterError> {
     I::update_test_indexer_after_append(
         &mut *rpc,
         &mut *indexer.lock().await,
         merkle_tree_pubkey,
         output_queue,
-        num_inserted_zkps,
     )
     .await;
 

program-libs/batched-merkle-tree/tests/merkle_tree.rs

@@ -1273,7 +1273,7 @@ pub async fn perform_address_update(
         &mt_pubkey,
     )
     .unwrap();
-    let (input_res, new_root, pre_next_full_batch) = {
+    let (input_res, new_root, _pre_next_full_batch) = {
         let mut account =
             BatchedMerkleTreeAccount::address_from_bytes(mt_account_data, &mt_pubkey).unwrap();
 
@@ -1285,7 +1285,9 @@ pub async fn perform_address_update(
             .batches
             .get(next_full_batch as usize)
             .unwrap();
-        let batch_start_index = batch.start_index;
+        let batch_start_index =
+            batch.start_index + batch.get_num_inserted_zkps() * batch.zkp_batch_size;
+        println!("batch start index {}", batch_start_index);
         let leaves_hash_chain = account
             .hash_chain_stores
             .get(next_full_batch as usize)
@@ -1298,12 +1300,14 @@ pub async fn perform_address_update(
                 account.get_metadata().queue_batches.batch_size as u32,
                 account.get_metadata().queue_batches.zkp_batch_size as u32,
                 *leaves_hash_chain,
-                next_index as usize,
+                next_index as usize, // % batch.zkp_batch_size as usize
                 batch_start_index as usize,
                 *current_root,
             )
             .await
             .unwrap();
+
+        mock_indexer.finalize_batch_address_update(10);
         let instruction_data = InstructionDataBatchNullifyInputs {
             new_root,
             compressed_proof: CompressedProof {
@@ -1334,16 +1338,6 @@ pub async fn perform_address_update(
     let account =
         BatchedMerkleTreeAccount::address_from_bytes(mt_account_data, &mt_pubkey).unwrap();
 
-    {
-        let batch = account
-            .queue_batches
-            .batches
-            .get(pre_next_full_batch as usize)
-            .unwrap();
-        if batch.get_state() == BatchState::Inserted {
-            mock_indexer.finalize_batch_address_update(batch.batch_size as usize);
-        }
-    }
     assert_address_merkle_tree_update(old_account, account, new_root);
 }
 

program-tests/account-compression-test/tests/batched_merkle_tree_test.rs

@@ -1699,6 +1699,7 @@ async fn test_batch_address_merkle_trees() {
         )
         .await
         .unwrap();
+        mock_indexer.finalize_batch_address_update(10);
     }
     // 5. Failing: invalid proof
     // 6. Failing: invalid new root
@@ -2027,7 +2028,8 @@ pub async fn update_batch_address_tree(
         .batches
         .get(next_full_batch as usize)
         .unwrap();
-    let batch_start_index = batch.start_index;
+    let batch_start_index =
+        batch.start_index + batch.get_num_inserted_zkps() * batch.zkp_batch_size;
     let leaves_hash_chain = zero_copy_account
         .hash_chain_stores
         .get(next_full_batch as usize)

program-tests/registry-test/tests/tests.rs

@@ -1918,18 +1918,18 @@ async fn test_batch_address_tree() {
             )
             .await
             .unwrap();
+            let mut account = rpc
+                .get_account(env.batch_address_merkle_tree)
+                .await
+                .unwrap()
+                .unwrap();
+            test_indexer
+                .finalize_batched_address_tree_update(
+                    env.batch_address_merkle_tree,
+                    account.data.as_mut_slice(),
+                )
+                .await;
         }
-        let mut account = rpc
-            .get_account(env.batch_address_merkle_tree)
-            .await
-            .unwrap()
-            .unwrap();
-        test_indexer
-            .finalize_batched_address_tree_update(
-                env.batch_address_merkle_tree,
-                account.data.as_mut_slice(),
-            )
-            .await;
     }
 
     // Non eligible forester.
@@ -1966,6 +1966,17 @@ async fn test_batch_address_tree() {
         )
         .await
         .unwrap();
+        let mut account = rpc
+            .get_account(env.batch_address_merkle_tree)
+            .await
+            .unwrap()
+            .unwrap();
+        test_indexer
+            .finalize_batched_address_tree_update(
+                env.batch_address_merkle_tree,
+                account.data.as_mut_slice(),
+            )
+            .await;
     }
     let mut account = rpc
         .get_account(env.batch_address_merkle_tree)

program-tests/system-cpi-test/tests/test.rs

@@ -146,7 +146,7 @@ async fn test_read_only_accounts() {
             // fails because of invalid leaves hash_chain in some iteration
             let instruction_data = create_batch_update_address_tree_instruction_data_with_proof(
                 &mut e2e_env.rpc,
-                &e2e_env.indexer,
+                &mut e2e_env.indexer,
                 env.batch_address_merkle_tree,
             )
             .await

program-tests/utils/src/e2e_test_env.rs

@@ -762,12 +762,13 @@ where
                                     .indexer
                                     .get_queue_elements(
                                         merkle_tree_pubkey.to_bytes(),
-                                        full_batch_index,
-                                        0,
                                         batch.batch_size,
+                                        None,
                                     )
                                     .await
                                     .unwrap();
+                                let addresses =
+                                    addresses.iter().map(|x| x.account_hash).collect::<Vec<_>>();
                                 // // local_leaves_hash_chain is only used for a test assertion.
                                 // let local_nullifier_hash_chain = create_hash_chain_from_array(&addresses);
                                 // assert_eq!(leaves_hash_chain, local_nullifier_hash_chain);