defuse_core/engine/state/
deltas.rs

1use crate::{
2    DefuseError, Nonce, NoncePrefix, Result, Salt,
3    amounts::Amounts,
4    fees::Pips,
5    intents::{
6        auth::AuthCall,
7        token_diff::TokenDeltas,
8        tokens::{
9            FtWithdraw, MtWithdraw, NativeWithdraw, NftWithdraw, NotifyOnTransfer, StorageDeposit,
10        },
11    },
12    token_id::TokenId,
13};
14use defuse_crypto::PublicKey;
15use defuse_map_utils::cleanup::DefaultMap;
16use defuse_nep245::{MtEvent, MtTransferEvent};
17use near_sdk::{AccountId, AccountIdRef, json_types::U128, near};
18use serde_with::{DisplayFromStr, serde_as};
19use std::{
20    borrow::Cow,
21    cmp::Reverse,
22    collections::{BTreeMap, HashMap},
23    iter,
24};
25
26use super::{State, StateView};
27
28pub struct Deltas<S> {
29    state: S,
30    deltas: TransferMatcher,
31}
32
33impl<S> Deltas<S> {
34    #[inline]
35    pub fn new(state: S) -> Self {
36        Self {
37            state,
38            deltas: TransferMatcher::new(),
39        }
40    }
41
42    #[inline]
43    pub fn finalize(self) -> Result<Transfers, InvariantViolated> {
44        self.deltas.finalize()
45    }
46}
47
48impl<S> StateView for Deltas<S>
49where
50    S: StateView,
51{
52    #[inline]
53    fn verifying_contract(&self) -> Cow<'_, AccountIdRef> {
54        self.state.verifying_contract()
55    }
56
57    #[inline]
58    fn wnear_id(&self) -> Cow<'_, AccountIdRef> {
59        self.state.wnear_id()
60    }
61
62    #[inline]
63    fn fee(&self) -> Pips {
64        self.state.fee()
65    }
66
67    #[inline]
68    fn fee_collector(&self) -> Cow<'_, AccountIdRef> {
69        self.state.fee_collector()
70    }
71
72    #[inline]
73    fn has_public_key(&self, account_id: &AccountIdRef, public_key: &PublicKey) -> bool {
74        self.state.has_public_key(account_id, public_key)
75    }
76
77    #[inline]
78    fn iter_public_keys(&self, account_id: &AccountIdRef) -> impl Iterator<Item = PublicKey> + '_ {
79        self.state.iter_public_keys(account_id)
80    }
81
82    #[inline]
83    fn is_nonce_used(&self, account_id: &AccountIdRef, nonce: Nonce) -> bool {
84        self.state.is_nonce_used(account_id, nonce)
85    }
86
87    #[inline]
88    fn balance_of(&self, account_id: &AccountIdRef, token_id: &TokenId) -> u128 {
89        self.state.balance_of(account_id, token_id)
90    }
91
92    #[inline]
93    fn is_account_locked(&self, account_id: &AccountIdRef) -> bool {
94        self.state.is_account_locked(account_id)
95    }
96
97    #[inline]
98    fn is_auth_by_predecessor_id_enabled(&self, account_id: &AccountIdRef) -> bool {
99        self.state.is_auth_by_predecessor_id_enabled(account_id)
100    }
101
102    #[inline]
103    fn is_valid_salt(&self, salt: Salt) -> bool {
104        self.state.is_valid_salt(salt)
105    }
106}
107
108impl<S> State for Deltas<S>
109where
110    S: State,
111{
112    #[inline]
113    fn add_public_key(&mut self, account_id: AccountId, public_key: PublicKey) -> Result<()> {
114        self.state.add_public_key(account_id, public_key)
115    }
116
117    #[inline]
118    fn remove_public_key(&mut self, account_id: AccountId, public_key: PublicKey) -> Result<()> {
119        self.state.remove_public_key(account_id, public_key)
120    }
121
122    #[inline]
123    fn commit_nonce(&mut self, account_id: AccountId, nonce: Nonce) -> Result<()> {
124        self.state.commit_nonce(account_id, nonce)
125    }
126
127    #[inline]
128    fn cleanup_nonce_by_prefix(
129        &mut self,
130        account_id: &AccountIdRef,
131        prefix: NoncePrefix,
132    ) -> Result<bool> {
133        self.state.cleanup_nonce_by_prefix(account_id, prefix)
134    }
135
136    fn internal_add_balance(
137        &mut self,
138        owner_id: AccountId,
139        tokens: impl IntoIterator<Item = (TokenId, u128)>,
140    ) -> Result<()> {
141        for (token_id, amount) in tokens {
142            self.state
143                .internal_add_balance(owner_id.clone(), [(token_id.clone(), amount)])?;
144            if !self.deltas.deposit(owner_id.clone(), token_id, amount) {
145                return Err(DefuseError::BalanceOverflow);
146            }
147        }
148        Ok(())
149    }
150
151    fn internal_sub_balance(
152        &mut self,
153        owner_id: &AccountIdRef,
154        tokens: impl IntoIterator<Item = (TokenId, u128)>,
155    ) -> Result<()> {
156        for (token_id, amount) in tokens {
157            self.state
158                .internal_sub_balance(owner_id, [(token_id.clone(), amount)])?;
159            if !self.deltas.withdraw(owner_id.to_owned(), token_id, amount) {
160                return Err(DefuseError::BalanceOverflow);
161            }
162        }
163        Ok(())
164    }
165
166    #[inline]
167    fn ft_withdraw(&mut self, owner_id: &AccountIdRef, withdraw: FtWithdraw) -> Result<()> {
168        self.state.ft_withdraw(owner_id, withdraw)
169    }
170
171    #[inline]
172    fn nft_withdraw(&mut self, owner_id: &AccountIdRef, withdraw: NftWithdraw) -> Result<()> {
173        self.state.nft_withdraw(owner_id, withdraw)
174    }
175
176    #[inline]
177    fn mt_withdraw(&mut self, owner_id: &AccountIdRef, withdraw: MtWithdraw) -> Result<()> {
178        self.state.mt_withdraw(owner_id, withdraw)
179    }
180
181    #[inline]
182    fn native_withdraw(&mut self, owner_id: &AccountIdRef, withdraw: NativeWithdraw) -> Result<()> {
183        self.state.native_withdraw(owner_id, withdraw)
184    }
185
186    #[inline]
187    fn notify_on_transfer(
188        &self,
189        sender_id: &AccountIdRef,
190        receiver_id: AccountId,
191        tokens: Amounts,
192        notification: NotifyOnTransfer,
193    ) {
194        self.state
195            .notify_on_transfer(sender_id, receiver_id, tokens, notification);
196    }
197
198    #[inline]
199    fn storage_deposit(
200        &mut self,
201        owner_id: &AccountIdRef,
202        storage_deposit: StorageDeposit,
203    ) -> Result<()> {
204        self.state.storage_deposit(owner_id, storage_deposit)
205    }
206
207    #[inline]
208    fn set_auth_by_predecessor_id(&mut self, account_id: AccountId, enable: bool) -> Result<bool> {
209        self.state.set_auth_by_predecessor_id(account_id, enable)
210    }
211
212    #[inline]
213    fn auth_call(&mut self, signer_id: &AccountIdRef, auth_call: AuthCall) -> Result<()> {
214        self.state.auth_call(signer_id, auth_call)
215    }
216}
217
218/// Accumulates internal deposits and withdrawals on different tokens
219/// to match transfers using `.finalize()`
220///
221/// Transfers in `TokenDiff` intents are represented as deltas without receivers.
222/// This struct accumulates tokens all transfers, and converts them from deltas, to
223/// a set of transfers from one account to another.
224/// Note that this doesn't touch account balances. The balances were already changed
225/// in an earlier stage while executing the intent.
226#[derive(Debug, Default)]
227pub struct TransferMatcher(HashMap<TokenId, TokenTransferMatcher>);
228
229impl TransferMatcher {
230    #[inline]
231    pub fn new() -> Self {
232        Self(HashMap::new())
233    }
234
235    #[inline]
236    pub fn deposit(&mut self, owner_id: AccountId, token_id: TokenId, amount: u128) -> bool {
237        self.0.entry_or_default(token_id).deposit(owner_id, amount)
238    }
239
240    #[inline]
241    pub fn withdraw(&mut self, owner_id: AccountId, token_id: TokenId, amount: u128) -> bool {
242        self.0.entry_or_default(token_id).withdraw(owner_id, amount)
243    }
244
245    #[inline]
246    pub fn add_delta(&mut self, owner_id: AccountId, token_id: TokenId, delta: i128) -> bool {
247        self.0.entry_or_default(token_id).add_delta(owner_id, delta)
248    }
249
250    // Finalizes all transfers, or returns unmatched deltas.
251    // If unmatched deltas overflow, then Err(None) is returned.
252    pub fn finalize(self) -> Result<Transfers, InvariantViolated> {
253        let mut transfers = Transfers::default();
254        let mut deltas = TokenDeltas::default();
255        for (token_id, transfer_matcher) in self.0 {
256            if let Err(unmatched) = transfer_matcher.finalize_into(&token_id, &mut transfers) {
257                if unmatched == 0 || deltas.apply_delta(token_id, unmatched).is_none() {
258                    return Err(InvariantViolated::Overflow);
259                }
260            }
261        }
262        if !deltas.is_empty() {
263            return Err(InvariantViolated::UnmatchedDeltas {
264                unmatched_deltas: deltas,
265            });
266        }
267        Ok(transfers)
268    }
269}
270
271type AccountAmounts = Amounts<HashMap<AccountId, u128>>;
272
273// Accumulates internal deposits and withdrawals on a single token
274#[derive(Debug, Default, PartialEq, Eq)]
275pub struct TokenTransferMatcher {
276    deposits: AccountAmounts,
277    withdrawals: AccountAmounts,
278}
279
280impl TokenTransferMatcher {
281    #[inline]
282    pub fn deposit(&mut self, owner_id: AccountId, amount: u128) -> bool {
283        Self::sub_add(&mut self.withdrawals, &mut self.deposits, owner_id, amount)
284    }
285
286    #[inline]
287    pub fn withdraw(&mut self, owner_id: AccountId, amount: u128) -> bool {
288        Self::sub_add(&mut self.deposits, &mut self.withdrawals, owner_id, amount)
289    }
290
291    #[inline]
292    pub fn add_delta(&mut self, owner_id: AccountId, delta: i128) -> bool {
293        let amount = delta.unsigned_abs();
294        if delta.is_negative() {
295            self.withdraw(owner_id, amount)
296        } else {
297            self.deposit(owner_id, amount)
298        }
299    }
300
301    fn sub_add(
302        sub: &mut AccountAmounts,
303        add: &mut AccountAmounts,
304        owner_id: AccountId,
305        mut amount: u128,
306    ) -> bool {
307        let s = sub.amount_for(&owner_id);
308        if s > 0 {
309            let a = s.min(amount);
310            sub.sub(owner_id.clone(), a)
311                .unwrap_or_else(|| unreachable!());
312            amount = amount.saturating_sub(a);
313            if amount == 0 {
314                return true;
315            }
316        }
317        add.add(owner_id, amount).is_some()
318    }
319
320    // Finalizes transfer of this token, or returns unmatched delta.
321    // If returned delta is zero, then overflow happened
322    pub fn finalize_into(self, token_id: &TokenId, transfers: &mut Transfers) -> Result<(), i128> {
323        // sort deposits and withdrawals in descending order
324        let [mut deposits, mut withdrawals] = [self.deposits, self.withdrawals].map(|amounts| {
325            let mut amounts: Vec<_> = amounts.into_iter().collect();
326            amounts.sort_unstable_by_key(|(_, amount)| Reverse(*amount));
327            amounts.into_iter()
328        });
329
330        // take first sender and receiver
331        let (mut deposit, mut withdraw) = (deposits.next(), withdrawals.next());
332
333        // as long as there is both: sender and receiver
334        while let (Some((sender, send)), Some((receiver, receive))) =
335            (withdraw.as_mut(), deposit.as_mut())
336        {
337            // get min amount and transfer
338            let transfer = (*send).min(*receive);
339            transfers
340                .transfer(sender.clone(), receiver.clone(), token_id.clone(), transfer)
341                // no error can happen since we add only one transfer for each
342                // combination of (sender, receiver, token_id)
343                .unwrap_or_else(|| unreachable!());
344
345            // subtract amount from sender and receiver
346            *send = send.saturating_sub(transfer);
347            *receive = receive.saturating_sub(transfer);
348
349            if *send == 0 {
350                // select next sender
351                withdraw = withdrawals.next();
352            }
353            if *receive == 0 {
354                // select next receiver
355                deposit = deposits.next();
356            }
357        }
358
359        // only sender(s) left
360        if let Some((_, send)) = withdraw {
361            return Err(withdrawals
362                .try_fold(send, |total, (_, s)| total.checked_add(s))
363                .and_then(|total| i128::try_from(total).ok())
364                .and_then(i128::checked_neg)
365                .unwrap_or_default());
366        }
367        // only receiver(s) left
368        if let Some((_, receive)) = deposit {
369            return Err(deposits
370                .try_fold(receive, |total, (_, r)| total.checked_add(r))
371                .and_then(|total| i128::try_from(total).ok())
372                .unwrap_or_default());
373        }
374
375        Ok(())
376    }
377}
378
379/// Raw transfers between accounts
380#[must_use]
381#[derive(Debug, Default, PartialEq, Eq)]
382pub struct Transfers(
383    /// `sender_id` -> `receiver_id` -> `token_id` -> `amount`
384    HashMap<AccountId, HashMap<AccountId, Amounts<HashMap<TokenId, u128>>>>,
385);
386
387impl Transfers {
388    #[must_use]
389    pub fn transfer(
390        &mut self,
391        sender_id: AccountId,
392        receiver_id: AccountId,
393        token_id: TokenId,
394        amount: u128,
395    ) -> Option<u128> {
396        let mut sender = self.0.entry_or_default(sender_id);
397        let mut receiver = sender.entry_or_default(receiver_id);
398        receiver.add(token_id, amount)
399    }
400
401    pub fn with_transfer(
402        mut self,
403        sender_id: AccountId,
404        receiver_id: AccountId,
405        token_id: TokenId,
406        amount: u128,
407    ) -> Option<Self> {
408        self.transfer(sender_id, receiver_id, token_id, amount)?;
409        Some(self)
410    }
411
412    pub fn as_mt_event(&self) -> Option<MtEvent<'_>> {
413        if self.0.is_empty() {
414            return None;
415        }
416        Some(MtEvent::MtTransfer(
417            self.0
418                .iter()
419                .flat_map(|(sender_id, transfers)| iter::repeat(sender_id).zip(transfers))
420                .map(|(sender_id, (receiver_id, transfers))| {
421                    let (token_ids, amounts) = transfers
422                        .iter()
423                        .map(|(token_id, amount)| (token_id.to_string(), U128(*amount)))
424                        .unzip();
425                    MtTransferEvent {
426                        authorized_id: None,
427                        old_owner_id: Cow::Borrowed(sender_id),
428                        new_owner_id: Cow::Borrowed(receiver_id),
429                        token_ids: Cow::Owned(token_ids),
430                        amounts: Cow::Owned(amounts),
431                        memo: None,
432                    }
433                })
434                .collect::<Vec<_>>()
435                .into(),
436        ))
437    }
438}
439
440#[near(serializers = [json])]
441#[serde(tag = "error", rename_all = "snake_case")]
442#[derive(Debug, Clone, PartialEq, Eq)]
443pub enum InvariantViolated {
444    UnmatchedDeltas {
445        #[serde_as(as = "Amounts<BTreeMap<_, DisplayFromStr>>")]
446        unmatched_deltas: TokenDeltas,
447    },
448    Overflow,
449}
450
451impl InvariantViolated {
452    #[inline]
453    pub const fn as_unmatched_deltas(&self) -> Option<&TokenDeltas> {
454        match self {
455            Self::UnmatchedDeltas {
456                unmatched_deltas: deltas,
457            } => Some(deltas),
458            Self::Overflow => None,
459        }
460    }
461
462    #[inline]
463    pub fn into_unmatched_deltas(self) -> Option<TokenDeltas> {
464        match self {
465            Self::UnmatchedDeltas {
466                unmatched_deltas: deltas,
467            } => Some(deltas),
468            Self::Overflow => None,
469        }
470    }
471}
472
473#[cfg(test)]
474#[allow(clippy::many_single_char_names)]
475mod tests {
476    use crate::token_id::nep141::Nep141TokenId;
477
478    use super::*;
479
480    #[test]
481    fn test_transfers() {
482        let mut transfers = TransferMatcher::default();
483        let [a, b, c, d, e, f, g]: [AccountId; 7] =
484            ["a", "b", "c", "d", "e", "f", "g"].map(|s| format!("{s}.near").parse().unwrap());
485        let [ft1, ft2] = ["ft1", "ft2"]
486            .map(|a| TokenId::from(Nep141TokenId::new(format!("{a}.near").parse().unwrap())));
487
488        let deltas: HashMap<AccountId, TokenDeltas> = [
489            (&a, [(&ft1, -5), (&ft2, 1)].as_slice()),
490            (&b, [(&ft1, 4), (&ft2, -1)].as_slice()),
491            (&c, [(&ft1, 3)].as_slice()),
492            (&d, [(&ft1, -10)].as_slice()),
493            (&e, [(&ft1, -1)].as_slice()),
494            (&f, [(&ft1, 10)].as_slice()),
495            (&g, [(&ft1, -1)].as_slice()),
496        ]
497        .into_iter()
498        .map(|(owner_id, deltas)| {
499            (
500                owner_id.clone(),
501                TokenDeltas::default()
502                    .with_apply_deltas(
503                        deltas
504                            .iter()
505                            .map(|(token_id, delta)| ((*token_id).clone(), *delta)),
506                    )
507                    .unwrap(),
508            )
509        })
510        .collect();
511
512        for (owner, (token_id, delta)) in deltas
513            .iter()
514            .flat_map(|(owner_id, deltas)| iter::repeat(owner_id).zip(deltas))
515        {
516            assert!(transfers.add_delta(owner.clone(), token_id.clone(), *delta));
517        }
518
519        let transfers = transfers.finalize().unwrap();
520        let mut new_deltas: HashMap<AccountId, TokenDeltas> = HashMap::new();
521
522        for (sender_id, transfers) in transfers.0 {
523            for (receiver_id, amounts) in transfers {
524                for (token_id, amount) in amounts {
525                    new_deltas
526                        .entry_or_default(sender_id.clone())
527                        .sub(token_id.clone(), amount)
528                        .unwrap();
529
530                    new_deltas
531                        .entry_or_default(receiver_id.clone())
532                        .add(token_id, amount)
533                        .unwrap();
534                }
535            }
536        }
537
538        assert_eq!(new_deltas, deltas);
539    }
540
541    #[test]
542    fn test_unmatched() {
543        let mut deltas = TransferMatcher::default();
544        let [a, b, _c, d, e, f, g]: [AccountId; 7] =
545            ["a", "b", "c", "d", "e", "f", "g"].map(|s| format!("{s}.near").parse().unwrap());
546        let [ft1, ft2] = ["ft1", "ft2"]
547            .map(|a| TokenId::from(Nep141TokenId::new(format!("{a}.near").parse().unwrap())));
548
549        for (owner, token_id, delta) in [
550            (&a, &ft1, -5),
551            (&b, &ft1, 4),
552            (&d, &ft1, -10),
553            (&e, &ft1, -1),
554            (&f, &ft1, 10),
555            (&g, &ft1, -1),
556            (&a, &ft2, -1),
557        ] {
558            assert!(deltas.add_delta(owner.clone(), token_id.clone(), delta));
559        }
560
561        assert_eq!(
562            deltas.finalize().unwrap_err(),
563            InvariantViolated::UnmatchedDeltas {
564                unmatched_deltas: TokenDeltas::default()
565                    .with_apply_delta(ft1, -3)
566                    .unwrap()
567                    .with_apply_delta(ft2, -1)
568                    .unwrap()
569            }
570        );
571    }
572}
defuse_core/engine/state/deltas.rs

defuse_core/engine/state/
deltas.rs
