feat: assign addresses (#1695)

* feat: assign addresses to compressed accounts

* feat: make tx hash optional

* chore: bump photon commit hash

Author: ananas-block

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

@@ -8,7 +8,9 @@ use super::compressed_account::{
 use crate::{
     hash_to_bn254_field_size_be,
     instruction_data::data::{
-        NewAddressParams, NewAddressParamsPacked, PackedReadOnlyAddress, ReadOnlyAddress,
+        pack_pubkey_usize, NewAddressParams, NewAddressParamsAssigned,
+        NewAddressParamsAssignedPacked, NewAddressParamsPacked, PackedReadOnlyAddress,
+        ReadOnlyAddress,
     },
     CompressedAccountError, Pubkey,
 };
@@ -57,6 +59,31 @@ pub fn add_and_get_remaining_account_indices(
     vec
 }
 
+pub fn pack_new_address_params_assigned(
+    new_address_params: &[NewAddressParamsAssigned],
+    remaining_accounts: &mut HashMap<Pubkey, usize>,
+) -> Vec<NewAddressParamsAssignedPacked> {
+    let mut vec = Vec::new();
+    for new_address_param in new_address_params.iter() {
+        let address_merkle_tree_account_index = pack_pubkey_usize(
+            &new_address_param.address_merkle_tree_pubkey,
+            remaining_accounts,
+        );
+        let address_queue_account_index =
+            pack_pubkey_usize(&new_address_param.address_queue_pubkey, remaining_accounts);
+        vec.push(NewAddressParamsAssignedPacked {
+            seed: new_address_param.seed,
+            address_queue_account_index,
+            address_merkle_tree_root_index: new_address_param.address_merkle_tree_root_index,
+            address_merkle_tree_account_index,
+            assigned_to_account: new_address_param.assigned_account_index.is_some(),
+            assigned_account_index: new_address_param.assigned_account_index.unwrap_or_default(),
+        });
+    }
+
+    vec
+}
+
 // Helper function to pack new address params for instruction data in rust clients
 pub fn pack_new_address_params(
     new_address_params: &[NewAddressParams],

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

@@ -38,6 +38,37 @@ pub struct NewAddressParamsPacked {
     pub address_merkle_tree_root_index: u16,
 }
 
+#[derive(Debug, PartialEq, Default, Clone, Copy, AnchorDeserialize, AnchorSerialize)]
+pub struct NewAddressParamsAssignedPacked {
+    pub seed: [u8; 32],
+    pub address_queue_account_index: u8,
+    pub address_merkle_tree_account_index: u8,
+    pub address_merkle_tree_root_index: u16,
+    pub assigned_to_account: bool,
+    pub assigned_account_index: u8,
+}
+
+impl NewAddressParamsAssignedPacked {
+    pub fn new(address_params: NewAddressParamsPacked, index: Option<u8>) -> Self {
+        Self {
+            seed: address_params.seed,
+            address_queue_account_index: address_params.address_queue_account_index,
+            address_merkle_tree_account_index: address_params.address_merkle_tree_account_index,
+            address_merkle_tree_root_index: address_params.address_merkle_tree_root_index,
+            assigned_to_account: index.is_some(),
+            assigned_account_index: index.unwrap_or_default(),
+        }
+    }
+
+    pub fn assigned_account_index(&self) -> Option<u8> {
+        if self.assigned_to_account {
+            Some(self.assigned_account_index)
+        } else {
+            None
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Default, Clone, AnchorDeserialize, AnchorSerialize)]
 pub struct NewAddressParams {
     pub seed: [u8; 32],
@@ -46,6 +77,15 @@ pub struct NewAddressParams {
     pub address_merkle_tree_root_index: u16,
 }
 
+#[derive(Debug, PartialEq, Default, Clone, AnchorDeserialize, AnchorSerialize)]
+pub struct NewAddressParamsAssigned {
+    pub seed: [u8; 32],
+    pub address_queue_pubkey: Pubkey,
+    pub address_merkle_tree_pubkey: Pubkey,
+    pub address_merkle_tree_root_index: u16,
+    pub assigned_account_index: Option<u8>,
+}
+
 #[derive(Debug, PartialEq, Default, Clone, Copy, AnchorDeserialize, AnchorSerialize)]
 pub struct PackedReadOnlyAddress {
     pub address: [u8; 32],
@@ -59,7 +99,7 @@ pub struct ReadOnlyAddress {
     pub address_merkle_tree_pubkey: Pubkey,
     pub address_merkle_tree_root_index: u16,
 }
-
+// TODO: move
 pub fn pack_pubkey(pubkey: &Pubkey, hash_set: &mut HashMap<Pubkey, u8>) -> u8 {
     match hash_set.get(pubkey) {
         Some(index) => *index,
@@ -70,3 +110,14 @@ pub fn pack_pubkey(pubkey: &Pubkey, hash_set: &mut HashMap<Pubkey, u8>) -> u8 {
         }
     }
 }
+
+pub fn pack_pubkey_usize(pubkey: &Pubkey, hash_set: &mut HashMap<Pubkey, usize>) -> u8 {
+    match hash_set.get(pubkey) {
+        Some(index) => (*index) as u8,
+        None => {
+            let index = hash_set.len();
+            hash_set.insert(*pubkey, index);
+            index as u8
+        }
+    }
+}

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

@@ -5,16 +5,13 @@ use zerocopy::Ref;
 use super::{
     compressed_proof::CompressedProof,
     cpi_context::CompressedCpiContext,
-    zero_copy::{
-        ZNewAddressParamsPacked, ZPackedMerkleContext, ZPackedReadOnlyAddress,
-        ZPackedReadOnlyCompressedAccount,
-    },
+    zero_copy::{ZPackedMerkleContext, ZPackedReadOnlyAddress, ZPackedReadOnlyCompressedAccount},
 };
 use crate::{compressed_account::CompressedAccountData, pubkey::Pubkey, CompressedAccountError};
 
 pub trait InstructionDataTrait<'a> {
     fn owner(&self) -> Pubkey;
-    fn new_addresses(&self) -> &[ZNewAddressParamsPacked];
+    fn new_addresses(&self) -> &[impl NewAddressParamsTrait<'a>];
     fn input_accounts(&self) -> &[impl InputAccountTrait<'a>];
     fn output_accounts(&self) -> &[impl OutputAccountTrait<'a>];
     fn read_only_accounts(&self) -> Option<&[ZPackedReadOnlyCompressedAccount]>;
@@ -25,6 +22,18 @@ pub trait InstructionDataTrait<'a> {
     fn cpi_context(&self) -> Option<CompressedCpiContext>;
     fn bump(&self) -> Option<u8>;
     fn account_option_config(&self) -> AccountOptions;
+    fn with_transaction_hash(&self) -> bool;
+}
+
+pub trait NewAddressParamsTrait<'a>
+where
+    Self: Debug,
+{
+    fn seed(&self) -> [u8; 32];
+    fn address_queue_index(&self) -> u8;
+    fn address_merkle_tree_account_index(&self) -> u8;
+    fn address_merkle_tree_root_index(&self) -> u16;
+    fn assigned_compressed_account_index(&self) -> Option<usize>;
 }
 
 pub trait InputAccountTrait<'a>

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

@@ -9,11 +9,15 @@ use zerocopy::{
 use super::{
     compressed_proof::CompressedProof,
     cpi_context::CompressedCpiContext,
-    data::{NewAddressParamsPacked, PackedReadOnlyAddress},
-    traits::{AccountOptions, InputAccountTrait, InstructionDataTrait, OutputAccountTrait},
+    data::{NewAddressParamsAssignedPacked, PackedReadOnlyAddress},
+    traits::{
+        AccountOptions, InputAccountTrait, InstructionDataTrait, NewAddressParamsTrait,
+        OutputAccountTrait,
+    },
+    with_readonly::ZInstructionDataInvokeCpiWithReadOnlyMeta,
     zero_copy::{
-        ZCompressedCpiContext, ZNewAddressParamsPacked, ZPackedMerkleContext,
-        ZPackedReadOnlyAddress, ZPackedReadOnlyCompressedAccount,
+        ZNewAddressParamsAssignedPacked, ZPackedMerkleContext, ZPackedReadOnlyAddress,
+        ZPackedReadOnlyCompressedAccount,
     },
 };
 use crate::{
@@ -285,9 +289,10 @@ pub struct InstructionDataInvokeCpiWithAccountInfo {
     /// -> expect account decompression_recipient
     pub is_decompress: bool,
     pub with_cpi_context: bool,
+    pub with_transaction_hash: bool,
     pub cpi_context: CompressedCpiContext,
     pub proof: Option<CompressedProof>,
-    pub new_address_params: Vec<NewAddressParamsPacked>,
+    pub new_address_params: Vec<NewAddressParamsAssignedPacked>,
     pub account_infos: Vec<CompressedAccountInfo>,
     pub read_only_addresses: Vec<PackedReadOnlyAddress>,
     pub read_only_accounts: Vec<PackedReadOnlyCompressedAccount>,
@@ -307,6 +312,10 @@ impl<'a> InstructionDataTrait<'a> for ZInstructionDataInvokeCpiWithAccountInfo<'
         }
     }
 
+    fn with_transaction_hash(&self) -> bool {
+        self.meta.with_transaction_hash()
+    }
+
     fn read_only_accounts(&self) -> Option<&[ZPackedReadOnlyCompressedAccount]> {
         Some(self.read_only_accounts.as_slice())
     }
@@ -319,7 +328,7 @@ impl<'a> InstructionDataTrait<'a> for ZInstructionDataInvokeCpiWithAccountInfo<'
         self.meta.invoking_program_id
     }
 
-    fn new_addresses(&self) -> &[ZNewAddressParamsPacked] {
+    fn new_addresses(&self) -> &[impl NewAddressParamsTrait<'a>] {
         self.new_address_params.as_slice()
     }
 
@@ -360,44 +369,17 @@ impl<'a> InstructionDataTrait<'a> for ZInstructionDataInvokeCpiWithAccountInfo<'
     }
 }
 
-#[repr(C)]
-#[derive(
-    Debug, Default, PartialEq, Clone, Copy, FromBytes, IntoBytes, Unaligned, Immutable, KnownLayout,
-)]
-pub struct ZInstructionDataInvokeCpiWithAccountInfoMeta {
-    /// 0 With program ids
-    /// 1 without program ids
-    pub mode: u8,
-    pub bump: u8,
-    pub invoking_program_id: Pubkey,
-    /// If compress_or_decompress_lamports > 0 -> expect sol_pool_pda
-    pub compress_or_decompress_lamports: U64,
-    /// -> expect account decompression_recipient
-    is_decompress: u8,
-    with_cpi_context: u8,
-    pub cpi_context: ZCompressedCpiContext,
-}
-
-impl ZInstructionDataInvokeCpiWithAccountInfoMeta {
-    pub fn is_decompress(&self) -> bool {
-        self.is_decompress > 0
-    }
-    pub fn with_cpi_context(&self) -> bool {
-        self.with_cpi_context > 0
-    }
-}
-
 pub struct ZInstructionDataInvokeCpiWithAccountInfo<'a> {
-    meta: Ref<&'a [u8], ZInstructionDataInvokeCpiWithAccountInfoMeta>,
+    meta: Ref<&'a [u8], ZInstructionDataInvokeCpiWithReadOnlyMeta>,
     pub proof: Option<Ref<&'a [u8], CompressedProof>>,
-    pub new_address_params: ZeroCopySliceBorsh<'a, ZNewAddressParamsPacked>,
+    pub new_address_params: ZeroCopySliceBorsh<'a, ZNewAddressParamsAssignedPacked>,
     pub account_infos: Vec<ZCAccountInfo<'a>>,
     pub read_only_addresses: ZeroCopySliceBorsh<'a, ZPackedReadOnlyAddress>,
     pub read_only_accounts: ZeroCopySliceBorsh<'a, ZPackedReadOnlyCompressedAccount>,
 }
 
 impl<'a> Deref for ZInstructionDataInvokeCpiWithAccountInfo<'a> {
-    type Target = Ref<&'a [u8], ZInstructionDataInvokeCpiWithAccountInfoMeta>;
+    type Target = Ref<&'a [u8], ZInstructionDataInvokeCpiWithReadOnlyMeta>;
 
     fn deref(&self) -> &Self::Target {
         &self.meta
@@ -408,10 +390,10 @@ impl<'a> Deserialize<'a> for InstructionDataInvokeCpiWithAccountInfo {
     type Output = ZInstructionDataInvokeCpiWithAccountInfo<'a>;
     fn zero_copy_at(bytes: &'a [u8]) -> Result<(Self::Output, &'a [u8]), ZeroCopyError> {
         let (meta, bytes) =
-            Ref::<&[u8], ZInstructionDataInvokeCpiWithAccountInfoMeta>::from_prefix(bytes)?;
+            Ref::<&[u8], ZInstructionDataInvokeCpiWithReadOnlyMeta>::from_prefix(bytes)?;
         let (proof, bytes) = Option::<Ref<&[u8], CompressedProof>>::zero_copy_at(bytes)?;
         let (new_address_params, bytes) =
-            ZeroCopySliceBorsh::<'a, ZNewAddressParamsPacked>::from_bytes_at(bytes)?;
+            ZeroCopySliceBorsh::<'a, ZNewAddressParamsAssignedPacked>::from_bytes_at(bytes)?;
         let (account_infos, bytes) = {
             let (num_slices, mut bytes) = Ref::<&[u8], U32>::from_prefix(bytes)?;
             let num_slices = u32::from(*num_slices) as usize;

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

@@ -10,12 +10,14 @@ use super::{
     compressed_proof::CompressedProof,
     cpi_context::CompressedCpiContext,
     data::{
-        NewAddressParamsPacked, OutputCompressedAccountWithPackedContext, PackedReadOnlyAddress,
+        NewAddressParamsAssignedPacked, OutputCompressedAccountWithPackedContext,
+        PackedReadOnlyAddress,
     },
-    traits::{AccountOptions, InputAccountTrait, InstructionDataTrait},
+    traits::{AccountOptions, InputAccountTrait, InstructionDataTrait, NewAddressParamsTrait},
     zero_copy::{
-        ZCompressedCpiContext, ZNewAddressParamsPacked, ZOutputCompressedAccountWithPackedContext,
-        ZPackedMerkleContext, ZPackedReadOnlyAddress, ZPackedReadOnlyCompressedAccount,
+        ZCompressedCpiContext, ZNewAddressParamsAssignedPacked,
+        ZOutputCompressedAccountWithPackedContext, ZPackedMerkleContext, ZPackedReadOnlyAddress,
+        ZPackedReadOnlyCompressedAccount,
     },
 };
 use crate::{
@@ -199,9 +201,10 @@ pub struct InstructionDataInvokeCpiWithReadOnly {
     /// -> expect account decompression_recipient
     pub is_decompress: bool,
     pub with_cpi_context: bool,
+    pub with_transaction_hash: bool,
     pub cpi_context: CompressedCpiContext,
     pub proof: Option<CompressedProof>,
-    pub new_address_params: Vec<NewAddressParamsPacked>,
+    pub new_address_params: Vec<NewAddressParamsAssignedPacked>,
     pub input_compressed_accounts: Vec<InAccount>,
     pub output_compressed_accounts: Vec<OutputCompressedAccountWithPackedContext>,
     pub read_only_addresses: Vec<PackedReadOnlyAddress>,
@@ -223,6 +226,7 @@ pub struct ZInstructionDataInvokeCpiWithReadOnlyMeta {
     /// -> expect account decompression_recipient
     is_decompress: u8,
     with_cpi_context: u8,
+    with_transaction_hash: u8,
     pub cpi_context: ZCompressedCpiContext,
 }
 
@@ -233,13 +237,16 @@ impl ZInstructionDataInvokeCpiWithReadOnlyMeta {
     pub fn with_cpi_context(&self) -> bool {
         self.with_cpi_context > 0
     }
+    pub fn with_transaction_hash(&self) -> bool {
+        self.with_transaction_hash > 0
+    }
 }
 
 #[derive(Debug, PartialEq)]
 pub struct ZInstructionDataInvokeCpiWithReadOnly<'a> {
     meta: Ref<&'a [u8], ZInstructionDataInvokeCpiWithReadOnlyMeta>,
     pub proof: Option<Ref<&'a [u8], CompressedProof>>,
-    pub new_address_params: ZeroCopySliceBorsh<'a, ZNewAddressParamsPacked>,
+    pub new_address_params: ZeroCopySliceBorsh<'a, ZNewAddressParamsAssignedPacked>,
     pub input_compressed_accounts: Vec<ZInAccount<'a>>,
     pub output_compressed_accounts: Vec<ZOutputCompressedAccountWithPackedContext<'a>>,
     pub read_only_addresses: ZeroCopySliceBorsh<'a, ZPackedReadOnlyAddress>,
@@ -256,6 +263,10 @@ impl<'a> InstructionDataTrait<'a> for ZInstructionDataInvokeCpiWithReadOnly<'a>
         }
     }
 
+    fn with_transaction_hash(&self) -> bool {
+        self.meta.with_transaction_hash()
+    }
+
     fn bump(&self) -> Option<u8> {
         Some(self.bump)
     }
@@ -271,7 +282,7 @@ impl<'a> InstructionDataTrait<'a> for ZInstructionDataInvokeCpiWithReadOnly<'a>
         self.meta.invoking_program_id
     }
 
-    fn new_addresses(&self) -> &[ZNewAddressParamsPacked] {
+    fn new_addresses(&self) -> &[impl NewAddressParamsTrait<'a>] {
         self.new_address_params.as_slice()
     }
 
@@ -328,7 +339,7 @@ impl<'a> Deserialize<'a> for InstructionDataInvokeCpiWithReadOnly {
             Ref::<&[u8], ZInstructionDataInvokeCpiWithReadOnlyMeta>::from_prefix(bytes)?;
         let (proof, bytes) = Option::<Ref<&[u8], CompressedProof>>::zero_copy_at(bytes)?;
         let (new_address_params, bytes) =
-            ZeroCopySliceBorsh::<'a, ZNewAddressParamsPacked>::from_bytes_at(bytes)?;
+            ZeroCopySliceBorsh::<'a, ZNewAddressParamsAssignedPacked>::from_bytes_at(bytes)?;
         let (input_compressed_accounts, bytes) = {
             let (num_slices, mut bytes) = Ref::<&[u8], U32>::from_prefix(bytes)?;
             let num_slices = u32::from(*num_slices) as usize;
@@ -417,6 +428,7 @@ impl PartialEq<InstructionDataInvokeCpiWithReadOnly> for ZInstructionDataInvokeC
     }
 }
 
+// TODO: add randomized tests.
 #[test]
 fn test_read_only_zero_copy() {
     let borsh_struct = InstructionDataInvokeCpiWithReadOnly {
@@ -426,17 +438,20 @@ fn test_read_only_zero_copy() {
         compress_or_decompress_lamports: 0,
         is_decompress: false,
         with_cpi_context: false,
+        with_transaction_hash: true,
         cpi_context: CompressedCpiContext {
             set_context: false,
             first_set_context: false,
             cpi_context_account_index: 0,
         },
         proof: None,
-        new_address_params: vec![NewAddressParamsPacked {
+        new_address_params: vec![NewAddressParamsAssignedPacked {
             seed: [1; 32],
             address_merkle_tree_account_index: 1,
             address_queue_account_index: 2,
             address_merkle_tree_root_index: 3,
+            assigned_to_account: true,
+            assigned_account_index: 2,
         }],
         input_compressed_accounts: vec![InAccount {
             discriminator: [1, 2, 3, 4, 5, 6, 7, 8],