defuse_core/engine/state/
cached.rs

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