penumbra_sdk_asset/
balance.rs

1use anyhow::anyhow;
2use ark_r1cs_std::prelude::*;
3use ark_r1cs_std::uint8::UInt8;
4use ark_relations::r1cs::SynthesisError;
5use penumbra_sdk_num::{Amount, AmountVar};
6use serde::{Deserialize, Serialize};
7use std::{
8    collections::{btree_map, BTreeMap},
9    fmt::{self, Debug, Formatter},
10    iter::FusedIterator,
11    mem,
12    num::NonZeroU128,
13    ops::{Add, AddAssign, Deref, Neg, Sub, SubAssign},
14};
15
16use crate::{
17    asset::{AssetIdVar, Id},
18    value::ValueVar,
19    Value,
20};
21
22pub mod commitment;
23pub use commitment::Commitment;
24
25mod imbalance;
26mod iter;
27use commitment::VALUE_BLINDING_GENERATOR;
28use decaf377::{r1cs::ElementVar, Fq, Fr};
29use imbalance::{Imbalance, Sign};
30
31use self::commitment::BalanceCommitmentVar;
32use penumbra_sdk_proto::{penumbra::core::asset::v1 as pb, DomainType};
33
34/// A `Balance` is a "vector of [`Value`]s", where some values may be required, while others may be
35/// provided. For a transaction to be valid, its balance must be zero.
36#[derive(Clone, Eq, Default, Serialize, Deserialize)]
37#[serde(try_from = "pb::Balance", into = "pb::Balance")]
38pub struct Balance {
39    pub negated: bool,
40    pub balance: BTreeMap<Id, Imbalance<NonZeroU128>>,
41}
42
43impl Balance {
44    fn from_signed_value(negated: bool, value: Value) -> Option<Self> {
45        let non_zero = NonZeroU128::try_from(value.amount.value()).ok()?;
46        Some(Self {
47            negated: false,
48            balance: BTreeMap::from([(
49                value.asset_id,
50                if negated {
51                    Imbalance::Required(non_zero)
52                } else {
53                    Imbalance::Provided(non_zero)
54                },
55            )]),
56        })
57    }
58}
59
60impl DomainType for Balance {
61    type Proto = pb::Balance;
62}
63
64/// Serialization should normalize the `Balance`, where the top-level
65/// negated field is excluded during serialization. Rather, the
66/// sign information is captured in the `SignedValue` pairs.  
67///
68/// Since the underlying BTreeMap can't hold multiple imbalances for
69/// the same asset ID, we implement an ordering-agnostic accumulation
70/// scheme that explicitly combines imbalances.
71impl TryFrom<pb::Balance> for Balance {
72    type Error = anyhow::Error;
73
74    fn try_from(balance: pb::Balance) -> Result<Self, Self::Error> {
75        let mut out = Self::default();
76        for v in balance.values {
77            let value = v.value.ok_or_else(|| anyhow!("missing value"))?;
78            if let Some(b) = Balance::from_signed_value(v.negated, Value::try_from(value)?) {
79                out += b;
80            }
81        }
82        Ok(out)
83    }
84}
85
86impl From<Balance> for pb::Balance {
87    fn from(v: Balance) -> Self {
88        let values = v
89            .balance
90            .into_iter()
91            .map(|(id, imbalance)| {
92                // Decompose imbalance into it sign and magnitude, and convert
93                // magnitude into raw amount and determine negation based on the sign.
94                let (sign, magnitude) = if v.negated { -imbalance } else { imbalance }.into_inner();
95                let amount = u128::from(magnitude);
96
97                pb::balance::SignedValue {
98                    value: Some(pb::Value {
99                        asset_id: Some(id.into()),
100                        amount: Some(Amount::from(amount).into()),
101                    }),
102                    negated: matches!(sign, Sign::Required),
103                }
104            })
105            .collect();
106
107        pb::Balance { values }
108    }
109}
110
111impl Debug for Balance {
112    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
113        f.debug_struct("Balance")
114            .field("required", &self.required().collect::<Vec<_>>())
115            .field("provided", &self.provided().collect::<Vec<_>>())
116            .finish()
117    }
118}
119
120impl Balance {
121    /// Make a new, zero balance.
122    pub fn zero() -> Self {
123        Self::default()
124    }
125
126    /// Check if this balance is zero.
127    pub fn is_zero(&self) -> bool {
128        self.balance.is_empty()
129    }
130
131    /// Find out how many distinct assets are represented in this balance.
132    pub fn dimension(&self) -> usize {
133        self.balance.len()
134    }
135
136    /// Iterate over all the requirements of the balance, as [`Value`]s.
137    pub fn required(
138        &self,
139    ) -> impl Iterator<Item = Value> + DoubleEndedIterator + FusedIterator + '_ {
140        self.iter().filter_map(Imbalance::required)
141    }
142
143    // Iterate over all the provisions of the balance, as [`Value`]s.
144    pub fn provided(
145        &self,
146    ) -> impl Iterator<Item = Value> + DoubleEndedIterator + FusedIterator + '_ {
147        self.iter().filter_map(Imbalance::provided)
148    }
149
150    /// Commit to a [`Balance`] using a provided blinding factor.
151    ///
152    /// This is like a vectorized [`Value::commit`].
153    #[allow(non_snake_case)]
154    pub fn commit(&self, blinding_factor: Fr) -> Commitment {
155        // Accumulate all the elements for the values
156        let mut commitment = decaf377::Element::default();
157        for imbalance in self.iter() {
158            let (sign, value) = imbalance.into_inner();
159            let G_v = value.asset_id.value_generator();
160
161            // Depending on the sign, either subtract or add
162            match sign {
163                imbalance::Sign::Required => {
164                    commitment -= G_v * Fr::from(value.amount);
165                }
166                imbalance::Sign::Provided => {
167                    commitment += G_v * Fr::from(value.amount);
168                }
169            }
170        }
171
172        // Add the blinding factor only once, after the accumulation
173        commitment += blinding_factor * VALUE_BLINDING_GENERATOR.deref();
174        Commitment(commitment)
175    }
176}
177
178impl PartialEq for Balance {
179    // Eq is implemented this way because there are two different representations for a `Balance`,
180    // to allow fast negation, so we check elements of the iterator against each other, because the
181    // iterator returns the values in canonical imbalance representation, in order
182    fn eq(&self, other: &Self) -> bool {
183        if self.dimension() != other.dimension() {
184            return false;
185        }
186
187        for (i, j) in self.iter().zip(other.iter()) {
188            if i != j {
189                return false;
190            }
191        }
192
193        true
194    }
195}
196
197impl Neg for Balance {
198    type Output = Self;
199
200    fn neg(self) -> Self {
201        Self {
202            negated: !self.negated,
203            balance: self.balance,
204        }
205    }
206}
207
208impl Add for Balance {
209    type Output = Self;
210
211    // This is a tricky function, because the representation of a `Balance` has a `negated` flag
212    // which inverts the meaning of the stored entry (this is so that you can negate balances in
213    // constant time, which makes subtraction fast to implement). As a consequence, however, we have
214    // to take care that when we access the raw storage, we negate the imbalance we retrieve if and
215    // only if we are in negated mode, and when we write back a value, we negate it again on writing
216    // it back if we are in negated mode.
217    fn add(mut self, mut other: Self) -> Self {
218        // Always iterate through the smaller of the two
219        if other.dimension() > self.dimension() {
220            mem::swap(&mut self, &mut other);
221        }
222
223        for imbalance in other.into_iter() {
224            // Convert back into an asset id key and imbalance value
225            let (sign, Value { asset_id, amount }) = imbalance.into_inner();
226            let (asset_id, mut imbalance) = if let Some(amount) = NonZeroU128::new(amount.into()) {
227                (asset_id, sign.imbalance(amount))
228            } else {
229                unreachable!("values stored in balance are always nonzero")
230            };
231
232            match self.balance.entry(asset_id) {
233                btree_map::Entry::Vacant(entry) => {
234                    // Important: if we are currently negated, we have to negate the imbalance
235                    // before we store it!
236                    if self.negated {
237                        imbalance = -imbalance;
238                    }
239                    entry.insert(imbalance);
240                }
241                btree_map::Entry::Occupied(mut entry) => {
242                    // Important: if we are currently negated, we have to negate the entry we just
243                    // pulled out!
244                    let mut existing_imbalance = *entry.get();
245                    if self.negated {
246                        existing_imbalance = -existing_imbalance;
247                    }
248
249                    if let Some(mut new_imbalance) = existing_imbalance + imbalance {
250                        // If there's still an imbalance, update the map entry, making sure to
251                        // negate the new imbalance if we are negated
252                        if self.negated {
253                            new_imbalance = -new_imbalance;
254                        }
255                        entry.insert(new_imbalance);
256                    } else {
257                        // If adding this imbalance zeroed out the balance for this asset, remove
258                        // the entry
259                        entry.remove();
260                    }
261                }
262            }
263        }
264
265        self
266    }
267}
268
269impl Add<Value> for Balance {
270    type Output = Balance;
271
272    fn add(self, value: Value) -> Self::Output {
273        self + Balance::from(value)
274    }
275}
276
277impl AddAssign for Balance {
278    fn add_assign(&mut self, other: Self) {
279        *self = mem::take(self) + other;
280    }
281}
282
283impl AddAssign<Value> for Balance {
284    fn add_assign(&mut self, other: Value) {
285        *self += Balance::from(other);
286    }
287}
288
289impl Sub for Balance {
290    type Output = Self;
291
292    fn sub(self, other: Self) -> Self {
293        self + -other
294    }
295}
296
297impl Sub<Value> for Balance {
298    type Output = Balance;
299
300    fn sub(self, value: Value) -> Self::Output {
301        self - Balance::from(value)
302    }
303}
304
305impl SubAssign for Balance {
306    fn sub_assign(&mut self, other: Self) {
307        *self = mem::take(self) - other;
308    }
309}
310
311impl SubAssign<Value> for Balance {
312    fn sub_assign(&mut self, other: Value) {
313        *self -= Balance::from(other);
314    }
315}
316
317impl From<Value> for Balance {
318    fn from(Value { amount, asset_id }: Value) -> Self {
319        let mut balance = BTreeMap::new();
320        if let Some(amount) = NonZeroU128::new(amount.into()) {
321            balance.insert(asset_id, Imbalance::Provided(amount));
322        }
323        Balance {
324            negated: false,
325            balance,
326        }
327    }
328}
329
330/// Represents a balance in a rank 1 constraint system.
331///
332/// A balance consists of a number of assets (represented
333/// by their asset ID), the amount of each asset, as
334/// well as a boolean var that represents their contribution to the
335/// transaction's balance.
336///
337/// True values represent assets that are being provided (positive sign).
338/// False values represent assets that are required (negative sign).
339#[derive(Clone)]
340pub struct BalanceVar {
341    pub inner: Vec<(AssetIdVar, (Boolean<Fq>, AmountVar))>,
342}
343
344impl AllocVar<Balance, Fq> for BalanceVar {
345    fn new_variable<T: std::borrow::Borrow<Balance>>(
346        cs: impl Into<ark_relations::r1cs::Namespace<Fq>>,
347        f: impl FnOnce() -> Result<T, SynthesisError>,
348        mode: ark_r1cs_std::prelude::AllocationMode,
349    ) -> Result<Self, SynthesisError> {
350        let ns = cs.into();
351        let cs = ns.cs();
352        let inner1 = f()?;
353        let inner = inner1.borrow();
354        match mode {
355            AllocationMode::Constant => unimplemented!(),
356            AllocationMode::Input => unimplemented!(),
357            AllocationMode::Witness => {
358                if !inner.negated {
359                    unimplemented!();
360                }
361
362                let mut inner_balance_vars = Vec::new();
363                for (asset_id, imbalance) in inner.balance.iter() {
364                    let (sign, amount) = imbalance.into_inner();
365
366                    let asset_id_var = AssetIdVar::new_witness(cs.clone(), || Ok(asset_id))?;
367                    let amount_var = AmountVar::new_witness(cs.clone(), || {
368                        Ok(Amount::from(u128::from(amount)))
369                    })?;
370
371                    let boolean_var = match sign {
372                        imbalance::Sign::Required => Boolean::constant(false),
373                        imbalance::Sign::Provided => Boolean::constant(true),
374                    };
375
376                    inner_balance_vars.push((asset_id_var, (boolean_var, amount_var)));
377                }
378
379                Ok(BalanceVar {
380                    inner: inner_balance_vars,
381                })
382            }
383        }
384    }
385}
386
387impl From<ValueVar> for BalanceVar {
388    fn from(ValueVar { amount, asset_id }: ValueVar) -> Self {
389        let mut balance_vec = Vec::new();
390        let sign = Boolean::constant(true);
391        balance_vec.push((asset_id, (sign, amount)));
392
393        BalanceVar { inner: balance_vec }
394    }
395}
396
397impl BalanceVar {
398    /// Commit to a [`BalanceVar`] using a provided blinding factor.
399    ///
400    /// This is like a vectorized [`ValueVar::commit`].
401    #[allow(non_snake_case)]
402    #[allow(clippy::assign_op_pattern)]
403    pub fn commit(
404        &self,
405        blinding_factor: Vec<UInt8<Fq>>,
406    ) -> Result<BalanceCommitmentVar, SynthesisError> {
407        // Access constraint system ref from one of the balance contributions
408        let cs = self
409            .inner
410            .get(0)
411            .expect("at least one contribution to balance")
412            .0
413            .asset_id
414            .cs();
415
416        // Begin by adding the blinding factor only once
417        let value_blinding_generator = ElementVar::new_constant(cs, *VALUE_BLINDING_GENERATOR)?;
418        let mut commitment =
419            value_blinding_generator.scalar_mul_le(blinding_factor.to_bits_le()?.iter())?;
420
421        // Accumulate all the elements for the values
422        for (asset_id, (sign, amount)) in self.inner.iter() {
423            let G_v = asset_id.value_generator()?;
424            // Access the inner `FqVar` on `AmountVar` for scalar mul
425            let value_amount = amount.amount.clone();
426
427            // We scalar mul first with value (small), _then_ negate [v]G_v if needed
428            let vG = G_v.scalar_mul_le(value_amount.to_bits_le()?.iter())?;
429            let minus_vG = vG.negate()?;
430            let to_add = ElementVar::conditionally_select(sign, &vG, &minus_vG)?;
431            // It seems like the AddAssign impl here doesn't match the Add impl
432            commitment = commitment + to_add;
433        }
434        Ok(BalanceCommitmentVar { inner: commitment })
435    }
436
437    /// Create a balance from a positive [`ValueVar`].
438    pub fn from_positive_value_var(value: ValueVar) -> Self {
439        value.into()
440    }
441
442    /// Create a balance from a negated [`ValueVar`].
443    pub fn from_negative_value_var(value: ValueVar) -> Self {
444        let mut balance_vec = Vec::new();
445        let sign = Boolean::constant(false);
446        balance_vec.push((value.asset_id, (sign, value.amount)));
447
448        BalanceVar { inner: balance_vec }
449    }
450}
451
452impl std::ops::Add for BalanceVar {
453    type Output = Self;
454
455    fn add(self, other: Self) -> Self {
456        let mut balance_vec = self.inner;
457        for (asset_id, (sign, amount)) in other.inner {
458            balance_vec.push((asset_id, (sign, amount)));
459        }
460        BalanceVar { inner: balance_vec }
461    }
462}
463
464#[cfg(test)]
465mod test {
466    use crate::{
467        asset::{self, Metadata},
468        STAKING_TOKEN_ASSET_ID,
469    };
470    use ark_ff::Zero;
471    use decaf377::Fr;
472    use once_cell::sync::Lazy;
473    use penumbra_sdk_proto::core::num::v1::Amount as ProtoAmount;
474    use proptest::prelude::*;
475    use rand_core::OsRng;
476
477    use super::*;
478
479    #[test]
480    fn provide_then_require() {
481        let mut balance = Balance::zero();
482        balance += Value {
483            amount: 1u64.into(),
484            asset_id: *STAKING_TOKEN_ASSET_ID,
485        };
486        balance -= Value {
487            amount: 1u64.into(),
488            asset_id: *STAKING_TOKEN_ASSET_ID,
489        };
490        assert!(balance.is_zero());
491    }
492
493    #[test]
494    fn require_then_provide() {
495        let mut balance = Balance::zero();
496        balance -= Value {
497            amount: 1u64.into(),
498            asset_id: *STAKING_TOKEN_ASSET_ID,
499        };
500        balance += Value {
501            amount: 1u64.into(),
502            asset_id: *STAKING_TOKEN_ASSET_ID,
503        };
504        assert!(balance.is_zero());
505    }
506
507    #[test]
508    fn provide_then_require_negative_zero() {
509        let mut balance = -Balance::zero();
510        balance += Value {
511            amount: 1u64.into(),
512            asset_id: *STAKING_TOKEN_ASSET_ID,
513        };
514        balance -= Value {
515            amount: 1u64.into(),
516            asset_id: *STAKING_TOKEN_ASSET_ID,
517        };
518        assert!(balance.is_zero());
519    }
520
521    #[test]
522    fn require_then_provide_negative_zero() {
523        let mut balance = -Balance::zero();
524        balance -= Value {
525            amount: 1u64.into(),
526            asset_id: *STAKING_TOKEN_ASSET_ID,
527        };
528        balance += Value {
529            amount: 1u64.into(),
530            asset_id: *STAKING_TOKEN_ASSET_ID,
531        };
532        assert!(balance.is_zero());
533    }
534
535    #[derive(Debug, Clone)]
536    enum Expression {
537        Value(Value),
538        Neg(Box<Expression>),
539        Add(Box<Expression>, Box<Expression>),
540        Sub(Box<Expression>, Box<Expression>),
541    }
542
543    impl Expression {
544        fn transparent_balance_commitment(&self) -> Commitment {
545            match self {
546                Expression::Value(value) => value.commit(Fr::zero()),
547                Expression::Neg(expr) => -expr.transparent_balance_commitment(),
548                Expression::Add(lhs, rhs) => {
549                    lhs.transparent_balance_commitment() + rhs.transparent_balance_commitment()
550                }
551                Expression::Sub(lhs, rhs) => {
552                    lhs.transparent_balance_commitment() - rhs.transparent_balance_commitment()
553                }
554            }
555        }
556
557        fn balance(&self) -> Balance {
558            match self {
559                Expression::Value(value) => Balance::from(*value),
560                Expression::Neg(expr) => -expr.balance(),
561                Expression::Add(lhs, rhs) => lhs.balance() + rhs.balance(),
562                Expression::Sub(lhs, rhs) => lhs.balance() - rhs.balance(),
563            }
564        }
565    }
566
567    // Two sample denom/asset id pairs, for testing
568    static DENOM_1: Lazy<Metadata> = Lazy::new(|| {
569        crate::asset::Cache::with_known_assets()
570            .get_unit("cube")
571            .unwrap()
572            .base()
573    });
574    static ASSET_ID_1: Lazy<Id> = Lazy::new(|| DENOM_1.id());
575
576    static DENOM_2: Lazy<Metadata> = Lazy::new(|| {
577        crate::asset::Cache::with_known_assets()
578            .get_unit("nala")
579            .unwrap()
580            .base()
581    });
582    static ASSET_ID_2: Lazy<Id> = Lazy::new(|| DENOM_2.id());
583
584    #[allow(clippy::arc_with_non_send_sync)]
585    fn gen_expression() -> impl proptest::strategy::Strategy<Value = Expression> {
586        (
587            (0u64..u32::MAX as u64), // limit amounts so that there is no overflow
588            prop_oneof![Just(*ASSET_ID_1), Just(*ASSET_ID_2)],
589        )
590            .prop_map(|(amount, asset_id)| {
591                Expression::Value(Value {
592                    amount: amount.into(),
593                    asset_id,
594                })
595            })
596            .prop_recursive(8, 256, 2, |inner| {
597                prop_oneof![
598                    inner
599                        .clone()
600                        .prop_map(|beneath| Expression::Neg(Box::new(beneath))),
601                    (inner.clone(), inner.clone()).prop_map(|(left, right)| {
602                        Expression::Add(Box::new(left), Box::new(right))
603                    }),
604                    (inner.clone(), inner).prop_map(|(left, right)| {
605                        Expression::Sub(Box::new(left), Box::new(right))
606                    }),
607                ]
608            })
609    }
610
611    proptest! {
612        /// Checks to make sure that any possible expression made of negation, addition, and
613        /// subtraction is a homomorphism with regard to the resultant balance commitment, which
614        /// should provide assurance that these operations are implemented correctly on the balance
615        /// type itself.
616        #[test]
617        fn all_expressions_correct_commitment(
618            expr in gen_expression()
619        ) {
620            // Compute the balance for the expression
621            let balance = expr.balance();
622
623            // Compute the transparent commitment for the expression
624            let commitment = expr.transparent_balance_commitment();
625
626            // Compute the transparent commitment for the balance
627            let mut balance_commitment = Commitment::default();
628            for required in balance.required() {
629                balance_commitment = balance_commitment - required.commit(Fr::zero());
630            }
631            for provided in balance.provided() {
632                balance_commitment = balance_commitment + provided.commit(Fr::zero());
633            }
634
635            assert_eq!(commitment, balance_commitment);
636        }
637    }
638
639    /// Implement fallible conversion (protobuf to domain type) for multiple entries
640    /// with the same asset ID.
641    #[test]
642    fn try_from_fallible_conversion_same_asset_id() {
643        let proto_balance_0 = pb::Balance {
644            values: vec![
645                pb::balance::SignedValue {
646                    value: Some(pb::Value {
647                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
648                        amount: Some(Amount::from(100u128).into()),
649                    }),
650                    negated: true,
651                },
652                pb::balance::SignedValue {
653                    value: Some(pb::Value {
654                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
655                        amount: Some(Amount::from(50u128).into()),
656                    }),
657                    negated: false,
658                },
659            ],
660        };
661
662        let proto_balance_1 = pb::Balance {
663            values: vec![
664                pb::balance::SignedValue {
665                    value: Some(pb::Value {
666                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
667                        amount: Some(Amount::from(100u128).into()),
668                    }),
669                    negated: true,
670                },
671                pb::balance::SignedValue {
672                    value: Some(pb::Value {
673                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
674                        amount: Some(Amount::from(200u128).into()),
675                    }),
676                    negated: false,
677                },
678            ],
679        };
680
681        let proto_balance_2 = pb::Balance {
682            values: vec![
683                pb::balance::SignedValue {
684                    value: Some(pb::Value {
685                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
686                        amount: Some(Amount::from(100u128).into()),
687                    }),
688                    negated: true,
689                },
690                pb::balance::SignedValue {
691                    value: Some(pb::Value {
692                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
693                        amount: Some(Amount::from(200u128).into()),
694                    }),
695                    negated: true,
696                },
697            ],
698        };
699
700        let proto_balance_3 = pb::Balance {
701            values: vec![
702                pb::balance::SignedValue {
703                    value: Some(pb::Value {
704                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
705                        amount: Some(Amount::from(100u128).into()),
706                    }),
707                    negated: false,
708                },
709                pb::balance::SignedValue {
710                    value: Some(pb::Value {
711                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
712                        amount: Some(Amount::from(50u128).into()),
713                    }),
714                    negated: true,
715                },
716            ],
717        };
718
719        let proto_balance_4 = pb::Balance {
720            values: vec![
721                pb::balance::SignedValue {
722                    value: Some(pb::Value {
723                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
724                        amount: Some(Amount::from(100u128).into()),
725                    }),
726                    negated: false,
727                },
728                pb::balance::SignedValue {
729                    value: Some(pb::Value {
730                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
731                        amount: Some(Amount::from(200u128).into()),
732                    }),
733                    negated: true,
734                },
735            ],
736        };
737
738        let proto_balance_5 = pb::Balance {
739            values: vec![
740                pb::balance::SignedValue {
741                    value: Some(pb::Value {
742                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
743                        amount: Some(Amount::from(100u128).into()),
744                    }),
745                    negated: false,
746                },
747                pb::balance::SignedValue {
748                    value: Some(pb::Value {
749                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
750                        amount: Some(Amount::from(200u128).into()),
751                    }),
752                    negated: false,
753                },
754            ],
755        };
756
757        let balance_0 = Balance::try_from(proto_balance_0).expect("fallible conversion");
758        let balance_1 = Balance::try_from(proto_balance_1).expect("fallible conversion");
759        let balance_2 = Balance::try_from(proto_balance_2).expect("fallible conversion");
760        let balance_3 = Balance::try_from(proto_balance_3).expect("fallible conversion");
761        let balance_4 = Balance::try_from(proto_balance_4).expect("fallible conversion");
762        let balance_5 = Balance::try_from(proto_balance_5).expect("fallible conversion");
763
764        assert!(matches!(
765            balance_0.balance.get(&STAKING_TOKEN_ASSET_ID),
766            Some(Imbalance::Required(amount)) if amount == &NonZeroU128::new(50).unwrap()
767        ));
768
769        assert!(matches!(
770            balance_1.balance.get(&STAKING_TOKEN_ASSET_ID),
771            Some(Imbalance::Provided(amount)) if amount == &NonZeroU128::new(100).unwrap()
772        ));
773
774        assert!(matches!(
775            balance_2.balance.get(&STAKING_TOKEN_ASSET_ID),
776            Some(Imbalance::Required(amount)) if amount == &NonZeroU128::new(300).unwrap()
777        ));
778
779        assert!(matches!(
780            balance_3.balance.get(&STAKING_TOKEN_ASSET_ID),
781            Some(Imbalance::Provided(amount)) if amount == &NonZeroU128::new(50).unwrap()
782        ));
783
784        assert!(matches!(
785            balance_4.balance.get(&STAKING_TOKEN_ASSET_ID),
786            Some(Imbalance::Required(amount)) if amount == &NonZeroU128::new(100).unwrap()
787        ));
788
789        assert!(matches!(
790            balance_5.balance.get(&STAKING_TOKEN_ASSET_ID),
791            Some(Imbalance::Provided(amount)) if amount == &NonZeroU128::new(300).unwrap()
792        ));
793    }
794
795    /// Implement fallible conversion (protobuf to domain type) for multiple entries
796    /// with different asset IDs.
797    #[test]
798    fn try_from_fallible_conversion_different_asset_id() {
799        let rand_asset_id = Id(Fq::rand(&mut OsRng));
800
801        let proto_balance = pb::Balance {
802            values: vec![
803                pb::balance::SignedValue {
804                    value: Some(pb::Value {
805                        asset_id: Some((*STAKING_TOKEN_ASSET_ID).into()),
806                        amount: Some(Amount::from(100u128).into()),
807                    }),
808                    negated: true,
809                },
810                pb::balance::SignedValue {
811                    value: Some(pb::Value {
812                        asset_id: Some(rand_asset_id.into()),
813                        amount: Some(Amount::from(50u128).into()),
814                    }),
815                    negated: false,
816                },
817            ],
818        };
819
820        let balance = Balance::try_from(proto_balance).expect("fallible conversion");
821
822        assert!(matches!(
823            balance.balance.get(&STAKING_TOKEN_ASSET_ID),
824            Some(Imbalance::Required(amount)) if amount == &NonZeroU128::new(100).unwrap()
825        ));
826
827        assert!(matches!(
828            balance.balance.get(&rand_asset_id),
829            Some(Imbalance::Provided(amount)) if amount == &NonZeroU128::new(50).unwrap()
830        ));
831    }
832
833    /// Implement fallible conversion (protobuf to domain type) with missing fields.
834    #[test]
835    fn try_from_fallible_conversion_failure() {
836        let proto_balance = pb::Balance {
837            values: vec![pb::balance::SignedValue {
838                value: None,
839                negated: false,
840            }],
841        };
842
843        assert!(Balance::try_from(proto_balance).is_err());
844    }
845
846    /// Implement infallible conversion (domain type to protobuf).
847    #[test]
848    fn from_infallible_conversion() {
849        let rand_asset_id = Id(Fq::rand(&mut OsRng));
850
851        let balance = Balance {
852            negated: false,
853            balance: [
854                (
855                    *STAKING_TOKEN_ASSET_ID,
856                    Imbalance::Provided(NonZeroU128::new(100).unwrap()),
857                ),
858                (
859                    rand_asset_id,
860                    Imbalance::Required(NonZeroU128::new(200).unwrap()),
861                ),
862            ]
863            .iter()
864            .cloned()
865            .collect(),
866        };
867
868        let proto_balance: pb::Balance = balance.into();
869
870        let first_value = proto_balance
871            .values
872            .iter()
873            .find(|v| v.value.as_ref().unwrap().asset_id == Some((*STAKING_TOKEN_ASSET_ID).into()))
874            .expect("asset should exist");
875        let second_value = proto_balance
876            .values
877            .iter()
878            .find(|v| v.value.as_ref().unwrap().asset_id == Some(rand_asset_id.into()))
879            .expect("asset should exist");
880
881        assert_eq!(proto_balance.values.len(), 2);
882
883        assert_eq!(
884            first_value.value.as_ref().unwrap().asset_id,
885            Some((*STAKING_TOKEN_ASSET_ID).into())
886        );
887        assert_eq!(
888            second_value.value.as_ref().unwrap().asset_id,
889            Some(rand_asset_id.into())
890        );
891
892        let proto_amount: ProtoAmount = Amount::from(100u128).into();
893        assert_eq!(
894            first_value.value.as_ref().unwrap().amount,
895            Some(proto_amount)
896        );
897
898        let proto_amount: ProtoAmount = Amount::from(200u128).into();
899        assert_eq!(
900            second_value.value.as_ref().unwrap().amount,
901            Some(proto_amount)
902        );
903    }
904
905    fn test_balance_serialization_round_tripping_example(
906        parts: Vec<(u8, i64)>,
907    ) -> anyhow::Result<()> {
908        let proto = pb::Balance {
909            values: parts
910                .into_iter()
911                .map(|(asset, amount)| pb::balance::SignedValue {
912                    value: Some(pb::Value {
913                        amount: Some(Amount::from(amount.abs() as u64).into()),
914                        asset_id: Some(asset::Id::try_from([asset; 32]).unwrap().into()),
915                    }),
916                    negated: amount < 0,
917                })
918                .collect(),
919        };
920        let p_to_d = Balance::try_from(proto)?;
921        let p_to_d_to_p = pb::Balance::from(p_to_d.clone());
922        let p_to_d_to_p_to_d = Balance::try_from(p_to_d_to_p)?;
923        assert_eq!(p_to_d, p_to_d_to_p_to_d);
924        Ok(())
925    }
926
927    proptest! {
928        #[test]
929        fn test_balance_serialization_roundtripping(parts in proptest::collection::vec(((0u8..16u8), any::<i64>()), 0..100)) {
930            // To get better errors
931            assert!(matches!(test_balance_serialization_round_tripping_example(parts), Ok(_)));
932        }
933    }
934}
penumbra_sdk_asset/balance.rs

penumbra_sdk_asset/
balance.rs
