chore: refactored insert_into_queues -> instruction_data, z_compressed_account hash -> compressed_account.rs (#1589)

Author: ananas-block

program-libs/compressed-account/src/compressed_account.rs

@@ -8,8 +8,12 @@ use light_hasher::{Hasher, Poseidon};
 use solana_program::pubkey::Pubkey;
 
 use crate::{
-    address::pack_account, hash_to_bn254_field_size_be,
-    instruction_data::data::OutputCompressedAccountWithPackedContext, CompressedAccountError,
+    address::pack_account,
+    hash_to_bn254_field_size_be,
+    instruction_data::{
+        data::OutputCompressedAccountWithPackedContext, zero_copy::ZCompressedAccount,
+    },
+    CompressedAccountError,
 };
 
 #[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
@@ -261,6 +265,68 @@ impl CompressedAccount {
     }
 }
 
+/// Hashing scheme:
+/// H(owner || leaf_index || merkle_tree_pubkey || lamports || address || data.discriminator || data.data_hash)
+impl ZCompressedAccount<'_> {
+    pub fn hash_with_hashed_values<H: Hasher>(
+        &self,
+        &owner_hashed: &[u8; 32],
+        &merkle_tree_hashed: &[u8; 32],
+        leaf_index: &u32,
+    ) -> Result<[u8; 32], CompressedAccountError> {
+        let capacity = 3
+            + std::cmp::min(u64::from(self.lamports), 1) as usize
+            + self.address.is_some() as usize
+            + self.data.is_some() as usize * 2;
+        let mut vec: Vec<&[u8]> = Vec::with_capacity(capacity);
+        vec.push(owner_hashed.as_slice());
+
+        // leaf index and merkle tree pubkey are used to make every compressed account hash unique
+        let leaf_index = leaf_index.to_le_bytes();
+        vec.push(leaf_index.as_slice());
+
+        vec.push(merkle_tree_hashed.as_slice());
+
+        // Lamports are only hashed if non-zero to safe CU
+        // For safety we prefix the lamports with 1 in 1 byte.
+        // Thus even if the discriminator has the same value as the lamports, the hash will be different.
+        let mut lamports_bytes = [1, 0, 0, 0, 0, 0, 0, 0, 0];
+        if self.lamports != 0 {
+            lamports_bytes[1..].copy_from_slice(&(u64::from(self.lamports)).to_le_bytes());
+            vec.push(lamports_bytes.as_slice());
+        }
+
+        if self.address.is_some() {
+            vec.push(self.address.as_ref().unwrap().as_slice());
+        }
+
+        let mut discriminator_bytes = [2, 0, 0, 0, 0, 0, 0, 0, 0];
+        if let Some(data) = &self.data {
+            discriminator_bytes[1..].copy_from_slice(data.discriminator.as_slice());
+            vec.push(&discriminator_bytes);
+            vec.push(data.data_hash.as_slice());
+        }
+        let hash = H::hashv(&vec)?;
+        Ok(hash)
+    }
+
+    pub fn hash<H: Hasher>(
+        &self,
+        &merkle_tree_pubkey: &Pubkey,
+        leaf_index: &u32,
+    ) -> Result<[u8; 32], CompressedAccountError> {
+        self.hash_with_hashed_values::<H>(
+            &hash_to_bn254_field_size_be(&self.owner.to_bytes())
+                .unwrap()
+                .0,
+            &hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
+                .unwrap()
+                .0,
+            leaf_index,
+        )
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use light_hasher::Poseidon;
@@ -468,4 +534,207 @@ mod tests {
             no_address_no_data_no_lamports_hash
         );
     }
+
+    /// Tests:
+    /// 1. functional with all inputs set
+    /// 2. no data
+    /// 3. no address
+    /// 4. no address and no lamports
+    /// 5. no address and no data
+    /// 6. no address, no data, no lamports
+    #[test]
+    fn test_zcompressed_account_hash() {
+        let owner = Pubkey::new_unique();
+        let address = [1u8; 32];
+        let data = CompressedAccountData {
+            discriminator: [1u8; 8],
+            data: vec![2u8; 32],
+            data_hash: [3u8; 32],
+        };
+        let lamports = 100;
+        let compressed_account = CompressedAccount {
+            owner,
+            lamports,
+            address: Some(address),
+            data: Some(data.clone()),
+        };
+        let merkle_tree_pubkey = Pubkey::new_unique();
+        let leaf_index = 1;
+        let hash = compressed_account
+            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
+            .unwrap();
+        let hash_manual = Poseidon::hashv(&[
+            hash_to_bn254_field_size_be(&owner.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            leaf_index.to_le_bytes().as_slice(),
+            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            [&[1u8], lamports.to_le_bytes().as_slice()]
+                .concat()
+                .as_slice(),
+            address.as_slice(),
+            [&[2u8], data.discriminator.as_slice()].concat().as_slice(),
+            &data.data_hash,
+        ])
+        .unwrap();
+        assert_eq!(hash, hash_manual);
+        assert_eq!(hash.len(), 32);
+
+        // no data
+        let compressed_account = CompressedAccount {
+            owner,
+            lamports,
+            address: Some(address),
+            data: None,
+        };
+        let no_data_hash = compressed_account
+            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
+            .unwrap();
+
+        let hash_manual = Poseidon::hashv(&[
+            hash_to_bn254_field_size_be(&owner.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            leaf_index.to_le_bytes().as_slice(),
+            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            [&[1u8], lamports.to_le_bytes().as_slice()]
+                .concat()
+                .as_slice(),
+            address.as_slice(),
+        ])
+        .unwrap();
+        assert_eq!(no_data_hash, hash_manual);
+        assert_ne!(hash, no_data_hash);
+
+        // no address
+        let compressed_account = CompressedAccount {
+            owner,
+            lamports,
+            address: None,
+            data: Some(data.clone()),
+        };
+        let no_address_hash = compressed_account
+            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
+            .unwrap();
+        let hash_manual = Poseidon::hashv(&[
+            hash_to_bn254_field_size_be(&owner.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            leaf_index.to_le_bytes().as_slice(),
+            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            [&[1u8], lamports.to_le_bytes().as_slice()]
+                .concat()
+                .as_slice(),
+            [&[2u8], data.discriminator.as_slice()].concat().as_slice(),
+            &data.data_hash,
+        ])
+        .unwrap();
+        assert_eq!(no_address_hash, hash_manual);
+        assert_ne!(hash, no_address_hash);
+        assert_ne!(no_data_hash, no_address_hash);
+
+        // no address no lamports
+        let compressed_account = CompressedAccount {
+            owner,
+            lamports: 0,
+            address: None,
+            data: Some(data.clone()),
+        };
+        let no_address_no_lamports_hash = compressed_account
+            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
+            .unwrap();
+        let hash_manual = Poseidon::hashv(&[
+            hash_to_bn254_field_size_be(&owner.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            leaf_index.to_le_bytes().as_slice(),
+            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            [&[2u8], data.discriminator.as_slice()].concat().as_slice(),
+            &data.data_hash,
+        ])
+        .unwrap();
+        assert_eq!(no_address_no_lamports_hash, hash_manual);
+        assert_ne!(hash, no_address_no_lamports_hash);
+        assert_ne!(no_data_hash, no_address_no_lamports_hash);
+        assert_ne!(no_address_hash, no_address_no_lamports_hash);
+
+        // no address and no data
+        let compressed_account = CompressedAccount {
+            owner,
+            lamports,
+            address: None,
+            data: None,
+        };
+        let no_address_no_data_hash = compressed_account
+            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
+            .unwrap();
+        let hash_manual = Poseidon::hashv(&[
+            hash_to_bn254_field_size_be(&owner.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            leaf_index.to_le_bytes().as_slice(),
+            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            [&[1u8], lamports.to_le_bytes().as_slice()]
+                .concat()
+                .as_slice(),
+        ])
+        .unwrap();
+        assert_eq!(no_address_no_data_hash, hash_manual);
+        assert_ne!(hash, no_address_no_data_hash);
+        assert_ne!(no_data_hash, no_address_no_data_hash);
+        assert_ne!(no_address_hash, no_address_no_data_hash);
+        assert_ne!(no_address_no_lamports_hash, no_address_no_data_hash);
+
+        // no address, no data, no lamports
+        let compressed_account = CompressedAccount {
+            owner,
+            lamports: 0,
+            address: None,
+            data: None,
+        };
+        let no_address_no_data_no_lamports_hash = compressed_account
+            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
+            .unwrap();
+        let hash_manual = Poseidon::hashv(&[
+            hash_to_bn254_field_size_be(&owner.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+            leaf_index.to_le_bytes().as_slice(),
+            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
+                .unwrap()
+                .0
+                .as_slice(),
+        ])
+        .unwrap();
+        assert_eq!(no_address_no_data_no_lamports_hash, hash_manual);
+        assert_ne!(no_address_no_data_hash, no_address_no_data_no_lamports_hash);
+        assert_ne!(hash, no_address_no_data_no_lamports_hash);
+        assert_ne!(no_data_hash, no_address_no_data_no_lamports_hash);
+        assert_ne!(no_address_hash, no_address_no_data_no_lamports_hash);
+        assert_ne!(
+            no_address_no_lamports_hash,
+            no_address_no_data_no_lamports_hash
+        );
+    }
 }

program-libs/compressed-account/src/event.rs

@@ -2,9 +2,10 @@ use borsh::{BorshDeserialize, BorshSerialize};
 use light_zero_copy::{borsh::Deserialize, errors::ZeroCopyError};
 use solana_program::pubkey::Pubkey;
 
-use super::{discriminators::*, insert_into_queues::InsertIntoQueuesInstructionData};
+use super::discriminators::*;
 use crate::instruction_data::{
     data::OutputCompressedAccountWithPackedContext,
+    insert_into_queues::InsertIntoQueuesInstructionData,
     zero_copy::{
         ZInstructionDataInvoke, ZInstructionDataInvokeCpi, ZInstructionDataInvokeCpiWithReadOnly,
     },

program-libs/compressed-account/src/insert_into_queues.rs renamed to program-libs/compressed-account/src/instruction_data/insert_into_queues.rs

@@ -1,5 +1,6 @@
 pub mod compressed_proof;
 pub mod cpi_context;
 pub mod data;
+pub mod insert_into_queues;
 pub mod invoke_cpi;
 pub mod zero_copy;

program-libs/compressed-account/src/instruction_data/mod.rs

@@ -12,10 +12,8 @@ pub mod compressed_account;
 pub mod discriminators;
 pub mod event;
 pub mod hash_chain;
-pub mod insert_into_queues;
 pub mod instruction_data;
 pub mod pubkey;
-pub mod z_compressed_account;
 
 #[derive(Debug, Error, PartialEq)]
 pub enum CompressedAccountError {