1use alloc::collections::BTreeMap;
2use alloc::collections::btree_map::Entry;
3use alloc::string::ToString;
4use alloc::vec::Vec;
5
6use super::{
7 AccountDeltaError,
8 ByteReader,
9 ByteWriter,
10 Deserializable,
11 DeserializationError,
12 Serializable,
13};
14use crate::account::{AccountId, AccountType};
15use crate::asset::{Asset, FungibleAsset, NonFungibleAsset};
16use crate::{Felt, LexicographicWord, ONE, Word, ZERO};
17
18const DOMAIN_ASSET: Felt = Felt::new(1);
23
24#[derive(Clone, Debug, Default, PartialEq, Eq)]
31pub struct AccountVaultDelta {
32 fungible: FungibleAssetDelta,
33 non_fungible: NonFungibleAssetDelta,
34}
35
36impl AccountVaultDelta {
37 pub const fn new(fungible: FungibleAssetDelta, non_fungible: NonFungibleAssetDelta) -> Self {
43 Self { fungible, non_fungible }
44 }
45
46 pub fn fungible(&self) -> &FungibleAssetDelta {
48 &self.fungible
49 }
50
51 pub fn non_fungible(&self) -> &NonFungibleAssetDelta {
53 &self.non_fungible
54 }
55
56 pub fn is_empty(&self) -> bool {
58 self.fungible.is_empty() && self.non_fungible.is_empty()
59 }
60
61 pub fn add_asset(&mut self, asset: Asset) -> Result<(), AccountDeltaError> {
63 match asset {
64 Asset::Fungible(asset) => self.fungible.add(asset),
65 Asset::NonFungible(asset) => self.non_fungible.add(asset),
66 }
67 }
68
69 pub fn remove_asset(&mut self, asset: Asset) -> Result<(), AccountDeltaError> {
71 match asset {
72 Asset::Fungible(asset) => self.fungible.remove(asset),
73 Asset::NonFungible(asset) => self.non_fungible.remove(asset),
74 }
75 }
76
77 pub fn merge(&mut self, other: Self) -> Result<(), AccountDeltaError> {
84 self.non_fungible.merge(other.non_fungible)?;
85 self.fungible.merge(other.fungible)
86 }
87
88 pub(super) fn append_delta_elements(&self, elements: &mut Vec<Felt>) {
91 self.fungible().append_delta_elements(elements);
92 self.non_fungible().append_delta_elements(elements);
93 }
94}
95
96#[cfg(any(feature = "testing", test))]
97impl AccountVaultDelta {
98 pub fn from_iters(
100 added_assets: impl IntoIterator<Item = crate::asset::Asset>,
101 removed_assets: impl IntoIterator<Item = crate::asset::Asset>,
102 ) -> Self {
103 use crate::asset::Asset;
104
105 let mut fungible = FungibleAssetDelta::default();
106 let mut non_fungible = NonFungibleAssetDelta::default();
107
108 for asset in added_assets {
109 match asset {
110 Asset::Fungible(asset) => {
111 fungible.add(asset).unwrap();
112 },
113 Asset::NonFungible(asset) => {
114 non_fungible.add(asset).unwrap();
115 },
116 }
117 }
118
119 for asset in removed_assets {
120 match asset {
121 Asset::Fungible(asset) => {
122 fungible.remove(asset).unwrap();
123 },
124 Asset::NonFungible(asset) => {
125 non_fungible.remove(asset).unwrap();
126 },
127 }
128 }
129
130 Self { fungible, non_fungible }
131 }
132
133 pub fn added_assets(&self) -> impl Iterator<Item = crate::asset::Asset> + '_ {
135 use crate::asset::{Asset, FungibleAsset, NonFungibleAsset};
136 self.fungible
137 .0
138 .iter()
139 .filter(|&(_, &value)| value >= 0)
140 .map(|(&faucet_id, &diff)| {
141 Asset::Fungible(FungibleAsset::new(faucet_id, diff.unsigned_abs()).unwrap())
142 })
143 .chain(self.non_fungible.filter_by_action(NonFungibleDeltaAction::Add).map(|key| {
144 Asset::NonFungible(unsafe { NonFungibleAsset::new_unchecked(key.into()) })
145 }))
146 }
147
148 pub fn removed_assets(&self) -> impl Iterator<Item = crate::asset::Asset> + '_ {
150 use crate::asset::{Asset, FungibleAsset, NonFungibleAsset};
151 self.fungible
152 .0
153 .iter()
154 .filter(|&(_, &value)| value < 0)
155 .map(|(&faucet_id, &diff)| {
156 Asset::Fungible(FungibleAsset::new(faucet_id, diff.unsigned_abs()).unwrap())
157 })
158 .chain(self.non_fungible.filter_by_action(NonFungibleDeltaAction::Remove).map(|key| {
159 Asset::NonFungible(unsafe { NonFungibleAsset::new_unchecked(key.into()) })
160 }))
161 }
162}
163
164impl Serializable for AccountVaultDelta {
165 fn write_into<W: ByteWriter>(&self, target: &mut W) {
166 target.write(&self.fungible);
167 target.write(&self.non_fungible);
168 }
169
170 fn get_size_hint(&self) -> usize {
171 self.fungible.get_size_hint() + self.non_fungible.get_size_hint()
172 }
173}
174
175impl Deserializable for AccountVaultDelta {
176 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
177 let fungible = source.read()?;
178 let non_fungible = source.read()?;
179
180 Ok(Self::new(fungible, non_fungible))
181 }
182}
183
184#[derive(Clone, Debug, Default, PartialEq, Eq)]
189pub struct FungibleAssetDelta(BTreeMap<AccountId, i64>);
190
191impl FungibleAssetDelta {
192 pub fn new(map: BTreeMap<AccountId, i64>) -> Result<Self, AccountDeltaError> {
197 let delta = Self(map);
198 delta.validate()?;
199
200 Ok(delta)
201 }
202
203 pub fn add(&mut self, asset: FungibleAsset) -> Result<(), AccountDeltaError> {
208 let amount: i64 = asset.amount().try_into().expect("Amount it too high");
209 self.add_delta(asset.faucet_id(), amount)
210 }
211
212 pub fn remove(&mut self, asset: FungibleAsset) -> Result<(), AccountDeltaError> {
217 let amount: i64 = asset.amount().try_into().expect("Amount it too high");
218 self.add_delta(asset.faucet_id(), -amount)
219 }
220
221 pub fn amount(&self, faucet_id: &AccountId) -> Option<i64> {
223 self.0.get(faucet_id).copied()
224 }
225
226 pub fn num_assets(&self) -> usize {
228 self.0.len()
229 }
230
231 pub fn is_empty(&self) -> bool {
233 self.0.is_empty()
234 }
235
236 pub fn iter(&self) -> impl Iterator<Item = (&AccountId, &i64)> {
238 self.0.iter()
239 }
240
241 pub fn merge(&mut self, other: Self) -> Result<(), AccountDeltaError> {
248 for (&faucet_id, &amount) in other.0.iter() {
254 self.add_delta(faucet_id, amount)?;
255 }
256
257 Ok(())
258 }
259
260 fn add_delta(&mut self, faucet_id: AccountId, delta: i64) -> Result<(), AccountDeltaError> {
269 match self.0.entry(faucet_id) {
270 Entry::Vacant(entry) => {
271 if delta != 0 {
273 entry.insert(delta);
274 }
275 },
276 Entry::Occupied(mut entry) => {
277 let old = *entry.get();
278 let new = old.checked_add(delta).ok_or(
279 AccountDeltaError::FungibleAssetDeltaOverflow {
280 faucet_id,
281 current: old,
282 delta,
283 },
284 )?;
285
286 if new == 0 {
287 entry.remove();
288 } else {
289 *entry.get_mut() = new;
290 }
291 },
292 }
293
294 Ok(())
295 }
296
297 fn validate(&self) -> Result<(), AccountDeltaError> {
302 for faucet_id in self.0.keys() {
303 if !matches!(faucet_id.account_type(), AccountType::FungibleFaucet) {
304 return Err(AccountDeltaError::NotAFungibleFaucetId(*faucet_id));
305 }
306 }
307
308 Ok(())
309 }
310
311 pub(super) fn append_delta_elements(&self, elements: &mut Vec<Felt>) {
322 for (faucet_id, amount_delta) in self.iter() {
323 debug_assert_ne!(
326 *amount_delta, 0,
327 "fungible asset iterator should never yield amount deltas of 0"
328 );
329
330 let asset = FungibleAsset::new(*faucet_id, amount_delta.unsigned_abs())
331 .expect("absolute amount delta should be less than i64::MAX");
332 let was_added = if *amount_delta > 0 { ONE } else { ZERO };
333
334 elements.extend_from_slice(&[DOMAIN_ASSET, was_added, ZERO, ZERO]);
335 elements.extend_from_slice(Word::from(asset).as_elements());
336 }
337 }
338}
339
340impl Serializable for FungibleAssetDelta {
341 fn write_into<W: ByteWriter>(&self, target: &mut W) {
342 target.write_usize(self.0.len());
343 target.write_many(self.0.iter().map(|(&faucet_id, &delta)| (faucet_id, delta as u64)));
347 }
348
349 fn get_size_hint(&self) -> usize {
350 self.0.len().get_size_hint() + self.0.len() * FungibleAsset::SERIALIZED_SIZE
351 }
352}
353
354impl Deserializable for FungibleAssetDelta {
355 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
356 let num_fungible_assets = source.read_usize()?;
357 let map = source
361 .read_many::<(AccountId, u64)>(num_fungible_assets)?
362 .into_iter()
363 .map(|(account_id, delta_as_u64)| (account_id, delta_as_u64 as i64))
364 .collect();
365
366 Self::new(map).map_err(|err| DeserializationError::InvalidValue(err.to_string()))
367 }
368}
369
370#[derive(Clone, Debug, Default, PartialEq, Eq)]
378pub struct NonFungibleAssetDelta(
379 BTreeMap<LexicographicWord<NonFungibleAsset>, NonFungibleDeltaAction>,
380);
381
382impl NonFungibleAssetDelta {
383 pub const fn new(
385 map: BTreeMap<LexicographicWord<NonFungibleAsset>, NonFungibleDeltaAction>,
386 ) -> Self {
387 Self(map)
388 }
389
390 pub fn add(&mut self, asset: NonFungibleAsset) -> Result<(), AccountDeltaError> {
395 self.apply_action(asset, NonFungibleDeltaAction::Add)
396 }
397
398 pub fn remove(&mut self, asset: NonFungibleAsset) -> Result<(), AccountDeltaError> {
403 self.apply_action(asset, NonFungibleDeltaAction::Remove)
404 }
405
406 pub fn num_assets(&self) -> usize {
408 self.0.len()
409 }
410
411 pub fn is_empty(&self) -> bool {
413 self.0.is_empty()
414 }
415
416 pub fn iter(&self) -> impl Iterator<Item = (&NonFungibleAsset, &NonFungibleDeltaAction)> {
418 self.0.iter().map(|(key, value)| (key.inner(), value))
419 }
420
421 pub fn merge(&mut self, other: Self) -> Result<(), AccountDeltaError> {
428 for (&key, &action) in other.0.iter() {
430 self.apply_action(key.into_inner(), action)?;
431 }
432
433 Ok(())
434 }
435
436 fn apply_action(
445 &mut self,
446 asset: NonFungibleAsset,
447 action: NonFungibleDeltaAction,
448 ) -> Result<(), AccountDeltaError> {
449 match self.0.entry(LexicographicWord::new(asset)) {
450 Entry::Vacant(entry) => {
451 entry.insert(action);
452 },
453 Entry::Occupied(entry) => {
454 let previous = *entry.get();
455 if previous == action {
456 return Err(AccountDeltaError::DuplicateNonFungibleVaultUpdate(asset));
458 }
459 entry.remove();
461 },
462 }
463
464 Ok(())
465 }
466
467 fn filter_by_action(
469 &self,
470 action: NonFungibleDeltaAction,
471 ) -> impl Iterator<Item = NonFungibleAsset> + '_ {
472 self.0
473 .iter()
474 .filter(move |&(_, cur_action)| cur_action == &action)
475 .map(|(key, _)| key.into_inner())
476 }
477
478 pub(super) fn append_delta_elements(&self, elements: &mut Vec<Felt>) {
481 for (asset, action) in self.iter() {
482 let was_added = match action {
483 NonFungibleDeltaAction::Remove => ZERO,
484 NonFungibleDeltaAction::Add => ONE,
485 };
486
487 elements.extend_from_slice(&[DOMAIN_ASSET, was_added, ZERO, ZERO]);
488 elements.extend_from_slice(Word::from(*asset).as_elements());
489 }
490 }
491}
492
493impl Serializable for NonFungibleAssetDelta {
494 fn write_into<W: ByteWriter>(&self, target: &mut W) {
495 let added: Vec<_> = self.filter_by_action(NonFungibleDeltaAction::Add).collect();
496 let removed: Vec<_> = self.filter_by_action(NonFungibleDeltaAction::Remove).collect();
497
498 target.write_usize(added.len());
499 target.write_many(added.iter());
500
501 target.write_usize(removed.len());
502 target.write_many(removed.iter());
503 }
504
505 fn get_size_hint(&self) -> usize {
506 let added = self.filter_by_action(NonFungibleDeltaAction::Add).count();
507 let removed = self.filter_by_action(NonFungibleDeltaAction::Remove).count();
508
509 added.get_size_hint()
510 + removed.get_size_hint()
511 + added * NonFungibleAsset::SERIALIZED_SIZE
512 + removed * NonFungibleAsset::SERIALIZED_SIZE
513 }
514}
515
516impl Deserializable for NonFungibleAssetDelta {
517 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
518 let mut map = BTreeMap::new();
519
520 let num_added = source.read_usize()?;
521 for _ in 0..num_added {
522 let added_asset = source.read()?;
523 map.insert(LexicographicWord::new(added_asset), NonFungibleDeltaAction::Add);
524 }
525
526 let num_removed = source.read_usize()?;
527 for _ in 0..num_removed {
528 let removed_asset = source.read()?;
529 map.insert(LexicographicWord::new(removed_asset), NonFungibleDeltaAction::Remove);
530 }
531
532 Ok(Self::new(map))
533 }
534}
535
536#[derive(Clone, Copy, Debug, PartialEq, Eq)]
537pub enum NonFungibleDeltaAction {
538 Add,
539 Remove,
540}
541
542#[cfg(test)]
546mod tests {
547 use super::{AccountVaultDelta, Deserializable, Serializable};
548 use crate::account::{AccountId, AccountIdPrefix};
549 use crate::asset::{Asset, FungibleAsset, NonFungibleAsset, NonFungibleAssetDetails};
550 use crate::testing::account_id::{
551 ACCOUNT_ID_PRIVATE_FUNGIBLE_FAUCET,
552 ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET,
553 };
554
555 #[test]
556 fn test_serde_account_vault() {
557 let asset_0 = FungibleAsset::mock(100);
558 let asset_1 = NonFungibleAsset::mock(&[10, 21, 32, 43]);
559 let delta = AccountVaultDelta::from_iters([asset_0], [asset_1]);
560
561 let serialized = delta.to_bytes();
562 let deserialized = AccountVaultDelta::read_from_bytes(&serialized).unwrap();
563 assert_eq!(deserialized, delta);
564 }
565
566 #[test]
567 fn test_is_empty_account_vault() {
568 let faucet = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap();
569 let asset: Asset = FungibleAsset::new(faucet, 123).unwrap().into();
570
571 assert!(AccountVaultDelta::default().is_empty());
572 assert!(!AccountVaultDelta::from_iters([asset], []).is_empty());
573 assert!(!AccountVaultDelta::from_iters([], [asset]).is_empty());
574 }
575
576 #[rstest::rstest]
577 #[case::pos_pos(50, 50, Some(100))]
578 #[case::neg_neg(-50, -50, Some(-100))]
579 #[case::empty_pos(0, 50, Some(50))]
580 #[case::empty_neg(0, -50, Some(-50))]
581 #[case::nullify_pos_neg(100, -100, Some(0))]
582 #[case::nullify_neg_pos(-100, 100, Some(0))]
583 #[case::overflow(FungibleAsset::MAX_AMOUNT as i64, FungibleAsset::MAX_AMOUNT as i64, None)]
584 #[case::underflow(-(FungibleAsset::MAX_AMOUNT as i64), -(FungibleAsset::MAX_AMOUNT as i64), None)]
585 #[test]
586 fn merge_fungible_aggregation(#[case] x: i64, #[case] y: i64, #[case] expected: Option<i64>) {
587 fn create_delta_with_fungible(account_id: AccountId, amount: i64) -> AccountVaultDelta {
590 let asset = FungibleAsset::new(account_id, amount.unsigned_abs()).unwrap().into();
591 match amount {
592 0 => AccountVaultDelta::default(),
593 x if x.is_positive() => AccountVaultDelta::from_iters([asset], []),
594 _ => AccountVaultDelta::from_iters([], [asset]),
595 }
596 }
597
598 let account_id = AccountId::try_from(ACCOUNT_ID_PRIVATE_FUNGIBLE_FAUCET).unwrap();
599
600 let mut delta_x = create_delta_with_fungible(account_id, x);
601 let delta_y = create_delta_with_fungible(account_id, y);
602
603 let result = delta_x.merge(delta_y);
604
605 if let Some(expected) = expected {
607 let expected = create_delta_with_fungible(account_id, expected);
608 assert_eq!(result.map(|_| delta_x).unwrap(), expected);
609 } else {
610 assert!(result.is_err());
611 }
612 }
613
614 #[rstest::rstest]
615 #[case::empty_removed(None, Some(false), Ok(Some(false)))]
616 #[case::empty_added(None, Some(true), Ok(Some(true)))]
617 #[case::add_remove(Some(true), Some(false), Ok(None))]
618 #[case::remove_add(Some(false), Some(true), Ok(None))]
619 #[case::double_add(Some(true), Some(true), Err(()))]
620 #[case::double_remove(Some(false), Some(false), Err(()))]
621 #[test]
622 fn merge_non_fungible_aggregation(
623 #[case] x: Option<bool>,
624 #[case] y: Option<bool>,
625 #[case] expected: Result<Option<bool>, ()>,
626 ) {
627 fn create_delta_with_non_fungible(
630 account_id_prefix: AccountIdPrefix,
631 added: Option<bool>,
632 ) -> AccountVaultDelta {
633 let asset: Asset = NonFungibleAsset::new(
634 &NonFungibleAssetDetails::new(account_id_prefix, vec![1, 2, 3]).unwrap(),
635 )
636 .unwrap()
637 .into();
638
639 match added {
640 Some(true) => AccountVaultDelta::from_iters([asset], []),
641 Some(false) => AccountVaultDelta::from_iters([], [asset]),
642 None => AccountVaultDelta::default(),
643 }
644 }
645
646 let account_id = NonFungibleAsset::mock_issuer().prefix();
647
648 let mut delta_x = create_delta_with_non_fungible(account_id, x);
649 let delta_y = create_delta_with_non_fungible(account_id, y);
650
651 let result = delta_x.merge(delta_y);
652
653 if let Ok(expected) = expected {
654 let expected = create_delta_with_non_fungible(account_id, expected);
655 assert_eq!(result.map(|_| delta_x).unwrap(), expected);
656 } else {
657 assert!(result.is_err());
658 }
659 }
660}