Add Sarkar2020 square root - useful for pasta

CHANGELOG.md

@@ -15,6 +15,7 @@
 - [\#1044](https://github.com/arkworks-rs/algebra/pull/1044), [\#1084](https://github.com/arkworks-rs/algebra/pull/1084), [\#1088](https://github.com/arkworks-rs/algebra/pull/1088) Add implementation for small field with native integer types
 - [\#1061](https://github.com/arkworks-rs/algebra/pull/1061) (`ark-poly`) Reduce allocations in `DenseMultilinearExtension::{concat, fix_variables, evaluate}`.
 - [\#1112](https://github.com/arkworks-rs/algebra/pull/1112) (`ark-ec`) Fix rayon::ThreadPoolBuilder panicking in wasm32 when parallel feature is enabled
+- [\#1119](https://github.com/arkworks-rs/algebra/pull/1119) (`ark-ff`, `ark-pallas`) Add Sarkar2020 square root - useful for pasta
 
 ### Breaking changes
 

curves/pallas/src/fields/fq.rs

@@ -3,5 +3,6 @@ use ark_ff::fields::{Fp256, MontBackend, MontConfig};
 #[derive(MontConfig)]
 #[modulus = "28948022309329048855892746252171976963363056481941560715954676764349967630337"]
 #[generator = "5"]
+#[sqrt_precomp = "crate::fields::fq_sqrt_table::SQRT_PRECOMP"]
 pub struct FqConfig;
 pub type Fq = Fp256<MontBackend<FqConfig, 4>>;

curves/pallas/src/fields/fr.rs

@@ -3,5 +3,6 @@ use ark_ff::fields::{Fp256, MontBackend, MontConfig};
 #[derive(MontConfig)]
 #[modulus = "28948022309329048855892746252171976963363056481941647379679742748393362948097"]
 #[generator = "5"]
+#[sqrt_precomp = "crate::fields::fr_sqrt_table::SQRT_PRECOMP"]
 pub struct FrConfig;
 pub type Fr = Fp256<MontBackend<FrConfig, 4>>;

curves/pallas/src/fields/mod.rs

@@ -2,11 +2,15 @@
 pub mod fq;
 #[cfg(feature = "base_field")]
 pub use self::fq::*;
+#[cfg(feature = "base_field")]
+mod fq_sqrt_table;
 
 #[cfg(feature = "scalar_field")]
 pub mod fr;
 #[cfg(feature = "scalar_field")]
 pub use self::fr::*;
+#[cfg(feature = "scalar_field")]
+mod fr_sqrt_table;
 
 #[cfg(all(feature = "curve", test))]
 mod tests;

curves/pallas/tests/gen_sqrt_tables.rs

@@ -0,0 +1,191 @@
+//! One-off generator for the Sarkar (2020) square-root tables used by
+//! `ark_ff::SqrtPrecomputation::Sarkar2020` for the Pasta fields.
+//!
+//! The table-building logic here is adapted from `SqrtTables::new` in
+//! `zcash/pasta_curves`:
+//! <https://github.com/zcash/pasta_curves/blob/main/src/arithmetic/fields.rs>
+//! It produces the same tables that pasta builds lazily at runtime, but emits
+//! them as `const` source so the runtime has no allocation/initialization cost
+//!
+//! Run with:
+//! ```text
+//! cargo test -p ark-pallas --test gen_sqrt_tables -- --ignored --nocapture
+//! ```
+//!
+//! It writes two generated modules:
+//! - `curves/pallas/src/fields/fq_sqrt_table.rs`  (p field)
+//! - `curves/pallas/src/fields/fr_sqrt_table.rs`  (q field)
+//!
+//! `ark_pallas::Fq` is the Pallas base field (p); `ark_pallas::Fr` is the
+//! Pallas scalar field (q). Vesta reuses Pallas's `FqConfig`/`FrConfig`
+//! (`ark_vesta::Fq == ark_pallas::Fr`, `ark_vesta::Fr == ark_pallas::Fq`).
+
+use ark_ff::{FftField, PrimeField};
+use std::fmt::Write as _;
+use std::fs;
+use std::path::PathBuf;
+
+/// Perfect-hash parameters taken verbatim from `zcash/pasta_curves` — the
+/// `lazy_static! { FP_TABLES / FQ_TABLES }` definitions in `src/fields/fp.rs`
+/// and `src/fields/fq.rs` — where they are originally produced by the
+/// `squareroottab.sage` in the `zcash/pasta` <https://github.com/zcash/pasta/blob/master/squareroottab.sage> repo:
+/// - <https://github.com/zcash/pasta_curves/blob/main/src/fields/fp.rs>
+/// - <https://github.com/zcash/pasta_curves/blob/main/src/fields/fq.rs>
+///
+/// They are valid because arkworks' `TWO_ADIC_ROOT_OF_UNITY` equals pasta's
+/// `ROOT_OF_UNITY` (both = `GENERATOR^T` with `GENERATOR = 5`). We do NOT
+/// re-run the sage search; instead `build` below asserts the resulting hash is
+/// collision-free over all 256 subgroup elements, so a wrong value would panic
+/// rather than emit a broken table.
+const P_FIELD_HASH: (u32, u32) = (0x11BE, 1098); // pasta Fp = Pallas base   = ark Pallas Fq
+const Q_FIELD_HASH: (u32, u32) = (0x116A9E, 1206); // pasta Fq = Vesta base   = ark Pallas Fr
+
+struct Dataset {
+    g0: Vec<String>,
+    g1: Vec<String>,
+    g2: Vec<String>,
+    g3: Vec<String>, // length 129
+    inv: Vec<u8>,
+    trace: Vec<u64>, // (T-1)/2 little-endian limbs
+    hash_xor: u32,
+    hash_mod: u32,
+}
+
+fn build<F: PrimeField + FftField>(hash_xor: u32, hash_mod: u32) -> Dataset {
+    assert_eq!(
+        F::TWO_ADICITY,
+        32,
+        "the Sarkar2020 variant splits S into four 8-bit windows and assumes S == 32"
+    );
+
+    // Follows from `zcash/pasta_curves`'s `SqrtTables`
+    let g = F::TWO_ADIC_ROOT_OF_UNITY;
+    let row = |base: F, n: usize| -> Vec<F> {
+        let mut v = Vec::with_capacity(n);
+        let mut acc = F::ONE;
+        for _ in 0..n {
+            v.push(acc);
+            acc *= base;
+        }
+        v
+    };
+
+    // g0[i] = g^i, g1[i] = g^(2^8 i), g2[i] = g^(2^16 i), g3full[i] = g^(2^24 i).
+    let g0 = row(g, 256);
+    let g_8 = g0[255] * g; // g^(2^8)
+    let g1 = row(g_8, 256);
+    let g_16 = g1[255] * g_8; // g^(2^16)
+    let g2 = row(g_16, 256);
+    let g_24 = g2[255] * g_16; // g^(2^24)
+    let g3full = row(g_24, 256);
+
+    let low32 = |x: &F| -> usize { (x.into_bigint().as_ref()[0] as u32) as usize };
+    let hash = |x: &F| -> usize { (low32(x) ^ (hash_xor as usize)) % (hash_mod as usize) };
+
+    // inv maps g^(2^24 j) -> (256 - j) & 0xFF over the order-256 subgroup.
+    let mut inv = vec![0u8; hash_mod as usize];
+    let mut used = vec![false; hash_mod as usize];
+    for j in 0..256usize {
+        let h = hash(&g3full[j]);
+        assert!(
+            !used[h],
+            "perfect-hash collision at j={j}: params (xor={hash_xor:#x}, mod={hash_mod}) are \
+             not valid for this field's ROOT_OF_UNITY"
+        );
+        used[h] = true;
+        inv[h] = ((256 - j) & 0xFF) as u8;
+    }
+
+    let dec = |v: &[F]| {
+        v.iter()
+            .map(|e| e.into_bigint().to_string())
+            .collect::<Vec<_>>()
+    };
+    let trace = F::TRACE_MINUS_ONE_DIV_TWO.as_ref().to_vec();
+
+    Dataset {
+        g0: dec(&g0),
+        g1: dec(&g1),
+        g2: dec(&g2),
+        g3: dec(&g3full[..129]),
+        inv,
+        trace,
+        hash_xor,
+        hash_mod,
+    }
+}
+
+fn emit(type_name: &str, d: &Dataset) -> String {
+    let mont_array = |name: &str, vals: &[String]| -> String {
+        let mut s = String::new();
+        write!(s, "const {name}: &[{type_name}] = &[\n").unwrap();
+        for v in vals {
+            write!(s, "    MontFp!(\"{v}\"),\n").unwrap();
+        }
+        s.push_str("];\n\n");
+        s
+    };
+
+    let mut out = String::new();
+    out.push_str(
+        "// @generated by `cargo test -p ark-pallas --test gen_sqrt_tables -- --ignored`.\n\
+         // Sarkar (2020) square-root tables; see `ark_ff::SqrtPrecomputation::Sarkar2020`.\n\
+         // Do not edit by hand.\n\
+         #![allow(clippy::all)]\n\n\
+         use ark_ff::{MontFp, SqrtPrecomputation};\n",
+    );
+    write!(out, "use super::{type_name};\n\n").unwrap();
+
+    let trace = d
+        .trace
+        .iter()
+        .map(|l| format!("0x{l:016x}"))
+        .collect::<Vec<_>>()
+        .join(", ");
+    write!(
+        out,
+        "pub(crate) const SQRT_PRECOMP: Option<SqrtPrecomputation<{type_name}>> =\n    \
+         Some(SqrtPrecomputation::Sarkar2020 {{\n        \
+         trace_minus_one_div_two: &[{trace}],\n        \
+         g0: G0,\n        g1: G1,\n        g2: G2,\n        g3: G3,\n        \
+         inv: INV,\n        hash_xor: {:#x},\n        hash_mod: {},\n    }});\n\n",
+        d.hash_xor, d.hash_mod
+    )
+    .unwrap();
+
+    out.push_str(&mont_array("G0", &d.g0));
+    out.push_str(&mont_array("G1", &d.g1));
+    out.push_str(&mont_array("G2", &d.g2));
+    out.push_str(&mont_array("G3", &d.g3));
+
+    out.push_str("const INV: &[u8] = &[\n");
+    for chunk in d.inv.chunks(16) {
+        out.push_str("    ");
+        for b in chunk {
+            write!(out, "{b}, ").unwrap();
+        }
+        out.push('\n');
+    }
+    out.push_str("];\n");
+    out
+}
+
+#[test]
+#[ignore = "regenerates committed source files; run explicitly"]
+fn generate() {
+    // ark_pallas::Fq is the p field; ark_pallas::Fr is the q field.
+    let p = build::<ark_pallas::Fq>(P_FIELD_HASH.0, P_FIELD_HASH.1);
+    let q = build::<ark_pallas::Fr>(Q_FIELD_HASH.0, Q_FIELD_HASH.1);
+
+    let pallas = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+
+    let targets = [
+        (pallas.join("src/fields/fq_sqrt_table.rs"), "Fq", &p),
+        (pallas.join("src/fields/fr_sqrt_table.rs"), "Fr", &q),
+    ];
+
+    for (path, ty, data) in targets {
+        fs::write(&path, emit(ty, data)).unwrap();
+        println!("wrote {}", path.display());
+    }
+}

curves/pallas/tests/sqrt_bench.rs

@@ -0,0 +1,56 @@
+//! Run with:
+//! `cargo test --release --manifest-path curves/pallas/Cargo.toml --test sqrt_bench -- --ignored --nocapture`
+
+use ark_ff::{FftField, PrimeField, SqrtPrecomputation};
+use ark_std::test_rng;
+use std::time::Instant;
+
+#[test]
+#[ignore = "timing comparison; run explicitly with --release --nocapture"]
+fn compare() {
+    bench::<ark_pallas::Fq>("Pallas Fq (p)");
+    bench::<ark_pallas::Fr>("Pallas Fr (q)");
+}
+
+fn bench<F: PrimeField + FftField>(name: &str) {
+    let rng = &mut test_rng();
+    // Squares only, so both algorithms take the "found a root" path.
+    let inputs: Vec<F> = (0..1000).map(|_| F::rand(rng).square()).collect();
+
+    // The installed precomputation (Sarkar2020 for these fields).
+    let sarkar_owned = F::SQRT_PRECOMP;
+    let sarkar = sarkar_owned.as_ref().expect("Sarkar precomp installed");
+
+    // A reference Tonelli-Shanks precomputation for the same field.
+    let trace: &'static [u64] = Box::leak(
+        F::TRACE_MINUS_ONE_DIV_TWO
+            .as_ref()
+            .to_vec()
+            .into_boxed_slice(),
+    );
+    let ts: SqrtPrecomputation<F> = SqrtPrecomputation::TonelliShanks {
+        two_adicity: F::TWO_ADICITY,
+        quadratic_nonresidue_to_trace: F::TWO_ADIC_ROOT_OF_UNITY,
+        trace_of_modulus_minus_one_div_two: trace,
+    };
+
+    let run = |p: &SqrtPrecomputation<F>| {
+        let t = Instant::now();
+        let mut acc = vec![];
+        for x in &inputs {
+            acc.push(p.sqrt(x).unwrap());
+        }
+        (t.elapsed(), acc)
+    };
+
+    let (ts_time, a1) = run(&ts);
+    let (sk_time, a2) = run(sarkar);
+    assert_eq!(a1, a2);
+
+    println!(
+        "{name}: tonelli-shanks {:?}, sarkar {:?}  ({:.2}x)",
+        ts_time,
+        sk_time,
+        ts_time.as_secs_f64() / sk_time.as_secs_f64()
+    );
+}

curves/pallas/tests/sqrt_sarkar.rs

@@ -0,0 +1,47 @@
+use ark_ff::{Field, LegendreSymbol};
+use ark_std::test_rng;
+
+fn check<F: Field>() {
+    assert!(F::SQRT_PRECOMP.is_some());
+
+    let rng = &mut test_rng();
+
+    assert_eq!(F::ZERO.sqrt(), Some(F::ZERO));
+    let one = F::ONE;
+    assert_eq!(one.sqrt().map(|r| r * r), Some(one));
+
+    let mut squares = 0u32;
+    let mut nonsquares = 0u32;
+    for _ in 0..3000 {
+        let x = F::rand(rng);
+        match x.legendre() {
+            LegendreSymbol::Zero => {},
+            LegendreSymbol::QuadraticResidue => {
+                squares += 1;
+                let r = x
+                    .sqrt()
+                    .expect("a quadratic residue must have a square root");
+                assert_eq!(r * r, x, "returned root does not square back to input");
+            },
+            LegendreSymbol::QuadraticNonResidue => {
+                nonsquares += 1;
+                assert!(
+                    x.sqrt().is_none(),
+                    "a quadratic non-residue must not yield a square root"
+                );
+            },
+        }
+    }
+    // Sanity that we actually exercised both branches.
+    assert!(squares > 0 && nonsquares > 0);
+}
+
+#[test]
+fn pallas_fq_p_field() {
+    check::<ark_pallas::Fq>();
+}
+
+#[test]
+fn pallas_fr_q_field() {
+    check::<ark_pallas::Fr>();
+}

ff-macros/src/lib.rs

@@ -86,6 +86,7 @@ pub fn define_field(input: TokenStream) -> TokenStream {
             generator_big,
             small_subgroup_base,
             small_subgroup_power,
+            None,
             config_name.clone(),
         );
 
@@ -110,7 +111,13 @@ pub fn define_field(input: TokenStream) -> TokenStream {
 // This code was adapted from the `PrimeField` Derive Macro in ff-derive.
 #[proc_macro_derive(
     MontConfig,
-    attributes(modulus, generator, small_subgroup_base, small_subgroup_power)
+    attributes(
+        modulus,
+        generator,
+        small_subgroup_base,
+        small_subgroup_power,
+        sqrt_precomp
+    )
 )]
 pub fn mont_config(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     // Parse the type definition
@@ -135,11 +142,17 @@ pub fn mont_config(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let small_subgroup_power: Option<u32> = fetch_attr("small_subgroup_power", &ast.attrs)
         .map(|s| s.parse().expect("small_subgroup_power should be a number"));
 
+    // Optional path to a `const SQRT_PRECOMP: Option<SqrtPrecomputation<F>>` that
+    // overrides the default (Tonelli-Shanks / Case3Mod4 / Case5Mod8) precomputation.
+    // Used to plug in a table-based square root for high-2-adicity fields.
+    let sqrt_precomp: Option<String> = fetch_attr("sqrt_precomp", &ast.attrs);
+
     montgomery::mont_config_helper(
         modulus,
         generator,
         small_subgroup_base,
         small_subgroup_power,
+        sqrt_precomp,
         ast.ident,
     )
     .into()