1use crate::{
2 AccountId, AccountIdRef, DefuseError, Lock, Nonce, NoncePrefix, Nonces, Result, Salt,
3 amounts::Amounts,
4 fees::Pips,
5 intents::{
6 auth::AuthCall,
7 tokens::{
8 FtWithdraw, MtWithdraw, NativeWithdraw, NftWithdraw, NotifyOnTransfer, StorageDeposit,
9 },
10 },
11 public_key::PublicKey,
12 token_id::{TokenId, nep141::Nep141TokenId, nep171::Nep171TokenId, nep245::Nep245TokenId},
13 tokens::BATCH_MAX_TOKENS,
14};
15use defuse_bitmap::{U248, U256};
16use std::{
17 borrow::Cow,
18 collections::{HashMap, HashSet},
19};
20
21use super::{State, StateView};
22
23#[derive(Debug)]
24pub struct CachedState<W: StateView> {
25 view: W,
26 accounts: CachedAccounts,
27}
28
29impl<W> CachedState<W>
30where
31 W: StateView,
32{
33 #[inline]
34 pub fn new(view: W) -> Self {
35 Self {
36 view,
37 accounts: CachedAccounts::new(),
38 }
39 }
40}
41
42impl<W> StateView for CachedState<W>
43where
44 W: StateView,
45{
46 #[inline]
47 fn verifying_contract(&self) -> Cow<'_, AccountIdRef> {
48 self.view.verifying_contract()
49 }
50
51 #[inline]
52 fn wnear_id(&self) -> Cow<'_, AccountIdRef> {
53 self.view.wnear_id()
54 }
55
56 #[inline]
57 fn fee(&self) -> Pips {
58 self.view.fee()
59 }
60
61 #[inline]
62 fn fee_collector(&self) -> Cow<'_, AccountIdRef> {
63 self.view.fee_collector()
64 }
65
66 fn has_public_key(&self, account_id: &AccountIdRef, public_key: &PublicKey) -> bool {
67 if let Some(account) = self.accounts.get(account_id).map(Lock::as_inner_unchecked) {
68 if account.public_keys_added.contains(public_key) {
69 return true;
70 }
71 if account.public_keys_removed.contains(public_key) {
72 return false;
73 }
74 }
75 self.view.has_public_key(account_id, public_key)
76 }
77
78 fn iter_public_keys(&self, account_id: &AccountIdRef) -> impl Iterator<Item = PublicKey> + '_ {
79 let account = self.accounts.get(account_id).map(Lock::as_inner_unchecked);
80 self.view
81 .iter_public_keys(account_id)
82 .filter(move |pk| account.is_none_or(|a| !a.public_keys_removed.contains(pk)))
83 .chain(
84 account
85 .map(|a| &a.public_keys_added)
86 .into_iter()
87 .flatten()
88 .copied(),
89 )
90 }
91
92 fn is_nonce_used(&self, account_id: &AccountIdRef, nonce: Nonce) -> bool {
93 self.accounts
94 .get(account_id)
95 .map(Lock::as_inner_unchecked)
96 .is_some_and(|account| account.is_nonce_used(nonce))
97 || self.view.is_nonce_used(account_id, nonce)
98 }
99
100 fn balance_of(&self, account_id: &AccountIdRef, token_id: &TokenId) -> u128 {
101 self.accounts
102 .get(account_id)
103 .map(Lock::as_inner_unchecked)
104 .and_then(|account| account.token_amounts.get(token_id).copied())
105 .unwrap_or_else(|| self.view.balance_of(account_id, token_id))
106 }
107
108 fn is_account_locked(&self, account_id: &AccountIdRef) -> bool {
109 self.accounts
110 .get(account_id)
111 .map_or_else(|| self.view.is_account_locked(account_id), Lock::is_locked)
112 }
113
114 fn is_auth_by_predecessor_id_enabled(&self, account_id: &AccountIdRef) -> bool {
115 let was_enabled = self.view.is_auth_by_predecessor_id_enabled(account_id);
116 let toggled = self
117 .accounts
118 .get(account_id)
119 .map(Lock::as_inner_unchecked)
120 .is_some_and(|a| a.auth_by_predecessor_id_toggled);
121 was_enabled ^ toggled
122 }
123
124 fn is_valid_salt(&self, salt: Salt) -> bool {
125 self.view.is_valid_salt(salt)
126 }
127}
128
129impl<W> State for CachedState<W>
130where
131 W: StateView,
132{
133 fn add_public_key(&mut self, account_id: AccountId, public_key: PublicKey) -> Result<()> {
134 let had = self.view.has_public_key(&account_id, &public_key);
135 let account = self
136 .accounts
137 .get_or_create(account_id.clone(), |account_id| {
138 self.view.is_account_locked(account_id)
139 })
140 .get_mut()
141 .ok_or_else(|| DefuseError::AccountLocked(account_id.clone()))?;
142 let added = if had {
143 account.public_keys_removed.remove(&public_key)
144 } else {
145 account.public_keys_added.insert(public_key)
146 };
147 if !added {
148 return Err(DefuseError::PublicKeyExists(account_id, public_key));
149 }
150 Ok(())
151 }
152
153 fn remove_public_key(&mut self, account_id: AccountId, public_key: PublicKey) -> Result<()> {
154 let had = self.view.has_public_key(&account_id, &public_key);
155 let account = self
156 .accounts
157 .get_or_create(account_id.clone(), |account_id| {
158 self.view.is_account_locked(account_id)
159 })
160 .get_mut()
161 .ok_or_else(|| DefuseError::AccountLocked(account_id.clone()))?;
162 let removed = if had {
163 account.public_keys_removed.insert(public_key)
164 } else {
165 account.public_keys_added.remove(&public_key)
166 };
167 if !removed {
168 return Err(DefuseError::PublicKeyNotExist(account_id, public_key));
169 }
170 Ok(())
171 }
172
173 fn commit_nonce(&mut self, account_id: AccountId, nonce: Nonce) -> Result<()> {
174 if self.view.is_nonce_used(&account_id, nonce) {
175 return Err(DefuseError::NonceUsed);
176 }
177
178 self.accounts
179 .get_or_create(account_id.clone(), |account_id| {
180 self.view.is_account_locked(account_id)
181 })
182 .get_mut()
183 .ok_or(DefuseError::AccountLocked(account_id))?
184 .commit_nonce(nonce)
185 }
186
187 fn cleanup_nonce_by_prefix(
188 &mut self,
189 account_id: &AccountIdRef,
190 prefix: NoncePrefix,
191 ) -> Result<bool> {
192 let account = self
193 .accounts
194 .get_mut(account_id)
195 .ok_or_else(|| DefuseError::AccountNotFound(account_id.to_owned()))?
196 .as_inner_unchecked_mut();
197
198 Ok(account.cleanup_nonce_by_prefix(prefix))
199 }
200
201 fn internal_add_balance(
202 &mut self,
203 owner_id: AccountId,
204 token_amounts: impl IntoIterator<Item = (TokenId, u128)>,
205 ) -> Result<()> {
206 let account = self
207 .accounts
208 .get_or_create(owner_id.clone(), |owner_id| {
209 self.view.is_account_locked(owner_id)
210 })
211 .as_inner_unchecked_mut();
212 for (token_id, amount) in token_amounts {
213 if account.token_amounts.get(&token_id).is_none() {
214 account
215 .token_amounts
216 .add(token_id.clone(), self.view.balance_of(&owner_id, &token_id))
217 .ok_or(DefuseError::BalanceOverflow)?;
218 }
219 account
220 .token_amounts
221 .add(token_id, amount)
222 .ok_or(DefuseError::BalanceOverflow)?;
223 }
224 Ok(())
225 }
226
227 fn internal_sub_balance(
228 &mut self,
229 owner_id: &AccountIdRef,
230 token_amounts: impl IntoIterator<Item = (TokenId, u128)>,
231 ) -> Result<()> {
232 let account = self
233 .accounts
234 .get_or_create(owner_id.to_owned(), |owner_id| {
235 self.view.is_account_locked(owner_id)
236 })
237 .get_mut()
238 .ok_or_else(|| DefuseError::AccountLocked(owner_id.to_owned()))?;
239 for (token_id, amount) in token_amounts {
240 if amount == 0 {
241 return Err(DefuseError::InvalidIntent);
242 }
243
244 if account.token_amounts.get(&token_id).is_none() {
245 account
246 .token_amounts
247 .add(token_id.clone(), self.view.balance_of(owner_id, &token_id))
248 .ok_or(DefuseError::BalanceOverflow)?;
249 }
250 account
251 .token_amounts
252 .sub(token_id, amount)
253 .ok_or(DefuseError::BalanceOverflow)?;
254 }
255 Ok(())
256 }
257
258 fn ft_withdraw(&mut self, owner_id: &AccountIdRef, withdraw: FtWithdraw) -> Result<()> {
259 self.internal_sub_balance(
260 owner_id,
261 std::iter::once((
262 Nep141TokenId::new(withdraw.token.clone()).into(),
263 withdraw.amount,
264 ))
265 .chain(withdraw.storage_deposit.map(|amount| {
266 (
267 Nep141TokenId::new(self.wnear_id().into_owned()).into(),
268 amount.as_yoctonear(),
269 )
270 })),
271 )
272 }
273
274 fn nft_withdraw(&mut self, owner_id: &AccountIdRef, withdraw: NftWithdraw) -> Result<()> {
275 self.internal_sub_balance(
276 owner_id,
277 std::iter::once((
278 Nep171TokenId::new(withdraw.token.clone(), withdraw.token_id.clone()).into(),
279 1,
280 ))
281 .chain(withdraw.storage_deposit.map(|amount| {
282 (
283 Nep141TokenId::new(self.wnear_id().into_owned()).into(),
284 amount.as_yoctonear(),
285 )
286 })),
287 )
288 }
289
290 fn mt_withdraw(&mut self, owner_id: &AccountIdRef, withdraw: MtWithdraw) -> Result<()> {
291 if withdraw.token_ids.len() != withdraw.amounts.len() || withdraw.token_ids.is_empty() {
292 return Err(DefuseError::InvalidIntent);
293 }
294
295 if withdraw.token_ids.len() > BATCH_MAX_TOKENS {
296 return Err(DefuseError::TooManyTokens(withdraw.token_ids.len()));
297 }
298
299 self.internal_sub_balance(
300 owner_id,
301 withdraw
302 .token_ids
303 .iter()
304 .cloned()
305 .map(|token_id| Nep245TokenId::new(withdraw.token.clone(), token_id))
306 .map(Into::into)
307 .zip(withdraw.amounts.iter().copied())
308 .chain(
309 withdraw
310 .storage_deposit
311 .map(|amount| (self.wnear_token_id(), amount.as_yoctonear())),
312 ),
313 )
314 }
315
316 fn native_withdraw(&mut self, owner_id: &AccountIdRef, withdraw: NativeWithdraw) -> Result<()> {
317 self.internal_sub_balance(
318 owner_id,
319 [(
320 Nep141TokenId::new(self.wnear_id().into_owned()).into(),
321 withdraw.amount.as_yoctonear(),
322 )],
323 )
324 }
325
326 #[inline]
328 fn notify_on_transfer(
329 &self,
330 _sender_id: &AccountIdRef,
331 _receiver_id: AccountId,
332 tokens: Amounts,
333 _notification: NotifyOnTransfer,
334 ) -> Result<()> {
335 if tokens.len() > BATCH_MAX_TOKENS {
336 return Err(DefuseError::TooManyTokens(tokens.len()));
337 }
338
339 Ok(())
340 }
341
342 fn storage_deposit(
343 &mut self,
344 owner_id: &AccountIdRef,
345 storage_deposit: StorageDeposit,
346 ) -> Result<()> {
347 self.internal_sub_balance(
348 owner_id,
349 [(
350 Nep141TokenId::new(self.wnear_id().into_owned()).into(),
351 storage_deposit.amount.as_yoctonear(),
352 )],
353 )
354 }
355
356 fn set_auth_by_predecessor_id(&mut self, account_id: AccountId, enable: bool) -> Result<bool> {
357 let was_enabled = self.is_auth_by_predecessor_id_enabled(&account_id);
358 let toggle = was_enabled ^ enable;
359 if toggle {
360 self.accounts
361 .get_or_create(account_id.clone(), |owner_id| {
362 self.view.is_account_locked(owner_id)
363 })
364 .get_mut()
365 .ok_or(DefuseError::AccountLocked(account_id))?
366 .auth_by_predecessor_id_toggled ^= true;
368 }
369 Ok(toggle)
370 }
371
372 fn auth_call(&mut self, signer_id: &AccountIdRef, auth_call: AuthCall) -> Result<()> {
373 if !auth_call.attached_deposit.is_zero() {
374 self.internal_sub_balance(
375 signer_id,
376 [(
377 Nep141TokenId::new(self.wnear_id().into_owned()).into(),
378 auth_call.attached_deposit.as_yoctonear(),
379 )],
380 )?;
381 }
382
383 Ok(())
384 }
385
386 #[inline]
387 fn mint(&mut self, owner_id: AccountId, tokens: Amounts, _memo: Option<String>) -> Result<()> {
388 self.internal_add_balance(owner_id, tokens)
389 }
390
391 #[inline]
392 fn burn(
393 &mut self,
394 owner_id: &AccountIdRef,
395 tokens: Amounts,
396 _memo: Option<String>,
397 ) -> Result<()> {
398 self.internal_sub_balance(owner_id, tokens)
399 }
400}
401
402#[derive(Debug, Default)]
403pub struct CachedAccounts(HashMap<AccountId, Lock<CachedAccount>>);
404
405impl CachedAccounts {
406 #[must_use]
407 #[inline]
408 pub fn new() -> Self {
409 Self(HashMap::new())
410 }
411
412 #[inline]
413 pub fn get(&self, account_id: &AccountIdRef) -> Option<&Lock<CachedAccount>> {
414 self.0.get(account_id)
415 }
416
417 #[inline]
418 pub fn get_mut(&mut self, account_id: &AccountIdRef) -> Option<&mut Lock<CachedAccount>> {
419 self.0.get_mut(account_id)
420 }
421
422 #[inline]
423 pub fn get_or_create(
424 &mut self,
425 account_id: AccountId,
426 is_initially_locked: impl FnOnce(&AccountId) -> bool,
427 ) -> &mut Lock<CachedAccount> {
428 self.0.entry(account_id).or_insert_with_key(|account_id| {
429 Lock::new(is_initially_locked(account_id), CachedAccount::default())
430 })
431 }
432}
433
434#[derive(Debug, Clone, Default)]
435pub struct CachedAccount {
436 nonces: Nonces<HashMap<U248, U256>>,
437
438 auth_by_predecessor_id_toggled: bool,
439
440 public_keys_added: HashSet<PublicKey>,
441 public_keys_removed: HashSet<PublicKey>,
442
443 token_amounts: Amounts<HashMap<TokenId, u128>>,
444}
445
446impl CachedAccount {
447 #[inline]
448 pub fn is_nonce_used(&self, nonce: U256) -> bool {
449 self.nonces.is_used(nonce)
450 }
451
452 #[inline]
453 pub fn commit_nonce(&mut self, n: U256) -> Result<()> {
454 self.nonces.commit(n)
455 }
456
457 #[inline]
458 pub fn cleanup_nonce_by_prefix(&mut self, prefix: NoncePrefix) -> bool {
459 self.nonces.cleanup_by_prefix(prefix)
460 }
461}