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