1#![allow(clippy::await_holding_refcell_ref)]
2
3use crate::assert_rollover::{
4 assert_rolledover_merkle_trees, assert_rolledover_merkle_trees_metadata,
5 assert_rolledover_queues_metadata,
6};
7use account_compression::NullifierQueueConfig;
8use account_compression::{
9 self, initialize_address_merkle_tree::AccountLoader, state::QueueAccount,
10 StateMerkleTreeAccount, StateMerkleTreeConfig, ID,
11};
12use anchor_lang::{InstructionData, Lamports, ToAccountMetas};
13use forester_utils::{create_account_instruction, get_hash_set};
14use light_client::rpc::errors::RpcError;
15use light_client::rpc::RpcConnection;
16use light_concurrent_merkle_tree::{
17 copy::ConcurrentMerkleTreeCopy, zero_copy::ConcurrentMerkleTreeZeroCopyMut,
18};
19use light_hasher::Poseidon;
20use solana_sdk::clock::Slot;
21use solana_sdk::{
22 account::AccountSharedData,
23 account_info::AccountInfo,
24 instruction::{AccountMeta, Instruction},
25 signature::{Keypair, Signer},
26 transaction::Transaction,
27};
28use solana_sdk::{account::WritableAccount, pubkey::Pubkey};
29use std::mem;
30
31pub enum StateMerkleTreeRolloverMode {
32 QueueInvalidSize,
33 TreeInvalidSize,
34}
35
36#[allow(clippy::too_many_arguments)]
37pub async fn perform_state_merkle_tree_roll_over<R: RpcConnection>(
38 rpc: &mut R,
39 new_nullifier_queue_keypair: &Keypair,
40 new_state_merkle_tree_keypair: &Keypair,
41 merkle_tree_pubkey: &Pubkey,
42 nullifier_queue_pubkey: &Pubkey,
43 merkle_tree_config: &StateMerkleTreeConfig,
44 queue_config: &NullifierQueueConfig,
45 mode: Option<StateMerkleTreeRolloverMode>,
46) -> Result<(solana_sdk::signature::Signature, Slot), RpcError> {
47 let payer_pubkey = rpc.get_payer().pubkey();
48 let mut size = QueueAccount::size(queue_config.capacity as usize).unwrap();
49 if let Some(StateMerkleTreeRolloverMode::QueueInvalidSize) = mode {
50 size += 1;
51 }
52 let create_nullifier_queue_instruction = create_account_instruction(
53 &payer_pubkey,
54 size,
55 rpc.get_minimum_balance_for_rent_exemption(size)
56 .await
57 .unwrap(),
58 &ID,
59 Some(new_nullifier_queue_keypair),
60 );
61 let mut state_tree_size = account_compression::state::StateMerkleTreeAccount::size(
62 merkle_tree_config.height as usize,
63 merkle_tree_config.changelog_size as usize,
64 merkle_tree_config.roots_size as usize,
65 merkle_tree_config.canopy_depth as usize,
66 );
67 if let Some(StateMerkleTreeRolloverMode::TreeInvalidSize) = mode {
68 state_tree_size += 1;
69 }
70 let create_state_merkle_tree_instruction = create_account_instruction(
71 &payer_pubkey,
72 state_tree_size,
73 rpc.get_minimum_balance_for_rent_exemption(state_tree_size)
74 .await
75 .unwrap(),
76 &ID,
77 Some(new_state_merkle_tree_keypair),
78 );
79 let instruction_data =
80 account_compression::instruction::RolloverStateMerkleTreeAndNullifierQueue {};
81 let accounts = account_compression::accounts::RolloverStateMerkleTreeAndNullifierQueue {
82 fee_payer: rpc.get_payer().pubkey(),
83 authority: rpc.get_payer().pubkey(),
84 registered_program_pda: None,
85 new_state_merkle_tree: new_state_merkle_tree_keypair.pubkey(),
86 new_nullifier_queue: new_nullifier_queue_keypair.pubkey(),
87 old_state_merkle_tree: *merkle_tree_pubkey,
88 old_nullifier_queue: *nullifier_queue_pubkey,
89 };
90 let instruction = Instruction {
91 program_id: account_compression::ID,
92 accounts: [
93 accounts.to_account_metas(Some(true)),
94 vec![AccountMeta::new(*merkle_tree_pubkey, false)],
95 ]
96 .concat(),
97 data: instruction_data.data(),
98 };
99 let blockhash = rpc.get_latest_blockhash().await.unwrap();
100 let transaction = Transaction::new_signed_with_payer(
101 &[
102 create_nullifier_queue_instruction,
103 create_state_merkle_tree_instruction,
104 instruction,
105 ],
106 Some(&rpc.get_payer().pubkey()),
107 &vec![
108 &rpc.get_payer(),
109 &new_nullifier_queue_keypair,
110 &new_state_merkle_tree_keypair,
111 ],
112 blockhash,
113 );
114 rpc.process_transaction_with_context(transaction).await
115}
116
117pub async fn set_state_merkle_tree_next_index<R: RpcConnection>(
118 rpc: &mut R,
119 merkle_tree_pubkey: &Pubkey,
120 next_index: u64,
121 lamports: u64,
122) {
123 let mut merkle_tree = rpc.get_account(*merkle_tree_pubkey).await.unwrap().unwrap();
124 {
125 let merkle_tree_deserialized =
126 &mut ConcurrentMerkleTreeZeroCopyMut::<Poseidon, 26>::from_bytes_zero_copy_mut(
127 &mut merkle_tree.data[8 + std::mem::size_of::<StateMerkleTreeAccount>()..],
128 )
129 .unwrap();
130 unsafe {
131 *merkle_tree_deserialized.next_index = next_index as usize;
132 }
133 }
134 let mut account_share_data = AccountSharedData::from(merkle_tree);
135 account_share_data.set_lamports(lamports);
136 rpc.set_account(merkle_tree_pubkey, &account_share_data);
137 let mut merkle_tree = rpc.get_account(*merkle_tree_pubkey).await.unwrap().unwrap();
138 let merkle_tree_deserialized =
139 ConcurrentMerkleTreeZeroCopyMut::<Poseidon, 26>::from_bytes_zero_copy_mut(
140 &mut merkle_tree.data[8 + std::mem::size_of::<StateMerkleTreeAccount>()..],
141 )
142 .unwrap();
143 assert_eq!(merkle_tree_deserialized.next_index() as u64, next_index);
144}
145
146#[allow(clippy::too_many_arguments)]
147pub async fn assert_rolled_over_pair<R: RpcConnection>(
148 payer: &Pubkey,
149 rpc: &mut R,
150 fee_payer_prior_balance: &u64,
151 old_merkle_tree_pubkey: &Pubkey,
152 old_nullifier_queue_pubkey: &Pubkey,
153 new_merkle_tree_pubkey: &Pubkey,
154 new_nullifier_queue_pubkey: &Pubkey,
155 current_slot: u64,
156 additional_rent: u64,
157 num_signatures: u64,
158) {
159 let mut new_mt_account = rpc
160 .get_account(*new_merkle_tree_pubkey)
161 .await
162 .unwrap()
163 .unwrap();
164 let mut new_mt_lamports = 0u64;
165 let old_account_info = AccountInfo::new(
166 new_merkle_tree_pubkey,
167 false,
168 false,
169 &mut new_mt_lamports,
170 &mut new_mt_account.data,
171 &ID,
172 false,
173 0u64,
174 );
175 let new_mt_account =
176 AccountLoader::<StateMerkleTreeAccount>::try_from(&old_account_info).unwrap();
177 let new_loaded_mt_account = new_mt_account.load().unwrap();
178
179 let mut old_mt_account = rpc
180 .get_account(*old_merkle_tree_pubkey)
181 .await
182 .unwrap()
183 .unwrap();
184
185 let mut old_mt_lamports = 0u64;
186 let new_account_info = AccountInfo::new(
187 old_merkle_tree_pubkey,
188 false,
189 false,
190 &mut old_mt_lamports,
191 &mut old_mt_account.data,
192 &account_compression::ID,
193 false,
194 0u64,
195 );
196 let old_mt_account =
197 AccountLoader::<StateMerkleTreeAccount>::try_from(&new_account_info).unwrap();
198 let old_loaded_mt_account = old_mt_account.load().unwrap();
199
200 assert_rolledover_merkle_trees_metadata(
201 &old_loaded_mt_account.metadata,
202 &new_loaded_mt_account.metadata,
203 current_slot,
204 new_nullifier_queue_pubkey,
205 );
206
207 let old_mt_data = old_account_info.try_borrow_data().unwrap();
208 let old_mt = ConcurrentMerkleTreeCopy::<Poseidon, 26>::from_bytes_copy(
209 &old_mt_data[8 + mem::size_of::<StateMerkleTreeAccount>()..],
210 )
211 .unwrap();
212 let new_mt_data = new_account_info.try_borrow_data().unwrap();
213 let new_mt = ConcurrentMerkleTreeCopy::<Poseidon, 26>::from_bytes_copy(
214 &new_mt_data[8 + mem::size_of::<StateMerkleTreeAccount>()..],
215 )
216 .unwrap();
217 assert_rolledover_merkle_trees(&old_mt, &new_mt);
218
219 {
220 let mut new_queue_account = rpc
221 .get_account(*new_nullifier_queue_pubkey)
222 .await
223 .unwrap()
224 .unwrap();
225 let mut new_mt_lamports = 0u64;
226 let account_info = AccountInfo::new(
227 new_nullifier_queue_pubkey,
228 false,
229 false,
230 &mut new_mt_lamports,
231 &mut new_queue_account.data,
232 &account_compression::ID,
233 false,
234 0u64,
235 );
236 let new_queue_account = AccountLoader::<QueueAccount>::try_from(&account_info).unwrap();
237 let new_loaded_queue_account = new_queue_account.load().unwrap();
238 let mut old_queue_account = rpc
239 .get_account(*old_nullifier_queue_pubkey)
240 .await
241 .unwrap()
242 .unwrap();
243 let mut old_mt_lamports = 0u64;
244 let account_info = AccountInfo::new(
245 old_nullifier_queue_pubkey,
246 false,
247 false,
248 &mut old_mt_lamports,
249 &mut old_queue_account.data,
250 &account_compression::ID,
251 false,
252 0u64,
253 );
254 let old_queue_account = AccountLoader::<QueueAccount>::try_from(&account_info).unwrap();
255 let old_loaded_queue_account = old_queue_account.load().unwrap();
256
257 assert_rolledover_queues_metadata(
258 &old_loaded_queue_account.metadata,
259 &new_loaded_queue_account.metadata,
260 current_slot,
261 new_merkle_tree_pubkey,
262 new_nullifier_queue_pubkey,
263 old_mt_account.get_lamports(),
264 new_mt_account.get_lamports(),
265 new_queue_account.get_lamports(),
266 );
267 }
268 let fee_payer_post_balance = rpc.get_account(*payer).await.unwrap().unwrap().lamports;
269 assert_eq!(
271 *fee_payer_prior_balance,
272 fee_payer_post_balance + 5000 * num_signatures + additional_rent
273 );
274 let old_address_queue =
275 unsafe { get_hash_set::<QueueAccount, R>(rpc, *old_nullifier_queue_pubkey).await };
276 let new_address_queue =
277 unsafe { get_hash_set::<QueueAccount, R>(rpc, *new_nullifier_queue_pubkey).await };
278
279 assert_eq!(
280 old_address_queue.get_capacity(),
281 new_address_queue.get_capacity()
282 );
283
284 assert_eq!(
285 old_address_queue.sequence_threshold,
286 new_address_queue.sequence_threshold,
287 );
288}