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    #[inline]
218    fn mint(&mut self, owner_id: AccountId, tokens: Amounts, memo: Option<String>) -> Result<()> {
219        self.state.mint(owner_id, tokens, memo)
220    }
221
222    #[inline]
223    fn burn(
224        &mut self,
225        owner_id: &AccountIdRef,
226        tokens: Amounts,
227        memo: Option<String>,
228    ) -> Result<()> {
229        self.state.burn(owner_id, tokens, memo)
230    }
231}
232
233/// Accumulates internal deposits and withdrawals on different tokens
234/// to match transfers using `.finalize()`
235///
236/// Transfers in `TokenDiff` intents are represented as deltas without receivers.
237/// This struct accumulates tokens all transfers, and converts them from deltas, to
238/// a set of transfers from one account to another.
239/// Note that this doesn't touch account balances. The balances were already changed
240/// in an earlier stage while executing the intent.
241#[derive(Debug, Default)]
242pub struct TransferMatcher(HashMap<TokenId, TokenTransferMatcher>);
243
244impl TransferMatcher {
245    #[inline]
246    pub fn new() -> Self {
247        Self(HashMap::new())
248    }
249
250    #[inline]
251    pub fn deposit(&mut self, owner_id: AccountId, token_id: TokenId, amount: u128) -> bool {
252        self.0.entry_or_default(token_id).deposit(owner_id, amount)
253    }
254
255    #[inline]
256    pub fn withdraw(&mut self, owner_id: AccountId, token_id: TokenId, amount: u128) -> bool {
257        self.0.entry_or_default(token_id).withdraw(owner_id, amount)
258    }
259
260    #[inline]
261    pub fn add_delta(&mut self, owner_id: AccountId, token_id: TokenId, delta: i128) -> bool {
262        self.0.entry_or_default(token_id).add_delta(owner_id, delta)
263    }
264
265    // Finalizes all transfers, or returns unmatched deltas.
266    // If unmatched deltas overflow, then Err(None) is returned.
267    pub fn finalize(self) -> Result<Transfers, InvariantViolated> {
268        let mut transfers = Transfers::default();
269        let mut deltas = TokenDeltas::default();
270        for (token_id, transfer_matcher) in self.0 {
271            if let Err(unmatched) = transfer_matcher.finalize_into(&token_id, &mut transfers) {
272                if unmatched == 0 || deltas.apply_delta(token_id, unmatched).is_none() {
273                    return Err(InvariantViolated::Overflow);
274                }
275            }
276        }
277        if !deltas.is_empty() {
278            return Err(InvariantViolated::UnmatchedDeltas {
279                unmatched_deltas: deltas,
280            });
281        }
282        Ok(transfers)
283    }
284}
285
286type AccountAmounts = Amounts<HashMap<AccountId, u128>>;
287
288// Accumulates internal deposits and withdrawals on a single token
289#[derive(Debug, Default, PartialEq, Eq)]
290pub struct TokenTransferMatcher {
291    deposits: AccountAmounts,
292    withdrawals: AccountAmounts,
293}
294
295impl TokenTransferMatcher {
296    #[inline]
297    pub fn deposit(&mut self, owner_id: AccountId, amount: u128) -> bool {
298        Self::sub_add(&mut self.withdrawals, &mut self.deposits, owner_id, amount)
299    }
300
301    #[inline]
302    pub fn withdraw(&mut self, owner_id: AccountId, amount: u128) -> bool {
303        Self::sub_add(&mut self.deposits, &mut self.withdrawals, owner_id, amount)
304    }
305
306    #[inline]
307    pub fn add_delta(&mut self, owner_id: AccountId, delta: i128) -> bool {
308        let amount = delta.unsigned_abs();
309        if delta.is_negative() {
310            self.withdraw(owner_id, amount)
311        } else {
312            self.deposit(owner_id, amount)
313        }
314    }
315
316    fn sub_add(
317        sub: &mut AccountAmounts,
318        add: &mut AccountAmounts,
319        owner_id: AccountId,
320        mut amount: u128,
321    ) -> bool {
322        let s = sub.amount_for(&owner_id);
323        if s > 0 {
324            let a = s.min(amount);
325            sub.sub(owner_id.clone(), a)
326                .unwrap_or_else(|| unreachable!());
327            amount = amount.saturating_sub(a);
328            if amount == 0 {
329                return true;
330            }
331        }
332        add.add(owner_id, amount).is_some()
333    }
334
335    // Finalizes transfer of this token, or returns unmatched delta.
336    // If returned delta is zero, then overflow happened
337    pub fn finalize_into(self, token_id: &TokenId, transfers: &mut Transfers) -> Result<(), i128> {
338        // sort deposits and withdrawals in descending order
339        let [mut deposits, mut withdrawals] = [self.deposits, self.withdrawals].map(|amounts| {
340            let mut amounts: Vec<_> = amounts.into_iter().collect();
341            amounts.sort_unstable_by_key(|(_, amount)| Reverse(*amount));
342            amounts.into_iter()
343        });
344
345        // take first sender and receiver
346        let (mut deposit, mut withdraw) = (deposits.next(), withdrawals.next());
347
348        // as long as there is both: sender and receiver
349        while let (Some((sender, send)), Some((receiver, receive))) =
350            (withdraw.as_mut(), deposit.as_mut())
351        {
352            // get min amount and transfer
353            let transfer = (*send).min(*receive);
354            transfers
355                .transfer(sender.clone(), receiver.clone(), token_id.clone(), transfer)
356                // no error can happen since we add only one transfer for each
357                // combination of (sender, receiver, token_id)
358                .unwrap_or_else(|| unreachable!());
359
360            // subtract amount from sender and receiver
361            *send = send.saturating_sub(transfer);
362            *receive = receive.saturating_sub(transfer);
363
364            if *send == 0 {
365                // select next sender
366                withdraw = withdrawals.next();
367            }
368            if *receive == 0 {
369                // select next receiver
370                deposit = deposits.next();
371            }
372        }
373
374        // only sender(s) left
375        if let Some((_, send)) = withdraw {
376            return Err(withdrawals
377                .try_fold(send, |total, (_, s)| total.checked_add(s))
378                .and_then(|total| i128::try_from(total).ok())
379                .and_then(i128::checked_neg)
380                .unwrap_or_default());
381        }
382        // only receiver(s) left
383        if let Some((_, receive)) = deposit {
384            return Err(deposits
385                .try_fold(receive, |total, (_, r)| total.checked_add(r))
386                .and_then(|total| i128::try_from(total).ok())
387                .unwrap_or_default());
388        }
389
390        Ok(())
391    }
392}
393
394/// Raw transfers between accounts
395#[must_use]
396#[derive(Debug, Default, PartialEq, Eq)]
397pub struct Transfers(
398    /// `sender_id` -> `receiver_id` -> `token_id` -> `amount`
399    HashMap<AccountId, HashMap<AccountId, Amounts<HashMap<TokenId, u128>>>>,
400);
401
402impl Transfers {
403    #[must_use]
404    pub fn transfer(
405        &mut self,
406        sender_id: AccountId,
407        receiver_id: AccountId,
408        token_id: TokenId,
409        amount: u128,
410    ) -> Option<u128> {
411        let mut sender = self.0.entry_or_default(sender_id);
412        let mut receiver = sender.entry_or_default(receiver_id);
413        receiver.add(token_id, amount)
414    }
415
416    pub fn with_transfer(
417        mut self,
418        sender_id: AccountId,
419        receiver_id: AccountId,
420        token_id: TokenId,
421        amount: u128,
422    ) -> Option<Self> {
423        self.transfer(sender_id, receiver_id, token_id, amount)?;
424        Some(self)
425    }
426
427    pub fn as_mt_event(&self) -> Option<MtEvent<'_>> {
428        if self.0.is_empty() {
429            return None;
430        }
431        Some(MtEvent::MtTransfer(
432            self.0
433                .iter()
434                .flat_map(|(sender_id, transfers)| iter::repeat(sender_id).zip(transfers))
435                .map(|(sender_id, (receiver_id, transfers))| {
436                    let (token_ids, amounts) = transfers
437                        .iter()
438                        .map(|(token_id, amount)| (token_id.to_string(), U128(*amount)))
439                        .unzip();
440                    MtTransferEvent {
441                        authorized_id: None,
442                        old_owner_id: Cow::Borrowed(sender_id),
443                        new_owner_id: Cow::Borrowed(receiver_id),
444                        token_ids: Cow::Owned(token_ids),
445                        amounts: Cow::Owned(amounts),
446                        memo: None,
447                    }
448                })
449                .collect::<Vec<_>>()
450                .into(),
451        ))
452    }
453}
454
455#[near(serializers = [json])]
456#[serde(tag = "error", rename_all = "snake_case")]
457#[derive(Debug, Clone, PartialEq, Eq)]
458pub enum InvariantViolated {
459    UnmatchedDeltas {
460        #[serde_as(as = "Amounts<BTreeMap<_, DisplayFromStr>>")]
461        unmatched_deltas: TokenDeltas,
462    },
463    Overflow,
464}
465
466impl InvariantViolated {
467    #[inline]
468    pub const fn as_unmatched_deltas(&self) -> Option<&TokenDeltas> {
469        match self {
470            Self::UnmatchedDeltas {
471                unmatched_deltas: deltas,
472            } => Some(deltas),
473            Self::Overflow => None,
474        }
475    }
476
477    #[inline]
478    pub fn into_unmatched_deltas(self) -> Option<TokenDeltas> {
479        match self {
480            Self::UnmatchedDeltas {
481                unmatched_deltas: deltas,
482            } => Some(deltas),
483            Self::Overflow => None,
484        }
485    }
486}
487
488#[cfg(test)]
489#[allow(clippy::many_single_char_names)]
490mod tests {
491    use crate::token_id::nep141::Nep141TokenId;
492
493    use super::*;
494
495    #[test]
496    fn test_transfers() {
497        let mut transfers = TransferMatcher::default();
498        let [a, b, c, d, e, f, g]: [AccountId; 7] =
499            ["a", "b", "c", "d", "e", "f", "g"].map(|s| format!("{s}.near").parse().unwrap());
500        let [ft1, ft2] = ["ft1", "ft2"].map(|a| {
501            TokenId::from(Nep141TokenId::new(
502                format!("{a}.near").parse::<AccountId>().unwrap(),
503            ))
504        });
505
506        let deltas: HashMap<AccountId, TokenDeltas> = [
507            (&a, [(&ft1, -5), (&ft2, 1)].as_slice()),
508            (&b, [(&ft1, 4), (&ft2, -1)].as_slice()),
509            (&c, [(&ft1, 3)].as_slice()),
510            (&d, [(&ft1, -10)].as_slice()),
511            (&e, [(&ft1, -1)].as_slice()),
512            (&f, [(&ft1, 10)].as_slice()),
513            (&g, [(&ft1, -1)].as_slice()),
514        ]
515        .into_iter()
516        .map(|(owner_id, deltas)| {
517            (
518                owner_id.clone(),
519                TokenDeltas::default()
520                    .with_apply_deltas(
521                        deltas
522                            .iter()
523                            .map(|(token_id, delta)| ((*token_id).clone(), *delta)),
524                    )
525                    .unwrap(),
526            )
527        })
528        .collect();
529
530        for (owner, (token_id, delta)) in deltas
531            .iter()
532            .flat_map(|(owner_id, deltas)| iter::repeat(owner_id).zip(deltas))
533        {
534            assert!(transfers.add_delta(owner.clone(), token_id.clone(), *delta));
535        }
536
537        let transfers = transfers.finalize().unwrap();
538        let mut new_deltas: HashMap<AccountId, TokenDeltas> = HashMap::new();
539
540        for (sender_id, transfers) in transfers.0 {
541            for (receiver_id, amounts) in transfers {
542                for (token_id, amount) in amounts {
543                    new_deltas
544                        .entry_or_default(sender_id.clone())
545                        .sub(token_id.clone(), amount)
546                        .unwrap();
547
548                    new_deltas
549                        .entry_or_default(receiver_id.clone())
550                        .add(token_id, amount)
551                        .unwrap();
552                }
553            }
554        }
555
556        assert_eq!(new_deltas, deltas);
557    }
558
559    #[test]
560    fn test_unmatched() {
561        let mut deltas = TransferMatcher::default();
562        let [a, b, _c, d, e, f, g]: [AccountId; 7] =
563            ["a", "b", "c", "d", "e", "f", "g"].map(|s| format!("{s}.near").parse().unwrap());
564        let [ft1, ft2] = ["ft1", "ft2"].map(|a| {
565            TokenId::from(Nep141TokenId::new(
566                format!("{a}.near").parse::<AccountId>().unwrap(),
567            ))
568        });
569
570        for (owner, token_id, delta) in [
571            (&a, &ft1, -5),
572            (&b, &ft1, 4),
573            (&d, &ft1, -10),
574            (&e, &ft1, -1),
575            (&f, &ft1, 10),
576            (&g, &ft1, -1),
577            (&a, &ft2, -1),
578        ] {
579            assert!(deltas.add_delta(owner.clone(), token_id.clone(), delta));
580        }
581
582        assert_eq!(
583            deltas.finalize().unwrap_err(),
584            InvariantViolated::UnmatchedDeltas {
585                unmatched_deltas: TokenDeltas::default()
586                    .with_apply_delta(ft1, -3)
587                    .unwrap()
588                    .with_apply_delta(ft2, -1)
589                    .unwrap()
590            }
591        );
592    }
593}
defuse_core/engine/state/deltas.rs

defuse_core/engine/state/
deltas.rs
