1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

use crate::{
    processor::{
        initialize_concurrent_merkle_tree::process_initialize_state_merkle_tree,
        initialize_nullifier_queue::process_initialize_nullifier_queue,
    },
    state::{
        queue::{queue_from_bytes_zero_copy_mut, QueueAccount},
        StateMerkleTreeAccount,
    },
    state_merkle_tree_from_bytes_zero_copy,
    utils::{
        check_account::check_account_balance_is_rent_exempt,
        check_signer_is_registered_or_authority::{
            check_signer_is_registered_or_authority, GroupAccounts,
        },
        transfer_lamports::transfer_lamports,
    },
    RegisteredProgram,
};
use anchor_lang::{prelude::*, solana_program::pubkey::Pubkey};

#[derive(Accounts)]
pub struct RolloverStateMerkleTreeAndNullifierQueue<'info> {
    #[account(mut)]
    /// Signer used to receive rollover accounts rentexemption reimbursement.
    pub fee_payer: Signer<'info>,
    pub authority: Signer<'info>,
    pub registered_program_pda: Option<Account<'info, RegisteredProgram>>,
    #[account(zero)]
    pub new_state_merkle_tree: AccountLoader<'info, StateMerkleTreeAccount>,
    #[account(zero)]
    pub new_nullifier_queue: AccountLoader<'info, QueueAccount>,
    #[account(mut)]
    pub old_state_merkle_tree: AccountLoader<'info, StateMerkleTreeAccount>,
    #[account(mut)]
    pub old_nullifier_queue: AccountLoader<'info, QueueAccount>,
}

impl<'info> GroupAccounts<'info> for RolloverStateMerkleTreeAndNullifierQueue<'info> {
    fn get_authority(&self) -> &Signer<'info> {
        &self.authority
    }
    fn get_registered_program_pda(&self) -> &Option<Account<'info, RegisteredProgram>> {
        &self.registered_program_pda
    }
}

/// Checks:
/// 1. Size of new accounts matches size old accounts
/// 2. Merkle tree is ready to be rolled over
/// 3. Merkle tree and nullifier queue are associated
/// 4. Merkle tree is not already rolled over Actions:
/// 1. mark Merkle tree as rolled over in this slot
/// 2. initialize new Merkle tree and nullifier queue with the same parameters
pub fn process_rollover_state_merkle_tree_nullifier_queue_pair<'a, 'b, 'c: 'info, 'info>(
    ctx: Context<'a, 'b, 'c, 'info, RolloverStateMerkleTreeAndNullifierQueue<'info>>,
) -> Result<()> {
    // TODO: rollover additional rent as well. (need to add a field to the metadata for this)
    let new_merkle_tree_account_info = ctx.accounts.new_state_merkle_tree.to_account_info();
    let merkle_tree_rent = check_account_balance_is_rent_exempt(
        &new_merkle_tree_account_info,
        ctx.accounts
            .old_state_merkle_tree
            .to_account_info()
            .data_len(),
    )?;
    let new_queue_account_info = ctx.accounts.new_nullifier_queue.to_account_info();
    let queue_rent = check_account_balance_is_rent_exempt(
        &new_queue_account_info,
        ctx.accounts
            .old_nullifier_queue
            .to_account_info()
            .data_len(),
    )?;

    let queue_metadata = {
        let (merkle_tree_metadata, queue_metadata) = {
            let mut merkle_tree_account_loaded = ctx.accounts.old_state_merkle_tree.load_mut()?;
            let mut queue_account_loaded = ctx.accounts.old_nullifier_queue.load_mut()?;
            check_signer_is_registered_or_authority::<
                RolloverStateMerkleTreeAndNullifierQueue,
                StateMerkleTreeAccount,
            >(&ctx, &merkle_tree_account_loaded)?;
            merkle_tree_account_loaded.metadata.rollover(
                ctx.accounts.old_nullifier_queue.key(),
                ctx.accounts.new_state_merkle_tree.key(),
            )?;
            queue_account_loaded.metadata.rollover(
                ctx.accounts.old_state_merkle_tree.key(),
                ctx.accounts.new_nullifier_queue.key(),
            )?;

            let merkle_tree_metadata = merkle_tree_account_loaded.metadata;
            let queue_metadata = queue_account_loaded.metadata;

            (merkle_tree_metadata, queue_metadata)
        };

        let merkle_tree = ctx.accounts.old_state_merkle_tree.to_account_info();
        let merkle_tree = merkle_tree.try_borrow_data()?;
        let merkle_tree = state_merkle_tree_from_bytes_zero_copy(&merkle_tree)?;

        let height = merkle_tree.height;

        if merkle_tree.next_index()
            < ((1 << height) * merkle_tree_metadata.rollover_metadata.rollover_threshold / 100)
                as usize
        {
            return err!(crate::errors::AccountCompressionErrorCode::NotReadyForRollover);
        }

        process_initialize_state_merkle_tree(
            &ctx.accounts.new_state_merkle_tree,
            merkle_tree_metadata.rollover_metadata.index,
            merkle_tree_metadata.access_metadata.owner,
            Some(merkle_tree_metadata.access_metadata.program_owner),
            Some(merkle_tree_metadata.access_metadata.forester),
            &(merkle_tree.height as u32),
            &(merkle_tree.changelog.capacity() as u64),
            &(merkle_tree.roots.capacity() as u64),
            &(merkle_tree.canopy_depth as u64),
            ctx.accounts.new_nullifier_queue.key(),
            merkle_tree_metadata.rollover_metadata.network_fee,
            Some(merkle_tree_metadata.rollover_metadata.rollover_threshold),
            Some(merkle_tree_metadata.rollover_metadata.close_threshold),
            merkle_tree_rent,
            queue_rent,
        )?;

        queue_metadata
    };
    {
        let nullifier_queue_account = ctx.accounts.old_nullifier_queue.to_account_info();
        let mut nullifier_queue = nullifier_queue_account.try_borrow_mut_data()?;
        let nullifier_queue = unsafe { queue_from_bytes_zero_copy_mut(&mut nullifier_queue)? };
        process_initialize_nullifier_queue(
            ctx.accounts.new_nullifier_queue.to_account_info(),
            &ctx.accounts.new_nullifier_queue,
            queue_metadata.rollover_metadata.index,
            queue_metadata.access_metadata.owner,
            Some(queue_metadata.access_metadata.program_owner),
            Some(queue_metadata.access_metadata.forester),
            ctx.accounts.new_state_merkle_tree.key(),
            nullifier_queue.hash_set.capacity as u16,
            nullifier_queue.hash_set.sequence_threshold as u64,
            Some(queue_metadata.rollover_metadata.rollover_threshold),
            Some(queue_metadata.rollover_metadata.close_threshold),
            queue_metadata.rollover_metadata.network_fee,
        )?;
    }
    let lamports = merkle_tree_rent + queue_rent;

    transfer_lamports(
        &ctx.accounts.old_state_merkle_tree.to_account_info(),
        &ctx.accounts.fee_payer.to_account_info(),
        lamports,
    )?;
    Ok(())
}