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#[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 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#[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 pub fn finalize_into(self, token_id: &TokenId, transfers: &mut Transfers) -> Result<(), i128> {
295 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 let (mut deposit, mut withdraw) = (deposits.next(), withdrawals.next());
304
305 while let (Some((sender, send)), Some((receiver, receive))) =
307 (withdraw.as_mut(), deposit.as_mut())
308 {
309 let transfer = (*send).min(*receive);
311 transfers
312 .transfer(sender.clone(), receiver.clone(), token_id.clone(), transfer)
313 .unwrap_or_else(|| unreachable!());
316
317 *send = send.saturating_sub(transfer);
319 *receive = receive.saturating_sub(transfer);
320
321 if *send == 0 {
322 withdraw = withdrawals.next();
324 }
325 if *receive == 0 {
326 deposit = deposits.next();
328 }
329 }
330
331 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 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#[must_use]
353#[derive(Debug, Default, PartialEq, Eq)]
354pub struct Transfers(
355 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}