tuktuk_program/
write_return_tasks.rs

1use anchor_lang::{
2    prelude::*,
3    solana_program::entrypoint::MAX_PERMITTED_DATA_INCREASE,
4    system_program::{self, transfer, Transfer},
5};
6
7use crate::TaskReturnV0;
8
9#[derive(AnchorSerialize, AnchorDeserialize, Clone, Default)]
10pub struct TasksAccountHeaderV0 {
11    pub num_tasks: u32,
12}
13
14pub struct WriteReturnTasksArgs<'info, I: Iterator<Item = TaskReturnV0>> {
15    pub program_id: Pubkey,
16    pub payer_info: PayerInfo<'info>,
17    pub accounts: Vec<AccountWithSeeds<'info>>,
18    pub tasks: I,
19    pub system_program: AccountInfo<'info>,
20}
21
22pub enum PayerInfo<'info> {
23    PdaPayer(AccountInfo<'info>),
24    SystemPayer {
25        account_info: AccountInfo<'info>,
26        seeds: Vec<Vec<u8>>,
27    },
28    Signer(AccountInfo<'info>),
29}
30
31#[derive(Clone)]
32pub struct AccountWithSeeds<'info> {
33    pub account: AccountInfo<'info>,
34    pub seeds: Vec<Vec<u8>>,
35}
36
37pub struct WriteReturnTasksReturn {
38    pub used_accounts: Vec<Pubkey>,
39    pub total_tasks: u32,
40}
41
42// Fills accounts with tasks up to the maximum length of 10kb, then moves on to the next account until it is out of tasks.
43// It should return a vector of the pubkeys of the accounts it used.
44// Note that tuktuk does not clean up these accounts, but you can reuse them with this method (it will overwrite)
45pub fn write_return_tasks<I>(args: WriteReturnTasksArgs<'_, I>) -> Result<WriteReturnTasksReturn>
46where
47    I: Iterator<Item = TaskReturnV0>,
48{
49    let WriteReturnTasksArgs {
50        program_id,
51        payer_info,
52        accounts,
53        mut tasks,
54        system_program,
55    } = args;
56    let mut used_accounts = Vec::with_capacity(accounts.len());
57    let mut original_sizes = Vec::with_capacity(accounts.len());
58
59    // Get the first task outside the loop to check if we have any tasks
60    let mut current_task = match tasks.next() {
61        Some(task) => task,
62        None => {
63            return Ok(WriteReturnTasksReturn {
64                used_accounts,
65                total_tasks: 0,
66            })
67        }
68    };
69
70    let mut total_tasks = 0;
71    for AccountWithSeeds { account, seeds } in accounts.iter() {
72        // Store original size before any reallocation
73        original_sizes.push(account.data_len());
74
75        let mut header = TasksAccountHeaderV0 { num_tasks: 0 };
76        let header_size = header.try_to_vec()?.len();
77        let mut total_size = header_size;
78
79        msg!("Assigning account {} and allocating space", account.key());
80        if account.owner == &system_program::ID {
81            // Assign account to our program
82            let seeds_refs: Vec<&[u8]> = seeds.iter().map(|s| s.as_slice()).collect();
83            let seeds_slice: &[&[u8]] = seeds_refs.as_slice();
84            system_program::assign(
85                CpiContext::new_with_signer(
86                    system_program.to_account_info(),
87                    system_program::Assign {
88                        account_to_assign: account.to_account_info(),
89                    },
90                    &[seeds_slice],
91                ),
92                &program_id,
93            )?;
94        }
95        account.realloc(MAX_PERMITTED_DATA_INCREASE, false)?;
96        let mut data = account.data.borrow_mut();
97
98        // Write tasks directly after header
99        let mut offset = header_size;
100        let mut num_tasks = 0;
101
102        loop {
103            let task_bytes = current_task.try_to_vec()?;
104            if offset + task_bytes.len() > MAX_PERMITTED_DATA_INCREASE {
105                break; // This task will be handled by the next account
106            }
107
108            data[offset..offset + task_bytes.len()].copy_from_slice(&task_bytes);
109            offset += task_bytes.len();
110            total_size += task_bytes.len();
111            num_tasks += 1;
112            total_tasks += 1;
113
114            // Get next task
115            match tasks.next() {
116                Some(task) => current_task = task,
117                None => {
118                    break;
119                }
120            }
121        }
122
123        if num_tasks > 0 {
124            header.num_tasks = num_tasks;
125
126            // Write header
127            let header_bytes = header.try_to_vec()?;
128            data[..header_size].copy_from_slice(&header_bytes);
129            drop(data);
130
131            // Resize account to actual size
132            account.realloc(total_size, false)?;
133            let rent = Rent::get()?.minimum_balance(total_size);
134            let current_lamports = account.lamports();
135            let rent_to_pay = rent.saturating_sub(current_lamports);
136            if rent_to_pay > 0 {
137                match &payer_info {
138                    PayerInfo::PdaPayer(account_info) => {
139                        if account_info.lamports()
140                            - Rent::get()?.minimum_balance(account_info.data_len())
141                            < rent_to_pay
142                        {
143                            // Reset all account sizes on error
144                            for (account, original_size) in
145                                accounts.iter().zip(original_sizes.iter())
146                            {
147                                account.account.realloc(*original_size, false)?;
148                            }
149                            return Err(error!(ErrorCode::ConstraintRentExempt));
150                        }
151                        account_info.sub_lamports(rent_to_pay)?;
152                        account.add_lamports(rent_to_pay)?;
153                    }
154                    PayerInfo::SystemPayer {
155                        account_info,
156                        seeds,
157                    } => {
158                        let payer_seeds_refs: Vec<&[u8]> =
159                            seeds.iter().map(|s| s.as_slice()).collect();
160
161                        if account_info.lamports()
162                            - Rent::get()?.minimum_balance(account_info.data_len())
163                            < rent_to_pay
164                        {
165                            // Reset all account sizes on error
166                            for (account, original_size) in
167                                accounts.iter().zip(original_sizes.iter())
168                            {
169                                account.account.realloc(*original_size, false)?;
170                            }
171                            return Err(error!(ErrorCode::ConstraintRentExempt));
172                        }
173                        transfer(
174                            CpiContext::new_with_signer(
175                                system_program.clone(),
176                                Transfer {
177                                    from: account_info.clone(),
178                                    to: account.clone(),
179                                },
180                                &[payer_seeds_refs.as_slice()],
181                            ),
182                            rent_to_pay,
183                        )?;
184                    }
185                    PayerInfo::Signer(account_info) => {
186                        if account_info.lamports()
187                            - Rent::get()?.minimum_balance(account_info.data_len())
188                            < rent_to_pay
189                        {
190                            // Reset all account sizes on error
191                            for (account, original_size) in
192                                accounts.iter().zip(original_sizes.iter())
193                            {
194                                account.account.realloc(*original_size, false)?;
195                            }
196                            return Err(error!(ErrorCode::ConstraintRentExempt));
197                        }
198                        transfer(
199                            CpiContext::new(
200                                system_program.clone(),
201                                Transfer {
202                                    from: account_info.clone(),
203                                    to: account.clone(),
204                                },
205                            ),
206                            rent_to_pay,
207                        )?;
208                    }
209                }
210            }
211
212            used_accounts.push(*account.key);
213        } else {
214            drop(data);
215            account.realloc(0, false)?;
216        }
217
218        // If we've processed all tasks, we can exit
219        if num_tasks == 0 || tasks.next().is_none() {
220            break;
221        }
222    }
223
224    // Check if we still have unprocessed tasks
225    if tasks.next().is_some() {
226        return Err(error!(ErrorCode::ConstraintRaw));
227    }
228
229    Ok(WriteReturnTasksReturn {
230        used_accounts,
231        total_tasks,
232    })
233}
tuktuk_program/write_return_tasks.rs

tuktuk_program/
write_return_tasks.rs
