Fix CheckedSigned::on_if behavior when sgn is a constant

sebastiencs · sebastiencs · commit af23b4a0715a · 2024-11-07T22:11:22.000+01:00
Make `CheckedSigned.sgn` a cvar to properly push it to the witness vector
Use `CheckedSigned::on_if` everywhere, instead of manually `match`
it (which was incorrect)
diff --git a/ledger/src/proofs/block.rs b/ledger/src/proofs/block.rs
@@ -35,6 +35,7 @@ use crate::{
         transaction_logic::protocol_state::{EpochLedger, ProtocolStateView},
     },
     staged_ledger::hash::StagedLedgerHash,
+    zkapps::intefaces::{SignedAmountBranchParam, SignedAmountInterface},
     Inputs, ToInputs,
 };
 
@@ -1637,10 +1638,14 @@ fn block_main<'a>(
         txn_statement_ledger_hashes_equal(s1, &s2, w)
     };
 
-    let supply_increase = w.exists_no_check(match txn_stmt_ledger_hashes_didn_t_change {
-        Boolean::True => CheckedSigned::const_zero(),
-        Boolean::False => txn_snark.supply_increase.to_checked(),
-    });
+    let supply_increase = CheckedSigned::on_if(
+        txn_stmt_ledger_hashes_didn_t_change.var(),
+        SignedAmountBranchParam {
+            on_true: &CheckedSigned::zero(),
+            on_false: &txn_snark.supply_increase.to_checked(),
+        },
+        w,
+    );
 
     let (updated_consensus_state, consensus_state) = consensus::next_state_checked(
         previous_state,
diff --git a/ledger/src/proofs/numbers/currency.rs b/ledger/src/proofs/numbers/currency.rs
@@ -1,4 +1,7 @@
-use crate::scan_state::currency::{self, Amount, Balance, Fee, Magnitude, MinMax, Sgn, Signed};
+use crate::{
+    proofs::field::CircuitVar,
+    scan_state::currency::{self, Amount, Balance, Fee, Magnitude, MinMax, Sgn, Signed},
+};
 use std::{cell::Cell, cmp::Ordering::Less};
 
 use crate::proofs::{
@@ -42,7 +45,7 @@ where
     T: CheckedCurrency<F>,
 {
     pub magnitude: T,
-    pub sgn: Sgn,
+    pub sgn: CircuitVar<Sgn>,
     pub value: Cell<Option<F>>,
 }
 
@@ -65,7 +68,7 @@ where
     F: FieldWitness + std::fmt::Debug,
     T: CheckedCurrency<F> + std::fmt::Debug,
 {
-    pub fn create(magnitude: T, sgn: Sgn, value: Option<F>) -> Self {
+    pub fn create(magnitude: T, sgn: CircuitVar<Sgn>, value: Option<F>) -> Self {
         Self {
             magnitude,
             sgn,
@@ -77,40 +80,46 @@ where
         let value = magnitude.to_field();
         Self {
             magnitude,
-            sgn: Sgn::Pos,
+            sgn: CircuitVar::Constant(Sgn::Pos),
             value: Cell::new(Some(value)),
         }
     }
 
     pub fn zero() -> Self {
-        Self::of_unsigned(T::zero())
+        Self {
+            magnitude: T::zero(),
+            sgn: CircuitVar::Constant(Sgn::Pos),
+            value: Cell::new(None),
+        }
     }
 
-    pub fn const_zero() -> Self {
+    // https://github.com/MinaProtocol/mina/blob/ca9c8c86aa21d3c346d28ea0be7ad4cb0c22bf7f/src/lib/transaction_snark/transaction_snark.ml#L1891-L1892
+    // https://github.com/MinaProtocol/mina/blob/ca9c8c86aa21d3c346d28ea0be7ad4cb0c22bf7f/src/lib/currency/currency.ml#L579
+    pub fn constant_zero() -> Self {
         Self {
             magnitude: T::zero(),
-            sgn: Sgn::Pos,
-            value: Cell::new(None),
+            sgn: CircuitVar::Constant(Sgn::Pos),
+            value: Cell::new(Some(T::zero().to_field())),
         }
     }
 
     pub fn negate(self) -> Self {
         Self {
             magnitude: self.magnitude,
-            sgn: self.sgn.negate(),
+            sgn: self.sgn.map(|sgn| sgn.negate()),
             value: Cell::new(self.value.get().map(|f| f.neg())),
         }
     }
 
     pub fn is_neg(&self) -> Boolean {
-        match self.sgn {
+        match self.sgn.value() {
             Sgn::Pos => Boolean::False,
             Sgn::Neg => Boolean::True,
         }
     }
 
     pub fn is_pos(&self) -> Boolean {
-        match self.sgn {
+        match self.sgn.value() {
             Sgn::Pos => Boolean::True,
             Sgn::Neg => Boolean::False,
         }
@@ -120,7 +129,7 @@ where
         match self.value.get() {
             Some(x) => x,
             None => {
-                let sgn: F = self.sgn.to_field();
+                let sgn: F = self.sgn.value().to_field();
                 let magnitude: F = self.magnitude.to_field();
                 let value = w.exists_no_check(magnitude * sgn);
                 self.value.replace(Some(value));
@@ -133,7 +142,7 @@ where
         match self.value.get() {
             Some(x) => x,
             None => {
-                let sgn: F = self.sgn.to_field();
+                let sgn: F = self.sgn.value().to_field();
                 let magnitude: F = self.magnitude.to_field();
                 magnitude * sgn
             }
@@ -144,7 +153,7 @@ where
         match self.value.get() {
             Some(_) => {}
             None => {
-                let sgn: F = self.sgn.to_field();
+                let sgn: F = self.sgn.value().to_field();
                 let magnitude: F = self.magnitude.to_field();
                 self.value.replace(Some(magnitude * sgn));
             }
@@ -158,7 +167,7 @@ where
     fn unchecked(&self) -> currency::Signed<T::Inner> {
         currency::Signed {
             magnitude: self.magnitude.to_inner(),
-            sgn: self.sgn,
+            sgn: *self.sgn.value(),
         }
     }
 
@@ -193,7 +202,7 @@ where
 
         let res = Self {
             magnitude: res_magnitude,
-            sgn,
+            sgn: CircuitVar::Var(sgn),
             value: Cell::new(Some(res_value)),
         };
         (res, overflow)
@@ -219,7 +228,11 @@ where
 
         range_check::<F, CURRENCY_NBITS>(magnitude, w);
 
-        Self::create(T::from_field(magnitude), sgn, Some(res_value))
+        Self::create(
+            T::from_field(magnitude),
+            CircuitVar::Var(sgn),
+            Some(res_value),
+        )
     }
 
     pub fn equal(&self, other: &Self, w: &mut Witness<F>) -> Boolean {
@@ -504,7 +517,7 @@ macro_rules! impl_currency {
             pub fn to_checked<F: FieldWitness>(&self) -> CheckedSigned<F, $name<F>> {
                 CheckedSigned {
                     magnitude: self.magnitude.to_checked(),
-                    sgn: self.sgn,
+                    sgn: CircuitVar::Var(self.sgn),
                     value: Cell::new(None),
                 }
             }
diff --git a/ledger/src/proofs/to_field_elements.rs b/ledger/src/proofs/to_field_elements.rs
@@ -32,7 +32,6 @@ use crate::{
 
 use super::{
     field::{Boolean, CircuitVar, FieldWitness, GroupAffine},
-    numbers::currency::{CheckedCurrency, CheckedSigned},
     step::PerProofWitness,
     transaction::{
         field_to_bits, InnerCurve, PlonkVerificationKeyEvals, StepMainProofState, StepMainStatement,
@@ -799,13 +798,6 @@ impl<F: FieldWitness, T: ToFieldElements<F>> ToFieldElements<F> for &T {
     }
 }
 
-impl<F: FieldWitness, T: CheckedCurrency<F>> ToFieldElements<F> for CheckedSigned<F, T> {
-    fn to_field_elements(&self, fields: &mut Vec<F>) {
-        self.sgn.to_field_elements(fields);
-        self.magnitude.to_field_elements(fields);
-    }
-}
-
 impl<F: FieldWitness> ToFieldElements<F> for InnerCurve<F> {
     fn to_field_elements(&self, fields: &mut Vec<F>) {
         let GroupAffine::<F> { x, y, .. } = self.to_affine();
@@ -826,6 +818,15 @@ impl<F: FieldWitness> ToFieldElements<F> for CircuitVar<Boolean> {
     }
 }
 
+impl<F: FieldWitness> ToFieldElements<F> for CircuitVar<Sgn> {
+    fn to_field_elements(&self, fields: &mut Vec<F>) {
+        match self {
+            CircuitVar::Constant(_) => (),
+            CircuitVar::Var(var) => var.to_field_elements(fields),
+        }
+    }
+}
+
 impl ToFieldElements<Fp> for StepMainStatement {
     fn to_field_elements(&self, fields: &mut Vec<Fp>) {
         let Self {
diff --git a/ledger/src/proofs/transaction.rs b/ledger/src/proofs/transaction.rs
@@ -2285,6 +2285,7 @@ pub mod transaction_snark {
             transaction_logic::{checked_cons_signed_command_payload, Coinbase},
         },
         sparse_ledger::SparseLedger,
+        zkapps::intefaces::{SignedAmountBranchParam, SignedAmountInterface},
         AccountId, Inputs, PermissionTo, PermsConst, Timing, TimingAsRecordChecked, ToInputs,
     };
     use ark_ff::Zero;
@@ -3080,7 +3081,11 @@ pub mod transaction_snark {
                             Boolean::True => Sgn::Neg,
                             Boolean::False => Sgn::Pos,
                         };
-                        CheckedSigned::create(CheckedAmount::of_fee(&fee), sgn, None)
+                        CheckedSigned::create(
+                            CheckedAmount::of_fee(&fee),
+                            CircuitVar::Constant(sgn),
+                            None,
+                        )
                     };
 
                     let account_creation_fee = {
@@ -3090,7 +3095,7 @@ pub mod transaction_snark {
                         } else {
                             CheckedAmount::zero()
                         };
-                        CheckedSigned::create(magnitude, Sgn::Neg, None)
+                        CheckedSigned::create(magnitude, CircuitVar::Constant(Sgn::Neg), None)
                     };
 
                     new_account_fees = account_creation_fee.clone();
@@ -3113,7 +3118,7 @@ pub mod transaction_snark {
 
                 let txn_global_slot = current_global_slot;
                 let timing = {
-                    let txn_amount = w.exists_no_check(match amount.sgn {
+                    let txn_amount = w.exists_no_check(match amount.sgn.value() {
                         Sgn::Neg => amount.magnitude,
                         Sgn::Pos => CheckedAmount::zero(),
                     });
@@ -3569,45 +3574,59 @@ pub mod transaction_snark {
                 let (fee_transfer_excess, fee_transfer_excess_overflowed) = {
                     let (magnitude, overflow) =
                         payload.body.amount.to_checked().add_flagged(&amount_fee, w);
-                    (CheckedSigned::create(magnitude, Sgn::Neg, None), overflow)
+                    (
+                        CheckedSigned::create(magnitude, CircuitVar::Constant(Sgn::Neg), None),
+                        overflow,
+                    )
                 };
 
                 Boolean::assert_any(
                     &[is_fee_transfer.neg(), fee_transfer_excess_overflowed.neg()],
                     w,
                 );
 
-                let value = match is_fee_transfer {
-                    Boolean::True => fee_transfer_excess,
-                    Boolean::False => user_command_excess,
-                };
-                w.exists_no_check(value.magnitude);
-                value
+                CheckedSigned::on_if(
+                    is_fee_transfer.var(),
+                    SignedAmountBranchParam {
+                        on_true: &fee_transfer_excess,
+                        on_false: &user_command_excess,
+                    },
+                    w,
+                )
             };
 
-            w.exists_no_check(match is_coinbase {
-                Boolean::True => then_value,
-                Boolean::False => else_value,
-            })
+            CheckedSigned::on_if(
+                is_coinbase.var(),
+                SignedAmountBranchParam {
+                    on_true: &then_value,
+                    on_false: &else_value,
+                },
+                w,
+            )
         };
 
         let supply_increase = {
-            let expected_supply_increase = match is_coinbase {
-                Boolean::True => CheckedSigned::of_unsigned(payload.body.amount.to_checked()),
-                Boolean::False => CheckedSigned::of_unsigned(CheckedAmount::zero()),
-            };
-            w.exists_no_check(expected_supply_increase.magnitude);
-            w.exists_no_check(expected_supply_increase.magnitude);
+            let expected_supply_increase = CheckedSigned::on_if(
+                is_coinbase.var(),
+                SignedAmountBranchParam {
+                    on_true: &CheckedSigned::of_unsigned(payload.body.amount.to_checked()),
+                    on_false: &CheckedSigned::of_unsigned(CheckedAmount::zero()),
+                },
+                w,
+            );
 
             let (amt0, _overflow0) = expected_supply_increase
                 .add_flagged(&CheckedSigned::of_unsigned(burned_tokens).negate(), w);
 
-            let new_account_fees_total = w.exists_no_check(match user_command_fails {
-                Boolean::True => zero_fee,
-                Boolean::False => new_account_fees,
-            });
+            let new_account_fees_total = CheckedSigned::on_if(
+                user_command_fails.var(),
+                SignedAmountBranchParam {
+                    on_true: &zero_fee,
+                    on_false: &new_account_fees,
+                },
+                w,
+            );
 
-            w.exists(new_account_fees_total.force_value()); // Made in the `add_flagged` call
             let (amt, _overflow) = amt0.add_flagged(&new_account_fees_total, w);
 
             amt
diff --git a/ledger/src/proofs/zkapp.rs b/ledger/src/proofs/zkapp.rs
@@ -1075,8 +1075,8 @@ fn zkapp_main(
                 ledger: witness.global_second_pass_ledger.clone(),
                 hash: statement.source.second_pass_ledger,
             },
-            fee_excess: CheckedSigned::zero(),
-            supply_increase: CheckedSigned::zero(),
+            fee_excess: CheckedSigned::constant_zero(),
+            supply_increase: CheckedSigned::constant_zero(),
             protocol_state: state_body.view(),
             block_global_slot,
         };
diff --git a/ledger/src/zkapps/intefaces.rs b/ledger/src/zkapps/intefaces.rs
@@ -142,11 +142,7 @@ where
     fn negate(&self) -> Self;
     fn add_flagged(&self, other: &Self, w: &mut Self::W) -> (Self, Self::Bool);
     fn of_unsigned(unsigned: Self::Amount) -> Self;
-    fn on_if<'a>(
-        b: Self::Bool,
-        param: SignedAmountBranchParam<&'a Self>,
-        w: &mut Self::W,
-    ) -> &'a Self;
+    fn on_if(b: Self::Bool, param: SignedAmountBranchParam<&Self>, w: &mut Self::W) -> Self;
 }
 
 pub trait BalanceInterface
@@ -226,6 +222,7 @@ pub struct StackFrameMakeParams<'a, Calls> {
     pub calls: &'a Calls,
 }
 
+#[derive(Debug)]
 pub struct SignedAmountBranchParam<T> {
     pub on_true: T,
     pub on_false: T,
@@ -583,8 +580,7 @@ where
     >;
     type SignedAmount: SignedAmountInterface<W = Self::WitnessGenerator, Bool = Self::Bool, Amount = Self::Amount>
         + std::fmt::Debug
-        + Clone
-        + ToFieldElements<Fp>;
+        + Clone;
     type Amount: AmountInterface<W = Self::WitnessGenerator, Bool = Self::Bool> + Clone;
     type Balance: BalanceInterface<
         W = Self::WitnessGenerator,
diff --git a/ledger/src/zkapps/non_snark.rs b/ledger/src/zkapps/non_snark.rs
@@ -401,15 +401,11 @@ impl SignedAmountInterface for Signed<Amount> {
     fn of_unsigned(unsigned: Self::Amount) -> Self {
         Self::of_unsigned(unsigned)
     }
-    fn on_if<'a>(
-        b: Self::Bool,
-        param: SignedAmountBranchParam<&'a Self>,
-        w: &mut Self::W,
-    ) -> &'a Self {
+    fn on_if(b: Self::Bool, param: SignedAmountBranchParam<&Self>, w: &mut Self::W) -> Self {
         let SignedAmountBranchParam { on_true, on_false } = param;
         match b {
-            true => on_true,
-            false => on_false,
+            true => *on_true,
+            false => *on_false,
         }
     }
 }
diff --git a/ledger/src/zkapps/snark.rs b/ledger/src/zkapps/snark.rs
diff --git a/ledger/src/zkapps/zkapp_logic.rs b/ledger/src/zkapps/zkapp_logic.rs

Original file line number	Diff line number	Diff line change
`@@ -401,15 +401,11 @@ impl SignedAmountInterface for Signed<Amount> {`
`401`	`401`	`fn of_unsigned(unsigned: Self::Amount) -> Self {`
`402`	`402`	`Self::of_unsigned(unsigned)`
`403`	`403`	`}`
`404`		`- fn on_if<'a>(`
`405`		`- b: Self::Bool,`
`406`		`- param: SignedAmountBranchParam<&'a Self>,`
`407`		`- w: &mut Self::W,`
`408`		`- ) -> &'a Self {`
	`404`	`+ fn on_if(b: Self::Bool, param: SignedAmountBranchParam<&Self>, w: &mut Self::W) -> Self {`
`409`	`405`	`let SignedAmountBranchParam { on_true, on_false } = param;`
`410`	`406`	`match b {`
`411`		`- true => on_true,`
`412`		`- false => on_false,`
	`407`	`+ true => *on_true,`
	`408`	`+ false => *on_false,`
`413`	`409`	`}`
`414`	`410`	`}`
`415`	`411`	`}`