1use {
2 crate::nonce_info::NonceInfo,
3 solana_account::{AccountSharedData, ReadableAccount, WritableAccount},
4 solana_clock::Epoch,
5 solana_pubkey::Pubkey,
6 solana_transaction_context::transaction_accounts::KeyedAccountSharedData,
7};
8
9#[derive(PartialEq, Eq, Debug, Clone)]
12pub enum RollbackAccounts {
13 FeePayerOnly {
14 fee_payer: KeyedAccountSharedData,
15 },
16 SameNonceAndFeePayer {
17 nonce: KeyedAccountSharedData,
18 },
19 SeparateNonceAndFeePayer {
20 nonce: KeyedAccountSharedData,
21 fee_payer: KeyedAccountSharedData,
22 },
23}
24
25#[cfg(feature = "dev-context-only-utils")]
26impl Default for RollbackAccounts {
27 fn default() -> Self {
28 Self::FeePayerOnly {
29 fee_payer: KeyedAccountSharedData::default(),
30 }
31 }
32}
33
34pub struct RollbackAccountsIter<'a> {
37 fee_payer: Option<&'a KeyedAccountSharedData>,
38 nonce: Option<&'a KeyedAccountSharedData>,
39}
40
41impl<'a> Iterator for RollbackAccountsIter<'a> {
42 type Item = &'a KeyedAccountSharedData;
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 KeyedAccountSharedData;
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 fee_payer(&self) -> &KeyedAccountSharedData {
104 match self {
105 Self::FeePayerOnly { fee_payer } => fee_payer,
106 Self::SameNonceAndFeePayer { nonce } => nonce,
107 Self::SeparateNonceAndFeePayer { fee_payer, .. } => fee_payer,
108 }
109 }
110
111 pub fn count(&self) -> usize {
113 match self {
114 Self::FeePayerOnly { .. } | Self::SameNonceAndFeePayer { .. } => 1,
115 Self::SeparateNonceAndFeePayer { .. } => 2,
116 }
117 }
118
119 pub fn iter(&self) -> RollbackAccountsIter<'_> {
121 match self {
122 Self::FeePayerOnly { fee_payer } => RollbackAccountsIter {
123 fee_payer: Some(fee_payer),
124 nonce: None,
125 },
126 Self::SameNonceAndFeePayer { nonce } => RollbackAccountsIter {
127 fee_payer: None,
128 nonce: Some(nonce),
129 },
130 Self::SeparateNonceAndFeePayer { nonce, fee_payer } => RollbackAccountsIter {
131 fee_payer: Some(fee_payer),
132 nonce: Some(nonce),
133 },
134 }
135 }
136
137 pub fn data_size(&self) -> usize {
140 let mut total_size: usize = 0;
141 for (_, account) in self.iter() {
142 total_size = total_size.saturating_add(account.data().len());
143 }
144 total_size
145 }
146}
147
148#[cfg(test)]
149mod tests {
150 use {
151 super::*,
152 solana_account::{ReadableAccount, WritableAccount},
153 solana_hash::Hash,
154 solana_nonce::{
155 state::{Data as NonceData, DurableNonce, State as NonceState},
156 versions::Versions as NonceVersions,
157 },
158 solana_sdk_ids::system_program,
159 };
160
161 #[test]
162 fn test_new_fee_payer_only() {
163 let fee_payer_address = Pubkey::new_unique();
164 let fee_payer_account = AccountSharedData::new(100, 0, &Pubkey::default());
165 let fee_payer_rent_epoch = fee_payer_account.rent_epoch();
166
167 let rent_epoch_updated_fee_payer_account = {
168 let mut account = fee_payer_account.clone();
169 account.set_lamports(fee_payer_account.lamports());
170 account.set_rent_epoch(fee_payer_rent_epoch + 1);
171 account
172 };
173
174 let rollback_accounts = RollbackAccounts::new(
175 None,
176 fee_payer_address,
177 rent_epoch_updated_fee_payer_account,
178 fee_payer_rent_epoch,
179 );
180
181 let expected_fee_payer = (fee_payer_address, fee_payer_account);
182 match rollback_accounts {
183 RollbackAccounts::FeePayerOnly { fee_payer } => {
184 assert_eq!(expected_fee_payer, fee_payer);
185 }
186 _ => panic!("Expected FeePayerOnly variant"),
187 }
188 }
189
190 #[test]
191 fn test_new_same_nonce_and_fee_payer() {
192 let nonce_address = Pubkey::new_unique();
193 let durable_nonce = DurableNonce::from_blockhash(&Hash::new_unique());
194 let lamports_per_signature = 42;
195 let nonce_account = AccountSharedData::new_data(
196 43,
197 &NonceVersions::new(NonceState::Initialized(NonceData::new(
198 Pubkey::default(),
199 durable_nonce,
200 lamports_per_signature,
201 ))),
202 &system_program::id(),
203 )
204 .unwrap();
205
206 let rent_epoch_updated_fee_payer_account = {
207 let mut account = nonce_account.clone();
208 account.set_lamports(nonce_account.lamports());
209 account
210 };
211
212 let nonce = NonceInfo::new(nonce_address, rent_epoch_updated_fee_payer_account.clone());
213 let rollback_accounts = RollbackAccounts::new(
214 Some(nonce),
215 nonce_address,
216 rent_epoch_updated_fee_payer_account,
217 u64::MAX, );
219
220 let expected_rollback_accounts = RollbackAccounts::SameNonceAndFeePayer {
221 nonce: (nonce_address, nonce_account),
222 };
223
224 assert_eq!(expected_rollback_accounts, rollback_accounts);
225 }
226
227 #[test]
228 fn test_separate_nonce_and_fee_payer() {
229 let nonce_address = Pubkey::new_unique();
230 let durable_nonce = DurableNonce::from_blockhash(&Hash::new_unique());
231 let lamports_per_signature = 42;
232 let nonce_account = AccountSharedData::new_data(
233 43,
234 &NonceVersions::new(NonceState::Initialized(NonceData::new(
235 Pubkey::default(),
236 durable_nonce,
237 lamports_per_signature,
238 ))),
239 &system_program::id(),
240 )
241 .unwrap();
242
243 let fee_payer_address = Pubkey::new_unique();
244 let fee_payer_account = AccountSharedData::new(44, 0, &Pubkey::default());
245
246 let rent_epoch_updated_fee_payer_account = {
247 let mut account = fee_payer_account.clone();
248 account.set_lamports(fee_payer_account.lamports());
249 account
250 };
251
252 let nonce = NonceInfo::new(nonce_address, nonce_account.clone());
253 let rollback_accounts = RollbackAccounts::new(
254 Some(nonce),
255 fee_payer_address,
256 rent_epoch_updated_fee_payer_account.clone(),
257 u64::MAX, );
259
260 let expected_nonce = (nonce_address, nonce_account);
261 let expected_fee_payer = (fee_payer_address, fee_payer_account);
262 match rollback_accounts {
263 RollbackAccounts::SeparateNonceAndFeePayer { nonce, fee_payer } => {
264 assert_eq!(expected_nonce, nonce);
265 assert_eq!(expected_fee_payer, fee_payer);
266 }
267 _ => panic!("Expected SeparateNonceAndFeePayer variant"),
268 }
269 }
270}