1use {
2 crate::nonce_info::NonceInfo,
3 solana_account::{AccountSharedData, ReadableAccount, WritableAccount},
4 solana_clock::Epoch,
5 solana_pubkey::Pubkey,
6 solana_transaction_context::TransactionAccount,
7};
8
9#[derive(PartialEq, Eq, Debug, Clone)]
12pub enum RollbackAccounts {
13 FeePayerOnly {
14 fee_payer: TransactionAccount,
15 },
16 SameNonceAndFeePayer {
17 nonce: TransactionAccount,
18 },
19 SeparateNonceAndFeePayer {
20 nonce: TransactionAccount,
21 fee_payer: TransactionAccount,
22 },
23}
24
25#[cfg(feature = "dev-context-only-utils")]
26impl Default for RollbackAccounts {
27 fn default() -> Self {
28 Self::FeePayerOnly {
29 fee_payer: TransactionAccount::default(),
30 }
31 }
32}
33
34pub struct RollbackAccountsIter<'a> {
37 fee_payer: Option<&'a TransactionAccount>,
38 nonce: Option<&'a TransactionAccount>,
39}
40
41impl<'a> Iterator for RollbackAccountsIter<'a> {
42 type Item = &'a TransactionAccount;
43
44 fn next(&mut self) -> Option<Self::Item> {
45 if let Some(fee_payer) = self.fee_payer.take() {
46 return Some(fee_payer);
47 }
48 if let Some(nonce) = self.nonce.take() {
49 return Some(nonce);
50 }
51 None
52 }
53}
54
55impl<'a> IntoIterator for &'a RollbackAccounts {
56 type Item = &'a TransactionAccount;
57 type IntoIter = RollbackAccountsIter<'a>;
58
59 fn into_iter(self) -> Self::IntoIter {
60 self.iter()
61 }
62}
63
64impl RollbackAccounts {
65 pub(crate) fn new(
66 nonce: Option<NonceInfo>,
67 fee_payer_address: Pubkey,
68 mut fee_payer_account: AccountSharedData,
69 fee_payer_loaded_rent_epoch: Epoch,
70 ) -> Self {
71 if let Some(nonce) = nonce {
72 if &fee_payer_address == nonce.address() {
73 fee_payer_account.set_data_from_slice(nonce.account().data());
78
79 RollbackAccounts::SameNonceAndFeePayer {
80 nonce: (fee_payer_address, fee_payer_account),
81 }
82 } else {
83 RollbackAccounts::SeparateNonceAndFeePayer {
84 nonce: (nonce.address, nonce.account),
85 fee_payer: (fee_payer_address, fee_payer_account),
86 }
87 }
88 } else {
89 fee_payer_account.set_rent_epoch(fee_payer_loaded_rent_epoch);
96 RollbackAccounts::FeePayerOnly {
97 fee_payer: (fee_payer_address, fee_payer_account),
98 }
99 }
100 }
101
102 pub fn count(&self) -> usize {
104 match self {
105 Self::FeePayerOnly { .. } | Self::SameNonceAndFeePayer { .. } => 1,
106 Self::SeparateNonceAndFeePayer { .. } => 2,
107 }
108 }
109
110 pub fn iter(&self) -> RollbackAccountsIter<'_> {
112 match self {
113 Self::FeePayerOnly { fee_payer } => RollbackAccountsIter {
114 fee_payer: Some(fee_payer),
115 nonce: None,
116 },
117 Self::SameNonceAndFeePayer { nonce } => RollbackAccountsIter {
118 fee_payer: None,
119 nonce: Some(nonce),
120 },
121 Self::SeparateNonceAndFeePayer { nonce, fee_payer } => RollbackAccountsIter {
122 fee_payer: Some(fee_payer),
123 nonce: Some(nonce),
124 },
125 }
126 }
127
128 pub fn data_size(&self) -> usize {
131 let mut total_size: usize = 0;
132 for (_, account) in self.iter() {
133 total_size = total_size.saturating_add(account.data().len());
134 }
135 total_size
136 }
137}
138
139#[cfg(test)]
140mod tests {
141 use {
142 super::*,
143 solana_account::{ReadableAccount, WritableAccount},
144 solana_hash::Hash,
145 solana_nonce::{
146 state::{Data as NonceData, DurableNonce, State as NonceState},
147 versions::Versions as NonceVersions,
148 },
149 solana_sdk_ids::system_program,
150 };
151
152 #[test]
153 fn test_new_fee_payer_only() {
154 let fee_payer_address = Pubkey::new_unique();
155 let fee_payer_account = AccountSharedData::new(100, 0, &Pubkey::default());
156 let fee_payer_rent_epoch = fee_payer_account.rent_epoch();
157
158 let rent_epoch_updated_fee_payer_account = {
159 let mut account = fee_payer_account.clone();
160 account.set_lamports(fee_payer_account.lamports());
161 account.set_rent_epoch(fee_payer_rent_epoch + 1);
162 account
163 };
164
165 let rollback_accounts = RollbackAccounts::new(
166 None,
167 fee_payer_address,
168 rent_epoch_updated_fee_payer_account,
169 fee_payer_rent_epoch,
170 );
171
172 let expected_fee_payer = (fee_payer_address, fee_payer_account);
173 match rollback_accounts {
174 RollbackAccounts::FeePayerOnly { fee_payer } => {
175 assert_eq!(expected_fee_payer, fee_payer);
176 }
177 _ => panic!("Expected FeePayerOnly variant"),
178 }
179 }
180
181 #[test]
182 fn test_new_same_nonce_and_fee_payer() {
183 let nonce_address = Pubkey::new_unique();
184 let durable_nonce = DurableNonce::from_blockhash(&Hash::new_unique());
185 let lamports_per_signature = 42;
186 let nonce_account = AccountSharedData::new_data(
187 43,
188 &NonceVersions::new(NonceState::Initialized(NonceData::new(
189 Pubkey::default(),
190 durable_nonce,
191 lamports_per_signature,
192 ))),
193 &system_program::id(),
194 )
195 .unwrap();
196
197 let rent_epoch_updated_fee_payer_account = {
198 let mut account = nonce_account.clone();
199 account.set_lamports(nonce_account.lamports());
200 account
201 };
202
203 let nonce = NonceInfo::new(nonce_address, rent_epoch_updated_fee_payer_account.clone());
204 let rollback_accounts = RollbackAccounts::new(
205 Some(nonce),
206 nonce_address,
207 rent_epoch_updated_fee_payer_account,
208 u64::MAX, );
210
211 let expected_rollback_accounts = RollbackAccounts::SameNonceAndFeePayer {
212 nonce: (nonce_address, nonce_account),
213 };
214
215 assert_eq!(expected_rollback_accounts, rollback_accounts);
216 }
217
218 #[test]
219 fn test_separate_nonce_and_fee_payer() {
220 let nonce_address = Pubkey::new_unique();
221 let durable_nonce = DurableNonce::from_blockhash(&Hash::new_unique());
222 let lamports_per_signature = 42;
223 let nonce_account = AccountSharedData::new_data(
224 43,
225 &NonceVersions::new(NonceState::Initialized(NonceData::new(
226 Pubkey::default(),
227 durable_nonce,
228 lamports_per_signature,
229 ))),
230 &system_program::id(),
231 )
232 .unwrap();
233
234 let fee_payer_address = Pubkey::new_unique();
235 let fee_payer_account = AccountSharedData::new(44, 0, &Pubkey::default());
236
237 let rent_epoch_updated_fee_payer_account = {
238 let mut account = fee_payer_account.clone();
239 account.set_lamports(fee_payer_account.lamports());
240 account
241 };
242
243 let nonce = NonceInfo::new(nonce_address, nonce_account.clone());
244 let rollback_accounts = RollbackAccounts::new(
245 Some(nonce),
246 fee_payer_address,
247 rent_epoch_updated_fee_payer_account.clone(),
248 u64::MAX, );
250
251 let expected_nonce = (nonce_address, nonce_account);
252 let expected_fee_payer = (fee_payer_address, fee_payer_account);
253 match rollback_accounts {
254 RollbackAccounts::SeparateNonceAndFeePayer { nonce, fee_payer } => {
255 assert_eq!(expected_nonce, nonce);
256 assert_eq!(expected_fee_payer, fee_payer);
257 }
258 _ => panic!("Expected SeparateNonceAndFeePayer variant"),
259 }
260 }
261}