light_compressed_account/instruction_data/
zero_copy_set.rs

1use light_program_profiler::profile;
2use light_zero_copy::traits::ZeroCopyAt;
3use zerocopy::little_endian::U16;
4
5use crate::{
6    compressed_account::PackedMerkleContext,
7    instruction_data::{
8        compressed_proof::CompressedProof,
9        data::{ZNewAddressParamsAssignedPackedMut, ZOutputCompressedAccountWithPackedContextMut},
10        with_readonly::{ZInAccountMut, ZInstructionDataInvokeCpiWithReadOnlyMut},
11    },
12    CompressedAccountError, Pubkey,
13};
14
15impl ZOutputCompressedAccountWithPackedContextMut<'_> {
16    #[profile]
17    #[inline]
18    pub fn set(
19        &mut self,
20        owner: Pubkey,
21        lamports: u64,
22        address: Option<[u8; 32]>,
23        merkle_tree_index: u8,
24        discriminator: [u8; 8],
25        data_hash: [u8; 32],
26    ) -> Result<(), CompressedAccountError> {
27        self.compressed_account.owner = owner;
28        self.compressed_account.lamports = lamports.into();
29        if let Some(self_address) = self.compressed_account.address.as_deref_mut() {
30            let input_address =
31                address.ok_or(CompressedAccountError::InstructionDataExpectedAddress)?;
32            *self_address = input_address;
33        }
34        if self.compressed_account.address.is_none() && address.is_some() {
35            return Err(CompressedAccountError::ZeroCopyExpectedAddress);
36        }
37        *self.merkle_tree_index = merkle_tree_index;
38        let data = self
39            .compressed_account
40            .data
41            .as_mut()
42            .ok_or(CompressedAccountError::CompressedAccountDataNotInitialized)?;
43        data.discriminator = discriminator;
44        *data.data_hash = data_hash;
45
46        Ok(())
47    }
48}
49
50impl ZInAccountMut<'_> {
51    #[inline]
52    #[profile]
53    pub fn set_z(
54        &mut self,
55        discriminator: [u8; 8],
56        data_hash: [u8; 32],
57        merkle_context: &<PackedMerkleContext as ZeroCopyAt>::ZeroCopyAt,
58        root_index: U16,
59        lamports: u64,
60        address: Option<&[u8; 32]>,
61    ) -> Result<(), CompressedAccountError> {
62        self.discriminator = discriminator;
63        // Set merkle context fields manually due to mutability constraints
64        self.merkle_context.merkle_tree_pubkey_index = merkle_context.merkle_tree_pubkey_index;
65        self.merkle_context.queue_pubkey_index = merkle_context.queue_pubkey_index;
66        self.merkle_context
67            .leaf_index
68            .set(merkle_context.leaf_index.get());
69        self.merkle_context.prove_by_index = merkle_context.prove_by_index() as u8;
70        *self.root_index = root_index;
71        self.data_hash = data_hash;
72        *self.lamports = lamports.into();
73        if let Some(address) = address {
74            self.address
75                .as_mut()
76                .ok_or(CompressedAccountError::InstructionDataExpectedAddress)?
77                .copy_from_slice(address);
78        }
79        if self.address.is_some() && address.is_none() {
80            return Err(CompressedAccountError::ZeroCopyExpectedAddress);
81        }
82        Ok(())
83    }
84
85    #[inline]
86    #[profile]
87    pub fn set(
88        &mut self,
89        discriminator: [u8; 8],
90        data_hash: [u8; 32],
91        merkle_context: &PackedMerkleContext,
92        root_index: U16,
93        lamports: u64,
94        address: Option<&[u8]>,
95    ) -> Result<(), CompressedAccountError> {
96        self.discriminator = discriminator;
97        // Set merkle context fields manually due to mutability constraints
98        self.merkle_context.merkle_tree_pubkey_index = merkle_context.merkle_tree_pubkey_index;
99        self.merkle_context.queue_pubkey_index = merkle_context.queue_pubkey_index;
100        self.merkle_context
101            .leaf_index
102            .set(merkle_context.leaf_index);
103        self.merkle_context.prove_by_index = merkle_context.prove_by_index as u8;
104        *self.root_index = root_index;
105        self.data_hash = data_hash;
106        *self.lamports = lamports.into();
107        if let Some(address) = address {
108            self.address
109                .as_mut()
110                .ok_or(CompressedAccountError::InstructionDataExpectedAddress)?
111                .copy_from_slice(address);
112        }
113        if self.address.is_some() && address.is_none() {
114            return Err(CompressedAccountError::ZeroCopyExpectedAddress);
115        }
116        Ok(())
117    }
118}
119
120pub trait CompressedCpiContextTrait {
121    fn set_context(&self) -> u8;
122    fn first_set_context(&self) -> u8;
123}
124
125impl ZInstructionDataInvokeCpiWithReadOnlyMut<'_> {
126    #[inline]
127    pub fn initialize(
128        &mut self,
129        bump: u8,
130        invoking_program_id: &Pubkey,
131        input_proof: Option<<CompressedProof as ZeroCopyAt>::ZeroCopyAt>,
132        cpi_context: &Option<impl CompressedCpiContextTrait>,
133    ) -> Result<(), CompressedAccountError> {
134        self.mode = 1; // Small ix mode
135        self.bump = bump;
136        self.invoking_program_id = *invoking_program_id;
137        if let Some(proof) = self.proof.as_deref_mut() {
138            let input_proof =
139                input_proof.ok_or(CompressedAccountError::InstructionDataExpectedProof)?;
140            proof.a = input_proof.a;
141            proof.b = input_proof.b;
142            proof.c = input_proof.c;
143        }
144        if self.proof.is_none() && input_proof.is_some() {
145            return Err(CompressedAccountError::ZeroCopyExpectedProof);
146        }
147        // self.cpi_context is constant, always allocated
148        //      -> no reverse ok_or check necessary
149        if let Some(cpi_context) = cpi_context {
150            self.with_cpi_context = 1;
151            self.cpi_context.cpi_context_account_index = 0;
152            self.cpi_context.first_set_context = cpi_context.first_set_context();
153            self.cpi_context.set_context = cpi_context.set_context();
154        }
155
156        Ok(())
157    }
158}
159
160impl ZNewAddressParamsAssignedPackedMut<'_> {
161    #[inline]
162    #[profile]
163    pub fn set(
164        &mut self,
165        seed: [u8; 32],
166        address_merkle_tree_root_index: U16,
167        assigned_account_index: Option<u8>,
168        address_merkle_tree_account_index: u8,
169    ) {
170        self.seed = seed;
171        self.address_merkle_tree_root_index = address_merkle_tree_root_index;
172        self.address_queue_account_index = 0; // always 0 for v2 address trees.
173        if let Some(assigned_account_index) = assigned_account_index {
174            self.assigned_account_index = assigned_account_index;
175            self.assigned_to_account = 1; // set to true
176        } else {
177            self.assigned_account_index = 0;
178            self.assigned_to_account = 0; // set to false
179        }
180        // Note we can skip address derivation since we are assigning it to the account in index 0.
181        self.address_merkle_tree_account_index = address_merkle_tree_account_index;
182    }
183}
light_compressed_account/instruction_data/zero_copy_set.rs

light_compressed_account/instruction_data/
zero_copy_set.rs
