defuse_core/engine/state/
cached.rs1use 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 .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}