Add balance struct (#22)

* add balance struct

* fix warning in rust

app/rust/.cargo/config.toml

@@ -18,7 +18,7 @@ rustflags = [
     "-C",
     "link-arg=-Tlink.ld",
     "-C",
-    "inline-threshold=0",
+    "llvm-args=--inline-threshold=0",
     "-C",
     "panic=abort",            # Add this line to disable unwinding
 ]
@@ -31,5 +31,7 @@ rustflags = [
     "-C",
     "link-arg=-Wl,--as-needed",
     "-C",
+    "llvm-args=--inline-threshold=0",
+    "-C",
     "panic=abort",                # Add this line to disable unwinding
 ]

app/rust/src/parser/fee.rs

@@ -2,7 +2,7 @@ use crate::constants::{AMOUNT_LEN_BYTES, ID_LEN_BYTES};
 use crate::parser::commitment::Commitment;
 use crate::parser::id::Id;
 use crate::parser::value::Sign;
-use crate::parser::value::{Value, ValueC};
+use crate::parser::value::{Value, ValueC, Imbalance, Balance};
 use crate::ParserError;
 use decaf377::Fq;
 use decaf377::Fr;
@@ -48,8 +48,12 @@ impl Fee {
     pub const LEN: usize = AMOUNT_LEN_BYTES + ID_LEN_BYTES;
 
     pub fn commit(&self, blinding: Fr) -> Result<Commitment, ParserError> {
-        let value = self.0.clone();
-        value.commit(blinding, Sign::Required)
+        let mut balance = Balance::new();
+        balance.add(Imbalance{
+            value: self.0.clone(),
+            sign: Sign::Required,
+        })?;
+        balance.commit(blinding)
     }
 
     pub fn to_bytes(&self) -> Result<[u8; Self::LEN], ParserError> {

app/rust/src/parser/penalty.rs

@@ -18,7 +18,7 @@ use crate::parser::{
     fixpoint::U128x128,
     ParserError,
     amount::Amount,
-    value::Imbalance,
+    value::{Balance, Imbalance, Sign},
     id::Id,
     value::Value,
     fee::STAKING_TOKEN_ASSET_ID_BYTES,
@@ -44,19 +44,25 @@ impl Penalty {
     /// This method takes the `unbonding_id` rather than the `UnbondingToken` so
     /// that it can be used in mock proof verification, where computation of the
     /// unbonding token's asset ID happens outside of the circuit.
-    pub fn balance_for_claim(&self, unbonding_id: Id, unbonding_amount: Amount) -> Imbalance {
+    pub fn balance_for_claim(&self, unbonding_id: Id, unbonding_amount: Amount) -> Result<Balance, ParserError> {
         // The undelegate claim action subtracts the unbonding amount and adds
         // the unbonded amount from the transaction's value balance.
-        Imbalance{
-            required_value: Value{
+        let mut balance = Balance::new();
+        balance.add(Imbalance{
+            value: Value{
                 amount: unbonding_amount,
                 asset_id: unbonding_id,
             },
-            provided_value: Value{
+            sign: Sign::Required,
+        })?;
+        balance.add(Imbalance{
+            value: Value{
                 amount: self.apply_to_amount(unbonding_amount),
                 asset_id: Id(Fq::from_le_bytes_mod_order(&STAKING_TOKEN_ASSET_ID_BYTES))
-            }
-        }
+            },
+            sign: Sign::Provided,
+        })?;
+        Ok(balance)
     }
 }
 

app/rust/src/parser/plans/output.rs

@@ -26,7 +26,7 @@ use crate::parser::{
     rseed::Rseed,
     symmetric::PayloadKey,
     symmetric::{OvkWrappedKey, WrappedMemoKey},
-    value::{Sign, Value, ValueC},
+    value::{Sign, Value, ValueC, Balance, Imbalance},
 };
 use crate::ParserError;
 use decaf377::Fr;
@@ -81,7 +81,7 @@ impl OutputPlanC {
         let ovk = fvk.outgoing();
         let note = self.output_note()?;
         let value = self.balance()?;
-        let balance_commitment = value.commit(self.get_value_blinding_fr()?, Sign::Required)?;
+        let balance_commitment = value.commit(self.get_value_blinding_fr()?)?;
 
         // Encrypt the note to the recipient...
         let esk = note.ephemeral_secret_key()?;
@@ -112,10 +112,13 @@ impl OutputPlanC {
         Note::from_parts(address, value, rseed)
     }
 
-    pub fn balance(&self) -> Result<Value, ParserError> {
-        // We should return a Balance struct here, but since we are currently managing only one value, it isn’t necessary for now
-        let value = Value::try_from(self.value.clone())?;
-        Ok(value)
+    pub fn balance(&self) -> Result<Balance, ParserError> {
+        let mut balance = Balance::new();
+        balance.add(Imbalance{
+            value: Value::try_from(self.value.clone())?,
+            sign: Sign::Required,
+        })?;
+        Ok(balance)
     }
 
     pub fn get_rseed(&self) -> Result<&[u8], ParserError> {

app/rust/src/parser/plans/spend.rs

@@ -21,7 +21,7 @@ use crate::parser::{
     effect_hash::{create_personalized_state, EffectHash},
     note::{Note, NoteC},
     nullifier::Nullifier,
-    value::{Sign, Value},
+    value::{Sign, Value, Balance, Imbalance},
 };
 use crate::ParserError;
 use decaf377::Fr;
@@ -67,16 +67,19 @@ impl SpendPlanC {
         Ok(Body {
             balance_commitment: self
                 .balance()?
-                .commit(self.get_value_blinding_fr()?, Sign::Provided)?,
+                .commit(self.get_value_blinding_fr()?)?,
             nullifier: self.nullifier(fvk)?,
             rk: self.rk(fvk)?,
         })
     }
 
-    pub fn balance(&self) -> Result<Value, ParserError> {
-        // We should return a Balance struct here, but since we are currently managing only one value, it isn’t necessary for now
-        let value = Value::try_from(self.note.value.clone())?;
-        Ok(value)
+    pub fn balance(&self) -> Result<Balance, ParserError> {
+        let mut balance = Balance::new();
+        balance.add(Imbalance{
+            value: Value::try_from(self.note.value.clone())?,
+            sign: Sign::Provided,
+        })?;
+        Ok(balance)
     }
 
     pub fn nullifier(&self, fvk: &FullViewingKey) -> Result<Nullifier, ParserError> {

app/rust/src/parser/plans/undelegate_claim.rs

@@ -21,7 +21,7 @@ use crate::parser::{
     penalty::{PenaltyC, Penalty},
     effect_hash::{create_personalized_state, EffectHash},
     commitment::Commitment,
-    value::Imbalance,
+    value::Balance,
     id::Id,
     id::AssetId,
 };
@@ -107,11 +107,11 @@ impl UndelegateClaimPlanC {
         Ok(Fr::from_le_bytes_mod_order(balance_blinding_bytes))
     }
 
-    pub fn balance(&self) -> Result<Imbalance, ParserError> {
+    pub fn balance(&self) -> Result<Balance, ParserError> {
         let penalty = Penalty::try_from(self.penalty.clone())?;
         let unbonding_amount = self.unbonding_amount.clone().try_into()?;
-        Ok(penalty
-            .balance_for_claim(self.unbonding_id()?, unbonding_amount))
+        penalty
+            .balance_for_claim(self.unbonding_id()?, unbonding_amount)
     }
 
     pub fn unbonding_id(&self) -> Result<Id, ParserError> {

app/rust/src/parser/value.rs

@@ -23,7 +23,8 @@ use crate::parser::{
 };
 use decaf377::Fr;
 
-// this should be in imbalance.rs. For now, it’s not necessary
+#[derive(Clone)]
+#[cfg_attr(any(feature = "derive-debug", test), derive(Debug))]
 pub enum Sign {
     Required,
     Provided,
@@ -40,35 +41,79 @@ pub struct Value {
 #[derive(Clone)]
 #[cfg_attr(any(feature = "derive-debug", test), derive(Debug))]
 pub struct Imbalance {
-    pub required_value: Value,
-    pub provided_value: Value,
+    pub value: Value,
+    pub sign: Sign,
 }
 
-// this should be implemented in the Balance, but since we are currently managing only one value, it isn’t necessary for now
-impl Value {
-    pub const LEN: usize = AMOUNT_LEN_BYTES + ID_LEN_BYTES;
-    pub fn commit(&self, blinding_factor: Fr, sign: Sign) -> Result<Commitment, ParserError> {
-        let mut commitment = decaf377::Element::IDENTITY;
-        let g_v = self.asset_id.value_generator();
-        let amount_fr: Fr = Into::into(self.amount);
+// Only two imbalances are supported for now
+const IMBALANCES_SIZE: usize = 2;
 
-        if amount_fr.ne(&Fr::ZERO) {
-            match sign {
-                Sign::Required => {
-                    commitment -= g_v * amount_fr;
-                }
-                Sign::Provided => {
-                    commitment += g_v * amount_fr;
-                }
+#[derive(Clone)]
+#[cfg_attr(any(feature = "derive-debug", test), derive(Debug))]
+pub struct Balance {
+    pub imbalances: [Option<Imbalance>; IMBALANCES_SIZE],
+}
+
+impl Balance {
+    pub fn new() -> Self {
+        Balance {
+            imbalances: [None, None],
+        }
+    }
+
+    pub fn add(&mut self, imbalance: Imbalance) -> Result<(), ParserError> {
+        for slot in &mut self.imbalances {
+            if slot.is_none() {
+                *slot = Some(imbalance);
+                return Ok(());
             }
         }
+        Err(ParserError::InvalidLength)
+    }
 
+    pub fn commit(&self, blinding_factor: Fr) -> Result<Commitment, ParserError> {
+        if !self.has_valid_imbalance() {
+            return Err(ParserError::InvalidLength);
+        }
+
+        let mut commitment = decaf377::Element::IDENTITY;
+
+        for imbalance in self.imbalances.iter().flatten() {
+            let g_v = imbalance.value.asset_id.value_generator();
+            let amount_fr: Fr = Into::into(imbalance.value.amount);
+
+            if amount_fr.ne(&Fr::ZERO) {
+                match imbalance.sign {
+                    Sign::Required => {
+                        commitment -= g_v * amount_fr;
+                    }
+                    Sign::Provided => {
+                        commitment += g_v * amount_fr;
+                    }
+                }
+            }
+        }
+
         let value_blinding_generator = Commitment::value_blinding_generator();
         commitment += blinding_factor * value_blinding_generator;
-
+    
         Ok(commitment.into())
     }
 
+    fn has_valid_imbalance(&self) -> bool {
+        self.imbalances.iter().any(|slot| slot.is_some())
+    }
+}
+
+impl Default for Balance {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Value {
+    pub const LEN: usize = AMOUNT_LEN_BYTES + ID_LEN_BYTES;
+
     pub fn to_bytes(&self) -> Result<[u8; Self::LEN], ParserError> {
         let mut bytes = [0; Self::LEN];
         bytes[0..AMOUNT_LEN_BYTES].copy_from_slice(&self.amount.to_le_bytes());
@@ -78,31 +123,6 @@ impl Value {
     }
 }
 
-impl Imbalance {
-    pub fn commit(&self, blinding_factor: Fr) -> Result<Commitment, ParserError> {
-        let mut commitment = decaf377::Element::IDENTITY;
-
-        // required value
-        let g_v = self.required_value.asset_id.value_generator();
-        let amount_fr: Fr = Into::into(self.required_value.amount);
-        if amount_fr.ne(&Fr::ZERO) {
-            commitment -= g_v * amount_fr;
-        }
-
-        // provided value
-        let g_v = self.provided_value.asset_id.value_generator();
-        let amount_fr: Fr = Into::into(self.provided_value.amount);
-        if amount_fr.ne(&Fr::ZERO) {
-            commitment += g_v * amount_fr;
-        }
-
-        let value_blinding_generator = Commitment::value_blinding_generator();
-        commitment += blinding_factor * value_blinding_generator;
-
-        Ok(commitment.into())
-    }
-}
-
 #[repr(C)]
 #[derive(Clone)]
 #[cfg_attr(any(feature = "derive-debug", test), derive(Debug))]