From d240cf9653e7a80032f883022d02e51702e8deaf Mon Sep 17 00:00:00 2001
From: sifnoc <sifnoc09@gmail.com>
Date: Tue, 4 Jun 2024 23:36:24 +0900
Subject: [PATCH] Fix selector and concatenated balance on nonzero-constrain

---
 backend/Cargo.lock                            |  2 +-
 prover/src/circuits/config/circuit_config.rs  |  5 +--
 .../circuits/config/no_range_check_config.rs  |  3 +-
 .../src/circuits/config/range_check_config.rs |  8 ++---
 prover/src/circuits/summa_circuit.rs          | 36 +++++++++++--------
 prover/src/circuits/tests.rs                  |  7 ++--
 prover/src/utils/dummy_entries.rs             |  4 ++-
 7 files changed, 37 insertions(+), 28 deletions(-)

diff --git a/backend/Cargo.lock b/backend/Cargo.lock
index 95318bca..24e50410 100644
--- a/backend/Cargo.lock
+++ b/backend/Cargo.lock
@@ -2427,7 +2427,7 @@ checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964"
 [[package]]
 name = "plonkish_backend"
 version = "0.1.0"
-source = "git+https://github.com/summa-dev/plonkish?branch=nonzero-constraints#e37ba53dcc8f8bd6e7add4e479d0e3295ee80661"
+source = "git+https://github.com/summa-dev/plonkish?branch=nonzero-constraints#c9dc12571d4a9aa06a419598ff2422b42770ae91"
 dependencies = [
  "bincode",
  "bitvec 1.0.1",
diff --git a/prover/src/circuits/config/circuit_config.rs b/prover/src/circuits/config/circuit_config.rs
index 0d7f5b76..8c2786a6 100644
--- a/prover/src/circuits/config/circuit_config.rs
+++ b/prover/src/circuits/config/circuit_config.rs
@@ -1,6 +1,6 @@
 use halo2_proofs::{
     circuit::{Layouter, Value},
-    plonk::{Advice, Column, ConstraintSystem, Error, Instance},
+    plonk::{Advice, Column, ConstraintSystem, Error, Instance, Selector},
 };
 
 use crate::{entry::Entry, utils::big_uint_to_fp};
@@ -19,6 +19,7 @@ pub trait CircuitConfig<const N_CURRENCIES: usize, const N_USERS: usize>: Clone
         meta: &mut ConstraintSystem<Fp>,
         username: Column<Advice>,
         concatenated_balance: Column<Advice>,
+        selector: Selector,
         balances: [Column<Advice>; N_CURRENCIES],
         instance: Column<Instance>,
     ) -> Self;
@@ -104,7 +105,7 @@ pub trait CircuitConfig<const N_CURRENCIES: usize, const N_USERS: usize>: Clone
                     || format!("concateneated total({} currencies)", N_CURRENCIES),
                     self.get_concatenated_balance(),
                     0,
-                    || Value::known(concatenated_grand_total.nag()),
+                    || Value::known(concatenated_grand_total.neg()),
                 )?;
 
                 Ok(balance_total)
diff --git a/prover/src/circuits/config/no_range_check_config.rs b/prover/src/circuits/config/no_range_check_config.rs
index 29a144fe..e8239aa7 100644
--- a/prover/src/circuits/config/no_range_check_config.rs
+++ b/prover/src/circuits/config/no_range_check_config.rs
@@ -1,6 +1,6 @@
 use halo2_proofs::{
     circuit::Layouter,
-    plonk::{Advice, Column, ConstraintSystem, Error, Instance},
+    plonk::{Advice, Column, ConstraintSystem, Error, Instance, Selector},
 };
 
 use crate::chips::range::range_check::RangeCheckU64Chip;
@@ -37,6 +37,7 @@ impl<const N_CURRENCIES: usize, const N_USERS: usize> CircuitConfig<N_CURRENCIES
         _: &mut ConstraintSystem<Fp>,
         username: Column<Advice>,
         concatenated_balance: Column<Advice>,
+        _selector: Selector,
         balances: [Column<Advice>; N_CURRENCIES],
         instance: Column<Instance>,
     ) -> NoRangeCheckConfig<N_CURRENCIES, N_USERS> {
diff --git a/prover/src/circuits/config/range_check_config.rs b/prover/src/circuits/config/range_check_config.rs
index 6c918dae..08ad6f94 100644
--- a/prover/src/circuits/config/range_check_config.rs
+++ b/prover/src/circuits/config/range_check_config.rs
@@ -1,6 +1,6 @@
 use halo2_proofs::{
     circuit::{Layouter, Value},
-    plonk::{Advice, Column, ConstraintSystem, Error, Fixed, Instance},
+    plonk::{Advice, Column, ConstraintSystem, Error, Fixed, Instance, Selector},
 };
 
 use crate::chips::range::range_check::{RangeCheckChipConfig, RangeCheckU64Chip};
@@ -40,6 +40,7 @@ impl<const N_CURRENCIES: usize, const N_USERS: usize> CircuitConfig<N_CURRENCIES
         meta: &mut ConstraintSystem<Fp>,
         username: Column<Advice>,
         concatenated_balance: Column<Advice>,
+        selector: Selector,
         balances: [Column<Advice>; N_CURRENCIES],
         instance: Column<Instance>,
     ) -> Self {
@@ -49,8 +50,6 @@ impl<const N_CURRENCIES: usize, const N_USERS: usize> CircuitConfig<N_CURRENCIES
 
         meta.annotate_lookup_any_column(range_u16, || "LOOKUP_MAXBITS_RANGE");
 
-        let range_check_selector = meta.complex_selector();
-
         // Create an empty array of range check configs
         let mut range_check_configs = Vec::with_capacity(N_CURRENCIES);
 
@@ -63,8 +62,7 @@ impl<const N_CURRENCIES: usize, const N_USERS: usize> CircuitConfig<N_CURRENCIES
                 meta.enable_equality(*column);
             }
 
-            let range_check_config =
-                RangeCheckU64Chip::configure(meta, z, zs, range_u16, range_check_selector);
+            let range_check_config = RangeCheckU64Chip::configure(meta, z, zs, range_u16, selector);
 
             range_check_configs.push(range_check_config);
         }
diff --git a/prover/src/circuits/summa_circuit.rs b/prover/src/circuits/summa_circuit.rs
index 9bc04003..ccdc4de2 100644
--- a/prover/src/circuits/summa_circuit.rs
+++ b/prover/src/circuits/summa_circuit.rs
@@ -5,7 +5,6 @@ use halo2_proofs::{
     plonk::{Circuit, ConstraintSystem, Error, Expression},
     poly::Rotation,
 };
-use num_bigint::BigUint;
 
 use crate::{entry::Entry, utils::big_uint_to_fp};
 
@@ -97,36 +96,43 @@ impl<
         }
 
         meta.create_gate("Concatenated balance validation check gate", |meta| {
-            let q_enable = meta.query_selector(q_enable);
+            let s = meta.query_selector(q_enable);
 
             let concatenated_balance = meta.query_advice(concatenated_balance, Rotation::cur());
-            let mut expr = Expression::Constant(Fp::zero());
+
+            // Right-most balance column is for the least significant balance in concatenated balance.
+            let mut balances_expr = meta.query_advice(balances[N_CURRENCIES - 1], Rotation::cur());
+
             // Base shift value for 84 bits.
             let base_shift = Fp::from(1 << 63).mul(&Fp::from(1 << 21));
 
-            let mut current_shift = Expression::Constant(Fp::one());
+            let mut current_shift = Expression::Constant(base_shift);
 
-            // Right-most balance column is for the least significant balance in concatenated balance.
-            for (i, balance_col) in balances.iter().rev().enumerate() {
-                let balance = meta.query_advice(*balance_col, Rotation::cur());
-                let shifted_balance = balance * current_shift.clone(); // Apply the current shift
-                expr = expr + shifted_balance; // Add to the expression
-
-                // Update the shift for the next iteration
-                if i < balances.len() - 1 {
-                    // Prevent updating shift unnecessarily on the last element
+            for i in (0..N_CURRENCIES - 1).rev() {
+                let balance = meta.query_advice(balances[i], Rotation::cur());
+                let shifted_balance = balance * current_shift.clone();
+                balances_expr = balances_expr + shifted_balance;
+
+                if i != 0 {
                     current_shift = current_shift * Expression::Constant(base_shift);
                 }
             }
 
             // Ensure that the whole expression equals to the concatenated_balance
-            vec![q_enable * (concatenated_balance - expr)]
+            vec![s * (concatenated_balance - balances_expr)]
         });
 
         let instance = meta.instance_column();
         meta.enable_equality(instance);
 
-        CONFIG::configure(meta, username, concatenated_balance, balances, instance)
+        CONFIG::configure(
+            meta,
+            username,
+            concatenated_balance,
+            q_enable,
+            balances,
+            instance,
+        )
     }
 
     fn synthesize(&self, config: Self::Config, layouter: impl Layouter<Fp>) -> Result<(), Error> {
diff --git a/prover/src/circuits/tests.rs b/prover/src/circuits/tests.rs
index 19a50f0b..d40e65ac 100644
--- a/prover/src/circuits/tests.rs
+++ b/prover/src/circuits/tests.rs
@@ -29,7 +29,7 @@ use crate::{
     },
 };
 const K: u32 = 17;
-const N_CURRENCIES: usize = 2;
+const N_CURRENCIES: usize = 3;
 // One row is reserved for the grand total.
 const N_USERS: usize = (1 << K) - 2;
 
@@ -85,8 +85,6 @@ fn test_summa_hyperplonk_e2e() {
         (witness_polys, proof_transcript)
     };
 
-    let num_points = 2;
-
     let proof = proof_transcript.into_proof();
 
     let mut transcript;
@@ -178,6 +176,9 @@ fn test_summa_hyperplonk_e2e() {
 
     let mut transcript = Keccak256Transcript::from_proof((), proof.as_slice());
 
+    // Username and Concatenated balance
+    let num_points = 2;
+
     let user_entry_commitments = MultilinearKzg::<Bn256>::read_commitments(
         &verifier_parameters.pcs,
         num_points,
diff --git a/prover/src/utils/dummy_entries.rs b/prover/src/utils/dummy_entries.rs
index bdd507d9..85e5b099 100644
--- a/prover/src/utils/dummy_entries.rs
+++ b/prover/src/utils/dummy_entries.rs
@@ -21,7 +21,6 @@ pub fn generate_dummy_entries<const N_USERS: usize, const N_CURRENCIES: usize>(
         let username: String = (0..10).map(|_| rng.sample(Alphanumeric) as char).collect();
 
         let balances: [BigUint; N_CURRENCIES] =
-            // std::array::from_fn(|_| BigUint::from(rng.gen_range(16777215..16777216) as u32));
             std::array::from_fn(|_| BigUint::from(rng.gen_range(10000..90000) as u32));
 
         *entry = Entry::new(username, balances).expect("Failed to create entry");
@@ -32,6 +31,9 @@ pub fn generate_dummy_entries<const N_USERS: usize, const N_CURRENCIES: usize>(
 
 #[cfg(test)]
 mod tests {
+    use crate::utils::big_uint_to_fp;
+    use halo2_proofs::halo2curves::bn256::Fr as Fp;
+
     use super::*;
 
     #[test]