Ensure Poseidon's alpha and F::MODULUS - 1 are coprime

crates/primitives/src/transcripts/poseidon/mod.rs

@@ -1,34 +1,139 @@
 //! Poseidon-based transcript configurations and implementations.
 
 use ark_crypto_primitives::sponge::poseidon::{PoseidonConfig, find_poseidon_ark_and_mds};
-use ark_ff::PrimeField;
+use ark_ff::{One, PrimeField};
+use num_bigint::BigUint;
+use num_integer::Integer;
 
 pub mod sponge;
 
-/// [`poseidon_custom_config`] produces a Poseidon configuration with custom
-/// parameters.
-pub fn poseidon_custom_config<F: PrimeField>(
-    full_rounds: usize,
-    partial_rounds: usize,
+fn log2_order<F: PrimeField>() -> f64 {
+    let x = F::MODULUS.into();
+    let bits = x.bits(); // bit length
+    if bits <= 53 {
+        // Fits in f64 mantissa exactly
+        let val: u64 = x.try_into().unwrap();
+        return (val as f64).log2();
+    }
+    // Shift right so only top ~53 bits remain
+    let shift = bits - 53;
+    let top = x >> shift;
+    let top_u64: u64 = top.try_into().unwrap();
+    (top_u64 as f64).log2() + shift as f64
+}
+
+fn sat_inequiv_alpha<F: PrimeField>(t: usize, r_f: u64, r_p: u64, alpha: u64, m: usize) -> bool {
+    let log2_p = log2_order::<F>();
+    let n = log2_p.ceil() as usize;
+    let m_f = m as f64;
+    let n_f = n as f64;
+    let t_f = t as f64;
+    let r_p_f = r_p as f64;
+    let r_f_f = r_f as f64;
+    let alpha_f = alpha as f64;
+    let log2_alpha = 2.0f64.ln() / alpha_f.ln();
+
+    let r_f_1: f64 = if m_f <= (log2_p - (alpha_f - 1.0) / 2.0).floor() * (t_f + 1.0) {
+        6.0
+    } else {
+        10.0
+    };
+
+    let r_f_2 = 1.0 + log2_alpha * m_f.min(n_f) + (t_f.ln() / alpha_f.ln()).ceil() - r_p_f;
+
+    let r_f_3 = 1.0 + log2_alpha * (m_f / 3.0).min(log2_p / 2.0) - r_p_f;
+
+    let r_f_4 = t_f - 1.0 + (log2_alpha * m_f / (t_f + 1.0)).min(log2_alpha * log2_p / 2.0) - r_p_f;
+
+    let r_f_max = r_f_1
+        .ceil()
+        .max(r_f_2.ceil())
+        .max(r_f_3.ceil())
+        .max(r_f_4.ceil());
+
+    r_f_f >= r_f_max
+}
+
+fn get_sbox_cost(r_f: u64, r_p: u64, _n: usize, t: usize) -> usize {
+    t * r_f as usize + r_p as usize
+}
+
+fn find_fd_round_numbers<F: PrimeField>(
+    t: usize,
+    alpha: u64,
+    m: usize,
+    cost_function: fn(u64, u64, usize, usize) -> usize,
+    security_margin: bool,
+) -> (u64, u64) {
+    let n = log2_order::<F>().ceil() as usize;
+    let n_total = n * t;
+
+    let mut r_p: u64 = 0;
+    let mut r_f: u64 = 0;
+    let mut min_cost = usize::MAX;
+    let mut max_cost_rf: u64 = 0;
+
+    for r_p_t in 1u64..500 {
+        for r_f_t in (4u64..100).step_by(2) {
+            if !sat_inequiv_alpha::<F>(t, r_f_t, r_p_t, alpha, m) {
+                continue;
+            }
+
+            let (r_f_eff, r_p_eff) = if security_margin {
+                (r_f_t + 2, (r_p_t as f64 * 1.075).ceil() as u64)
+            } else {
+                (r_f_t, r_p_t)
+            };
+
+            let cost = cost_function(r_f_eff, r_p_eff, n_total, t);
+            if cost < min_cost || (cost == min_cost && r_f_eff < max_cost_rf) {
+                r_p = r_p_eff;
+                r_f = r_f_eff;
+                min_cost = cost;
+                max_cost_rf = r_f;
+            }
+        }
+    }
+
+    (r_f, r_p)
+}
+
+/// [`poseidon_paper_config`] produces a Poseidon configuration which agrees
+/// with the paper's reference implementation.
+///
+/// TODO(@winderica): alpha = -1 is not supported yet.
+pub fn poseidon_paper_config<F: PrimeField, const SECURITY_BITS: usize>(
     alpha: u64,
     rate: usize,
-    capacity: usize,
 ) -> PoseidonConfig<F> {
+    assert_eq!(
+        BigUint::from(alpha).gcd(&(-F::one()).into()),
+        BigUint::one()
+    );
+    let (full_rounds, partial_rounds) =
+        find_fd_round_numbers::<F>(rate, alpha, SECURITY_BITS, get_sbox_cost, true);
     let (ark, mds) = find_poseidon_ark_and_mds::<F>(
         F::MODULUS_BIT_SIZE as u64,
         rate,
-        full_rounds as u64,
-        partial_rounds as u64,
+        full_rounds,
+        partial_rounds,
         0,
     );
 
-    PoseidonConfig::new(full_rounds, partial_rounds, alpha, mds, ark, rate, capacity)
+    PoseidonConfig::new(
+        full_rounds as usize,
+        partial_rounds as usize,
+        alpha,
+        mds,
+        ark,
+        rate,
+        1,
+    )
 }
 
-/// [`poseidon_canonical_config`] produces a Poseidon configuration with default
-/// parameters, which agrees with Circom's Poseidon(4) when `F` is the scalar
-/// field of BN254.
-pub fn poseidon_canonical_config<F: PrimeField>() -> PoseidonConfig<F> {
+/// [`poseidon_circom_config`] produces a Poseidon configuration which agrees
+/// with Circom's Poseidon(4) when `F` is the scalar field of BN254.
+pub fn poseidon_circom_config<F: PrimeField>() -> PoseidonConfig<F> {
     // 120 bit security target as in
     // https://eprint.iacr.org/2019/458.pdf
     // t = rate + 1
@@ -38,5 +143,13 @@ pub fn poseidon_canonical_config<F: PrimeField>() -> PoseidonConfig<F> {
     let alpha = 5;
     let rate = 4;
 
-    poseidon_custom_config(full_rounds, partial_rounds, alpha, rate, 1)
+    let (ark, mds) = find_poseidon_ark_and_mds::<F>(
+        F::MODULUS_BIT_SIZE as u64,
+        rate,
+        full_rounds as u64,
+        partial_rounds as u64,
+        0,
+    );
+
+    PoseidonConfig::new(full_rounds, partial_rounds, alpha, mds, ark, rate, 1)
 }

crates/primitives/src/transcripts/poseidon/sponge.rs

@@ -95,13 +95,13 @@ mod tests {
 
     use crate::{
         algebra::group::emulated::EmulatedAffineVar,
-        transcripts::{Transcript, TranscriptGadget, poseidon::poseidon_canonical_config},
+        transcripts::{Transcript, TranscriptGadget, poseidon::poseidon_circom_config},
     };
 
     // Test with value taken from https://github.com/iden3/circomlibjs/blob/43cc582b100fc3459cf78d903a6f538e5d7f38ee/test/poseidon.js#L32
     #[test]
     fn check_against_circom_poseidon() -> Result<(), Box<dyn Error>> {
-        let config = poseidon_canonical_config::<Fr>();
+        let config = poseidon_circom_config::<Fr>();
         let mut poseidon_sponge = PoseidonSponge::new(&config);
         let v = vec![1, 2, 3, 4]
             .into_iter()
@@ -122,7 +122,7 @@ mod tests {
     #[test]
     fn test_challenge_field_element() -> Result<(), Box<dyn Error>> {
         // Create a transcript outside of the circuit
-        let config = poseidon_canonical_config::<Fr>();
+        let config = poseidon_circom_config::<Fr>();
         let mut tr = PoseidonSponge::<Fr>::new(&config);
         tr.add(&Fr::from(42_u32));
         let c = tr.challenge_field_element();
@@ -145,7 +145,7 @@ mod tests {
         let nbits = 128;
 
         // Create a transcript outside of the circuit
-        let config = poseidon_canonical_config::<Fq>();
+        let config = poseidon_circom_config::<Fq>();
         let mut tr = PoseidonSponge::<Fq>::new(&config);
         tr.add(&Fq::from(42_u32));
         let c = tr.challenge_bits(nbits);
@@ -166,7 +166,7 @@ mod tests {
     #[test]
     fn test_absorb_canonical_point() -> Result<(), Box<dyn Error>> {
         // Create a transcript outside of the circuit
-        let config = poseidon_canonical_config::<Fq>();
+        let config = poseidon_circom_config::<Fq>();
         let mut tr = PoseidonSponge::<Fq>::new(&config);
         let rng = &mut thread_rng();
 
@@ -190,7 +190,7 @@ mod tests {
     #[test]
     fn test_absorb_emulated_point() -> Result<(), Box<dyn Error>> {
         // Create a transcript outside of the circuit
-        let config = poseidon_canonical_config::<Fr>();
+        let config = poseidon_circom_config::<Fr>();
         let mut tr = PoseidonSponge::<Fr>::new(&config);
         let rng = &mut thread_rng();