defuse_core/engine/state/
deltas.rs

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

defuse_core/engine/state/
deltas.rs
