wip

Lightprotocol · Nov 27, 2024 · 2db5b62 · 2db5b62
1 parent d7fdd1d
commit 2db5b62
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 11 deletions.
diff --git a/forester-utils/src/indexer/mod.rs b/forester-utils/src/indexer/mod.rs
@@ -173,6 +173,11 @@ pub trait Indexer<R: RpcConnection>: Sync + Send + Debug {
     ) {
         unimplemented!()
     }
+
+    async fn update_test_indexer_in_nullification(&mut self, _merkle_tree_pubkey: Pubkey, _nullifier: &[u8; 32], _index: usize) {
+        unimplemented!()
+    }
+
 }
 
 #[derive(Debug, Clone)]

diff --git a/forester/src/batched_ops.rs b/forester/src/batched_ops.rs
@@ -33,6 +33,7 @@ use solana_sdk::signer::Signer;
 use std::sync::Arc;
 use tokio::sync::Mutex;
 use tracing::{debug, error, info};
+use account_compression::batch::BatchState;
 
 pub struct BatchedOperations<R: RpcConnection, I: Indexer<R>> {
     pub rpc_pool: Arc<SolanaRpcPool<R>>,
@@ -138,15 +139,15 @@ impl<R: RpcConnection, I: Indexer<R>> BatchedOperations<R, I> {
             )
         };
 
-        self.indexer
-            .lock()
-            .await
-            .update_test_indexer_after_nullification(
-                &mut rpc,
-                self.merkle_tree,
-                batch_index as usize,
-            )
-            .await;
+        // self.indexer
+        //     .lock()
+        //     .await
+        //     .update_test_indexer_after_nullification(
+        //         &mut rpc,
+        //         self.merkle_tree,
+        //         batch_index as usize,
+        //     )
+        //     .await;
 
         info!("Batch nullify completed successfully: {:?}", result);
         Ok(batch_size as usize)
@@ -292,8 +293,7 @@ impl<R: RpcConnection, I: Indexer<R>> BatchedOperations<R, I> {
     async fn get_batched_nullify_ix_data(&self) -> Result<InstructionDataBatchNullifyInputs> {
         let mut rpc = self.rpc_pool.get_connection().await.unwrap();
 
-        // Get merkle tree data
-        let (zkp_batch_size, _batch_index, old_root, old_root_index, leaves_hashchain) = {
+        let (zkp_batch_size, batch_index, old_root, old_root_index, leaves_hashchain) = {
             let mut account = rpc.get_account(self.merkle_tree).await.unwrap().unwrap();
             let merkle_tree =
                 ZeroCopyBatchedMerkleTreeAccount::from_bytes_mut(account.data.as_mut_slice())
@@ -322,6 +322,15 @@ impl<R: RpcConnection, I: Indexer<R>> BatchedOperations<R, I> {
         let mut path_indices = Vec::new();
         let mut merkle_proofs = Vec::new();
 
+        let batch = {
+            let mut merkle_tree_account = rpc.get_account(self.merkle_tree).await.unwrap().unwrap();
+            let merkle_tree = ZeroCopyBatchedMerkleTreeAccount::from_bytes_mut(
+                merkle_tree_account.data.as_mut_slice(),
+            )
+                .unwrap();
+            merkle_tree.batches[batch_index].clone()
+        };
+
         for (index, leaf, tx_hash) in leaf_indices_tx_hashes.iter() {
             path_indices.push(*index);
             leaves.push(*leaf);
@@ -333,6 +342,17 @@ impl<R: RpcConnection, I: Indexer<R>> BatchedOperations<R, I> {
                 .get_proof_by_index(self.merkle_tree, *index as usize);
             old_leaves.push(proof.leaf);
             merkle_proofs.push(proof.proof);
+
+            tx_hashes.push(*tx_hash);
+
+            if batch.get_state() == BatchState::Inserted
+                || batch.get_state() == BatchState::ReadyToUpdateTree {
+                let index_bytes = index.to_be_bytes();
+                use light_hasher::Hasher;
+                let leaf = *leaf;
+                let nullifier = Poseidon::hashv(&[&leaf, &index_bytes, tx_hash]).unwrap();
+                self.indexer.lock().await.update_test_indexer_in_nullification(self.merkle_tree, &nullifier, *index as usize).await;
+            }
         }
 
         let inputs = get_batch_update_inputs::<26>(

diff --git a/test-utils/src/indexer/test_indexer.rs b/test-utils/src/indexer/test_indexer.rs
@@ -738,6 +738,15 @@ impl<R: RpcConnection + Send + Sync + 'static> Indexer<R> for TestIndexer<R> {
         }
     }
 
+    async fn update_test_indexer_in_nullification(&mut self, merkle_tree_pubkey: Pubkey, nullifier: &[u8; 32], index: usize) {
+        let mut state_merkle_tree_bundle = self
+            .state_merkle_trees
+            .iter_mut()
+            .find(|x| x.accounts.merkle_tree == merkle_tree_pubkey)
+            .unwrap();
+        state_merkle_tree_bundle.merkle_tree.update(&nullifier, index).unwrap();
+    }
+
     async fn update_test_indexer_after_nullification(
         &mut self,
         rpc: &mut R,