Fix CheckedSigned::on_if

Make `CheckedSigned.sgn` a cvar to properly push it to the witness vector

Author: sebastiencs

ledger/src/proofs/block.rs

@@ -1638,7 +1638,7 @@ fn block_main<'a>(
     };
 
     let supply_increase = w.exists_no_check(match txn_stmt_ledger_hashes_didn_t_change {
-        Boolean::True => CheckedSigned::const_zero(),
+        Boolean::True => CheckedSigned::zero(),
         Boolean::False => txn_snark.supply_increase.to_checked(),
     });
 

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,36 @@ 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())
-    }
-
-    pub fn const_zero() -> Self {
         Self {
             magnitude: T::zero(),
-            sgn: Sgn::Pos,
+            sgn: CircuitVar::Constant(Sgn::Pos),
             value: Cell::new(None),
         }
     }
 
     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 +119,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 +132,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 +143,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 +157,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 +192,7 @@ where
 
         let res = Self {
             magnitude: res_magnitude,
-            sgn,
+            sgn: CircuitVar::Var(sgn),
             value: Cell::new(Some(res_value)),
         };
         (res, overflow)
@@ -219,7 +218,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 +507,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),
                 }
             }

ledger/src/proofs/to_field_elements.rs

@@ -826,6 +826,24 @@ 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<F: FieldWitness, T: ToFieldElements<F>> ToFieldElements<F> for CircuitVar<T> {
+//     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 {

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

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<'a>(b: Self::Bool, param: SignedAmountBranchParam<&'a 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,

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<'a>(b: Self::Bool, param: SignedAmountBranchParam<&'a Self>, w: &mut Self::W) -> Self {
         let SignedAmountBranchParam { on_true, on_false } = param;
         match b {
-            true => on_true,
-            false => on_false,
+            true => on_true.clone(),
+            false => on_false.clone(),
         }
     }
 }

ledger/src/zkapps/snark.rs

@@ -1,4 +1,4 @@
-use std::{fmt::Write, marker::PhantomData};
+use std::{cell::Cell, fmt::Write, marker::PhantomData};
 
 use ark_ff::Zero;
 use mina_hasher::Fp;
@@ -189,7 +189,7 @@ impl SignedAmountInterface for CheckedSigned<Fp, CheckedAmount<Fp>> {
     type Amount = SnarkAmount;
 
     fn zero() -> Self {
-        CheckedSigned::zero()
+        CheckedSigned::of_unsigned(<CheckedAmount<_> as CheckedCurrency<Fp>>::zero())
     }
     fn is_neg(&self) -> Self::Bool {
         CheckedSigned::is_neg(self).var()
@@ -210,21 +210,32 @@ impl SignedAmountInterface for CheckedSigned<Fp, CheckedAmount<Fp>> {
     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<'a>(b: Self::Bool, param: SignedAmountBranchParam<&'a Self>, w: &mut Self::W) -> Self {
         let SignedAmountBranchParam { on_true, on_false } = param;
 
-        let amount = w.exists_no_check(match b.as_boolean() {
+        let amount = match b.as_boolean() {
             Boolean::True => on_true,
             Boolean::False => on_false,
-        });
-        if on_true.try_get_value().is_some() && on_false.try_get_value().is_some() {
-            w.exists_no_check(amount.force_value());
+        };
+
+        // TODO: This should be moved in a `Sgn::on_if`
+        let sgn = match (on_true.sgn, on_false.sgn) {
+            (CircuitVar::Constant(_), CircuitVar::Constant(_)) => {
+                CircuitVar::Var(*amount.sgn.value())
+            }
+            _ => CircuitVar::Var(w.exists_no_check(*amount.sgn.value())),
+        };
+        w.exists_no_check(&amount.magnitude);
+
+        let value = match (on_true.try_get_value(), on_false.try_get_value()) {
+            (Some(_), Some(_)) => Some(w.exists_no_check(amount.force_value())),
+            _ => None,
+        };
+        Self {
+            value: Cell::new(value),
+            sgn,
+            ..amount.clone()
         }
-        amount
     }
 }