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}