feat: read only accounts

circuit-lib/verifier/src/lib.rs

@@ -173,6 +173,7 @@ pub fn verify_merkle_proof_zkp(
     leaves: &[[u8; 32]],
     compressed_proof: &CompressedProof,
 ) -> Result<(), VerifierError> {
+    // TODO: replace with poseidon hash which avoids copying the data
     let public_inputs = [roots, leaves].concat();
 
     // The public inputs are expected to be a multiple of 2

programs/account-compression/src/state/batch.rs

@@ -233,7 +233,11 @@ impl Batch {
         let num_zkp_batches = self.get_num_zkp_batches();
 
         self.num_inserted_zkps += 1;
-
+        msg!(
+            "Marking batch as inserted in the merkle tree. num_inserted_zkps: {}",
+            self.num_inserted_zkps
+        );
+        msg!("num_zkp_batches: {}", num_zkp_batches);
         // Batch has been successfully inserted into the tree.
         if self.num_inserted_zkps == num_zkp_batches {
             self.current_zkp_batch_index = 0;

programs/account-compression/src/state/batched_queue.rs

@@ -274,21 +274,36 @@ impl ZeroCopyBatchedQueueAccount {
         Ok(())
     }
 
-    pub fn prove_inclusion_by_index_and_zero_out_leaf(
-        &mut self,
-        leaf_index: u64,
-        value: &[u8; 32],
-    ) -> Result<()> {
-        self.prove_inclusion_by_index_option_zero_out::<true>(leaf_index, value)
+    pub fn prove_inclusion_by_index(&mut self, leaf_index: u64, value: &[u8; 32]) -> Result<bool> {
+        for (batch_index, batch) in self.batches.iter().enumerate() {
+            if batch.value_is_inserted_in_batch(leaf_index)? {
+                let index = batch.get_value_index_in_batch(leaf_index)?;
+                let element = self.value_vecs[batch_index]
+                    .get_mut(index as usize)
+                    .ok_or(AccountCompressionErrorCode::InclusionProofByIndexFailed)?;
+
+                if element == value {
+                    return Ok(true);
+                } else {
+                    return err!(AccountCompressionErrorCode::InclusionProofByIndexFailed);
+                }
+            }
+        }
+        Ok(false)
     }
 
-    pub fn prove_inclusion_by_index(&mut self, leaf_index: u64, value: &[u8; 32]) -> Result<()> {
-        self.prove_inclusion_by_index_option_zero_out::<false>(leaf_index, value)
+    pub fn could_exist_in_batches(&mut self, leaf_index: u64) -> Result<()> {
+        for batch in self.batches.iter() {
+            if batch.value_is_inserted_in_batch(leaf_index)? {
+                return Ok(());
+            }
+        }
+        err!(AccountCompressionErrorCode::InclusionProofByIndexFailed)
     }
 
     /// Zero out a leaf by index if it exists in the queues value vec. If
     /// checked fail if leaf is not found.
-    fn prove_inclusion_by_index_option_zero_out<const ZERO_OUT_LEAF: bool>(
+    pub fn prove_inclusion_by_index_and_zero_out_leaf(
         &mut self,
         leaf_index: u64,
         value: &[u8; 32],
@@ -301,16 +316,14 @@ impl ZeroCopyBatchedQueueAccount {
                     .ok_or(AccountCompressionErrorCode::InclusionProofByIndexFailed)?;
 
                 if element == value {
-                    if ZERO_OUT_LEAF {
-                        *element = [0u8; 32];
-                    }
+                    *element = [0u8; 32];
                     return Ok(());
                 } else {
                     return err!(AccountCompressionErrorCode::InclusionProofByIndexFailed);
                 }
             }
         }
-        err!(AccountCompressionErrorCode::InclusionProofByIndexFailed)
+        Ok(())
     }
 
     pub fn get_batch_num_inserted_in_current_batch(&self) -> u64 {
@@ -365,6 +378,7 @@ pub fn insert_into_current_batch(
         println!("value {:?}", value);
 
         if wipe {
+            msg!("wipe");
             if let Some(blomfilter_stores) = bloom_filter_stores.as_mut() {
                 if !current_batch.bloom_filter_is_wiped {
                     (*blomfilter_stores)
@@ -811,32 +825,39 @@ pub mod tests {
         {
             zero_copy_account.insert_into_current_batch(&value).unwrap();
             assert_eq!(
-                zero_copy_account.prove_inclusion_by_index_option_zero_out::<false>(1, &value),
+                zero_copy_account.prove_inclusion_by_index(1, &value),
                 anchor_lang::err!(AccountCompressionErrorCode::InclusionProofByIndexFailed)
             );
             assert_eq!(
-                zero_copy_account.prove_inclusion_by_index_option_zero_out::<true>(1, &value),
+                zero_copy_account.prove_inclusion_by_index_and_zero_out_leaf(1, &value),
                 anchor_lang::err!(AccountCompressionErrorCode::InclusionProofByIndexFailed)
             );
             assert_eq!(
-                zero_copy_account.prove_inclusion_by_index_option_zero_out::<true>(0, &value2),
+                zero_copy_account.prove_inclusion_by_index_and_zero_out_leaf(0, &value2),
                 anchor_lang::err!(AccountCompressionErrorCode::InclusionProofByIndexFailed)
             );
             assert!(zero_copy_account
-                .prove_inclusion_by_index_option_zero_out::<false>(0, &value)
+                .prove_inclusion_by_index(0, &value)
                 .is_ok());
+            // prove inclusion for value out of range returns false
+            assert_eq!(
+                zero_copy_account
+                    .prove_inclusion_by_index(100000, &[0u8; 32])
+                    .unwrap(),
+                false
+            );
             assert!(zero_copy_account
-                .prove_inclusion_by_index_option_zero_out::<true>(0, &value)
+                .prove_inclusion_by_index_and_zero_out_leaf(0, &value)
                 .is_ok());
         }
         // 2. Functional does not succeed on second invocation
         {
             assert_eq!(
-                zero_copy_account.prove_inclusion_by_index_option_zero_out::<true>(0, &value),
+                zero_copy_account.prove_inclusion_by_index_and_zero_out_leaf(0, &value),
                 anchor_lang::err!(AccountCompressionErrorCode::InclusionProofByIndexFailed)
             );
             assert_eq!(
-                zero_copy_account.prove_inclusion_by_index_option_zero_out::<false>(0, &value),
+                zero_copy_account.prove_inclusion_by_index(0, &value),
                 anchor_lang::err!(AccountCompressionErrorCode::InclusionProofByIndexFailed)
             );
         }
@@ -848,17 +869,17 @@ pub mod tests {
                 .unwrap();
 
             assert_eq!(
-                zero_copy_account.prove_inclusion_by_index_option_zero_out::<true>(0, &value2),
+                zero_copy_account.prove_inclusion_by_index_and_zero_out_leaf(0, &value2),
                 anchor_lang::err!(AccountCompressionErrorCode::InclusionProofByIndexFailed)
             );
             assert!(zero_copy_account
-                .prove_inclusion_by_index_option_zero_out::<true>(1, &value2)
+                .prove_inclusion_by_index_and_zero_out_leaf(1, &value2)
                 .is_ok());
         }
         // 4. Functional does not succeed on second invocation
         {
             assert_eq!(
-                zero_copy_account.prove_inclusion_by_index_option_zero_out::<true>(1, &value2),
+                zero_copy_account.prove_inclusion_by_index_and_zero_out_leaf(1, &value2),
                 anchor_lang::err!(AccountCompressionErrorCode::InclusionProofByIndexFailed)
             );
         }

programs/system/src/errors.rs

@@ -69,4 +69,5 @@ pub enum SystemProgramError {
     OutputMerkleTreeNotUnique,
     DataFieldUndefined,
     ReadOnlyAddressAlreadyExists,
+    ReadOnlyAccountDoesNotExist,
 }

programs/system/src/invoke/instruction.rs

@@ -136,7 +136,7 @@ pub struct PackedReadOnlyAddress {
 }
 
 #[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
-pub struct ReadOnlyAddressParams {
+pub struct ReadOnlyAddress {
     pub address: [u8; 32],
     pub address_merkle_tree_pubkey: Pubkey,
     pub address_merkle_tree_root_index: u16,

programs/system/src/invoke/nullify_state.rs

@@ -24,10 +24,9 @@ pub fn insert_nullifiers<
 >(
     input_compressed_accounts_with_merkle_context: &'a [PackedCompressedAccountWithMerkleContext],
     ctx: &'a Context<'a, 'b, 'c, 'info, A>,
-    input_compressed_account_hashes: &'a [[u8; 32]],
+    nullifiers: &'a [[u8; 32]],
     invoking_program: &Option<Pubkey>,
     tx_hash: [u8; 32],
-    num_read_only: usize,
 ) -> Result<Option<(u8, u64)>> {
     light_heap::bench_sbf_start!("cpda_insert_nullifiers_prep_accs");
     let mut account_infos = vec![
@@ -57,19 +56,13 @@ pub fn insert_nullifiers<
     // nullifier queue is mutable. The network fee field in the queue is not
     // used.
     let mut network_fee_bundle = None;
-    let mut nullifiers =
-        Vec::with_capacity(input_compressed_accounts_with_merkle_context.len() - num_read_only);
-    for (account, account_hash) in input_compressed_accounts_with_merkle_context
-        .iter()
-        .zip(input_compressed_account_hashes)
-    {
+    for account in input_compressed_accounts_with_merkle_context.iter() {
         msg!("nullify state: account: {:?}", account);
         // Don't nullify read-only accounts.
         if account.read_only {
             continue;
         }
         leaf_indices.push(account.merkle_context.leaf_index);
-        nullifiers.push(*account_hash);
 
         let account_info =
             &ctx.remaining_accounts[account.merkle_context.nullifier_queue_pubkey_index as usize];
@@ -103,7 +96,7 @@ pub fn insert_nullifiers<
     light_heap::bench_sbf_start!("cpda_instruction_data");
 
     let instruction_data = account_compression::instruction::InsertIntoNullifierQueues {
-        nullifiers,
+        nullifiers: nullifiers.to_vec(),
         leaf_indices,
         tx_hash: Some(tx_hash),
     };

programs/system/src/invoke/processor.rs

@@ -14,11 +14,15 @@ use crate::{
         sum_check::sum_check,
         verify_state_proof::{
             create_tx_hash, fetch_input_compressed_account_roots, fetch_roots_address_merkle_tree,
-            hash_input_compressed_accounts, verify_read_only_account,
-            verify_read_only_address_queue_non_inclusion, verify_state_proof,
+            hash_input_compressed_accounts, verify_input_accounts_proof_by_index,
+            verify_read_only_account_inclusion, verify_read_only_address_queue_non_inclusion,
+            verify_state_proof,
         },
     },
-    sdk::accounts::{InvokeAccounts, SignerAccounts},
+    sdk::{
+        accounts::{InvokeAccounts, SignerAccounts},
+        compressed_account::PackedReadOnlyCompressedAccount,
+    },
     InstructionDataInvoke,
 };
 
@@ -61,13 +65,14 @@ pub fn process<
     ctx: Context<'a, 'b, 'c, 'info, A>,
     cpi_context_inputs: usize,
     read_only_addresses: Option<Vec<PackedReadOnlyAddress>>,
+    read_only_accounts: Option<Vec<PackedReadOnlyCompressedAccount>>,
 ) -> Result<()> {
     if inputs.relay_fee.is_some() {
         unimplemented!("Relay fee is not implemented yet.");
     }
     // Sum check ---------------------------------------------------
     bench_sbf_start!("cpda_sum_check");
-    let (num_read_only_input_accounts, num_prove_by_index_input_accounts) = sum_check(
+    let num_prove_by_index_input_accounts = sum_check(
         &inputs.input_compressed_accounts_with_merkle_context,
         &inputs.output_compressed_accounts,
         &inputs.relay_fee,
@@ -136,11 +141,6 @@ pub fn process<
             );
             return err!(SystemProgramError::InvalidCapacity);
         }
-        verify_read_only_account(
-            &ctx.remaining_accounts,
-            &inputs.input_compressed_accounts_with_merkle_context,
-            &input_compressed_account_hashes,
-        )?;
     }
 
     bench_sbf_end!("cpda_hash_input_compressed_accounts");
@@ -207,7 +207,7 @@ pub fn process<
     // in certain cases we zero out roots in batched input queues.
     // These roots need to be zero prior to proof verification.
     bench_sbf_start!("cpda_nullifiers");
-    let input_network_fee_bundle = if num_input_compressed_accounts > num_read_only_input_accounts {
+    let input_network_fee_bundle = if num_input_compressed_accounts != 0 {
         // Access the current slot
         let current_slot = Clock::get()?.slot;
         let tx_hash = create_tx_hash(
@@ -224,7 +224,6 @@ pub fn process<
             &input_compressed_account_hashes,
             &invoking_program,
             tx_hash,
-            num_read_only_input_accounts,
         )?
     } else {
         None
@@ -239,8 +238,14 @@ pub fn process<
         output_network_fee_bundle,
     )?;
 
+    // Proof inputs order:
+    // 1. input compressed accounts
+    // 2. read only compressed accounts
+    // 3. new addresses
+    // 4. read only addresses
     if num_prove_by_index_input_accounts < num_input_compressed_accounts
         || new_address_roots.capacity() != 0
+        || read_only_accounts.as_ref().map_or(false, |x| !x.is_empty())
     {
         bench_sbf_start!("cpda_verify_state_proof");
         if let Some(proof) = inputs.proof.as_ref() {
@@ -250,13 +255,26 @@ pub fn process<
                 b: proof.b,
                 c: proof.c,
             };
+            verify_input_accounts_proof_by_index(
+                &ctx.remaining_accounts,
+                &inputs.input_compressed_accounts_with_merkle_context,
+            )?;
             let mut input_compressed_account_roots =
                 Vec::with_capacity(num_input_compressed_accounts);
             fetch_input_compressed_account_roots(
                 &inputs.input_compressed_accounts_with_merkle_context,
                 &ctx,
                 &mut input_compressed_account_roots,
             )?;
+            let read_only_accounts = read_only_accounts.unwrap_or_default();
+            let mut read_only_accounts_roots = Vec::with_capacity(read_only_accounts.len());
+            fetch_input_compressed_account_roots(
+                &read_only_accounts,
+                &ctx,
+                &mut read_only_accounts_roots,
+            )?;
+            verify_read_only_account_inclusion(&ctx.remaining_accounts, &read_only_accounts)?;
+
             fetch_roots_address_merkle_tree(
                 &inputs.new_address_params,
                 &read_only_addresses,
@@ -272,12 +290,16 @@ pub fn process<
                 new_addresses.push(read_only_address.address);
             }
             msg!("read_only_addresses {:?}", read_only_addresses);
+            msg!("read_only_accounts {:?}", read_only_accounts);
+            msg!("read_only_accounts_roots {:?}", read_only_accounts_roots);
             match verify_state_proof(
                 &inputs.input_compressed_accounts_with_merkle_context,
-                &input_compressed_account_roots,
+                input_compressed_account_roots,
                 &input_compressed_account_hashes,
                 &new_address_roots,
                 &new_addresses,
+                &read_only_accounts,
+                &read_only_accounts_roots,
                 &compressed_verifier_proof,
             ) {
                 Ok(_) => Ok(()),
@@ -287,7 +309,7 @@ pub fn process<
                         "input_compressed_account_hashes {:?}",
                         input_compressed_account_hashes
                     );
-                    msg!("input roots {:?}", input_compressed_account_roots);
+                    // msg!("input roots {:?}", input_compressed_account_roots);
                     msg!(
                         "input_compressed_accounts_with_merkle_context: {:?}",
                         inputs.input_compressed_accounts_with_merkle_context