defuse_core/engine/state/
deltas.rs

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

defuse_core/engine/state/
deltas.rs
