perf(ecsm): 7.7x faster witness generation

crypto/ecsm/Cargo.toml

@@ -7,8 +7,14 @@ license.workspace = true
 
 [dependencies]
 num-bigint = "0.4.6"
+num-integer = "0.1.46"
 num-traits = "0.2.19"
+rayon = { version = "1.8.0", optional = true }
 # Audited secp256k1 arithmetic (host-side witness generation only; never in the
-# constraint system). Used for executor scalar multiplication and for the projective
-# double-and-add replay + batch inversion that builds ECDAS step witnesses efficiently.
+# constraint system). Used for executor scalar multiplication and for the
+# hand-rolled Jacobian double-and-add replay that builds ECDAS step witnesses
+# efficiently.
 k256 = { version = "0.13", default-features = false, features = ["arithmetic", "expose-field"] }
+
+[features]
+parallel = ["dep:rayon"]

crypto/ecsm/examples/bench_witness.rs

@@ -0,0 +1,41 @@
+//! Timing harness for `compute_witness` (one ECSM ecall's witness).
+//! Run: cargo run --release --example bench_witness -p ecsm
+
+use std::time::Instant;
+
+// secp256k1 generator x-coordinate, big-endian.
+const GX_BE: [u8; 32] = [
+    0x79, 0xbe, 0x66, 0x7e, 0xf9, 0xdc, 0xbb, 0xac, 0x55, 0xa0, 0x62, 0x95, 0xce, 0x87, 0x0b, 0x07,
+    0x02, 0x9b, 0xfc, 0xdb, 0x2d, 0xce, 0x28, 0xd9, 0x59, 0xf2, 0x81, 0x5b, 0x16, 0xf8, 0x17, 0x98,
+];
+
+fn le32(be: &[u8; 32]) -> [u8; 32] {
+    let mut out = [0u8; 32];
+    for i in 0..32 {
+        out[i] = be[31 - i];
+    }
+    out
+}
+
+fn main() {
+    // Worst-case-ish scalar: high popcount → ~380 double/add steps.
+    let k_be: [u8; 32] = [
+        0xde, 0xad, 0xbe, 0xef, 0xca, 0xfe, 0xba, 0xbe, 0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd,
+        0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32, 0x10, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff,
+        0x11, 0x22,
+    ];
+    let k_le = le32(&k_be);
+    let xg_le = le32(&GX_BE);
+
+    for _ in 0..2 {
+        std::hint::black_box(ecsm::compute_witness(&k_le, &xg_le).unwrap());
+    }
+
+    const N: u32 = 20;
+    let t = Instant::now();
+    for _ in 0..N {
+        std::hint::black_box(ecsm::compute_witness(&k_le, &xg_le).unwrap());
+    }
+    let d = t.elapsed() / N;
+    println!("compute_witness: {d:?} per call ({N} runs)");
+}

crypto/ecsm/src/curve.rs

@@ -58,17 +58,19 @@ pub fn msb_position(k: &BigUint) -> u32 {
 }
 
 // =========================================================================
-// k256-backed fast path: projective double-and-add replay + batch inversion.
+// k256-backed fast path: hand-rolled Jacobian double-and-add replay (dbl-2009-l /
+// madd-2007-bl) over k256's public field arithmetic, plus two Montgomery batch
+// inversions (z-normalization and slope denominators).
 //
 // The witness generator is untrusted (the ECDAS chip re-proves every step), so
-// any audited arithmetic is sound here. We replay the schedule in k256
-// projective coordinates (no per-op inversion), `batch_normalize` all points to
-// affine in one shot, and batch-invert the slope denominators — replacing the
-// ~2*len_k Fermat inversions of the reference with two batched inversions.
+// any audited arithmetic is sound here. We replay the schedule in Jacobian
+// coordinates (no per-op inversion), batch-invert every z at once for the
+// Jacobian→affine conversion, and batch-invert the slope denominators —
+// replacing the ~2*len_k Fermat inversions of the reference with two batched
+// inversions.
 // =========================================================================
 
 use k256::elliptic_curve::ff::PrimeField as _;
-use k256::elliptic_curve::group::Curve as _;
 use k256::elliptic_curve::sec1::{FromEncodedPoint, ToEncodedPoint};
 use k256::{AffinePoint as K256Affine, EncodedPoint, FieldElement, ProjectivePoint, Scalar};
 
@@ -158,49 +160,122 @@ pub fn scalar_mul_affine_x(k: &BigUint, g: &AffinePoint) -> BigUint {
     from_k256_affine(&r).x
 }
 
-/// Replays the ECDAS double-and-add for `k·g` using k256 projective arithmetic and
-/// batched inversion. Produces the identical `StepPts` sequence as the BigUint
-/// reference replay (validated by the parity test in `tests::curve_tests`), but with
-/// two batched inversions instead of one per double/add step.
+/// Jacobian doubling (dbl-2009-l) for `y² = x³ + 7`: on `(X:Y:Z)` with
+/// `x = X/Z²`, `y = Y/Z³`. Intermediates are normalized where a later
+/// subtraction would otherwise negate a high-magnitude lazy value, and no
+/// subtraction ever negates a magnitude-2 value: k256's `negate(1)` requires
+/// the operand's magnitude to be ≤ 1 — a contract k256 *enforces with an
+/// assertion in debug builds* (release builds have slack, dev-profile tests
+/// don't).
+fn jac_double(
+    x: FieldElement,
+    y: FieldElement,
+    z: FieldElement,
+) -> (FieldElement, FieldElement, FieldElement) {
+    let a = x * x; // X1²
+    let b = y * y; // Y1²
+    let c = b * b; // B²
+    let d = ((x + b) * (x + b) - a - c).double().normalize(); // 2·((X1+B)² − A − C)
+    let e = a.double() + a; // 3A
+    let f = e * e; // E²
+    // F − 2D as two subtractions of the normalized d (never `d.double()`: that
+    // operand would be magnitude 2 and break negate(1)'s contract).
+    let x3 = (f - d - d).normalize(); // F − 2D
+    let c8 = FieldElement::from_u64(8) * c; // 8C (mul output, subtraction-safe)
+    let y3 = (e * (d - x3) - c8).normalize(); // E·(D − X3) − 8C
+    let z3 = (y * z).double(); // 2·Y1·Z1
+    (x3, y3, z3)
+}
+
+/// Mixed Jacobian+affine addition (madd-2007-bl); the affine operand has Z2 = 1.
+/// Same lazy-magnitude caveat as [`jac_double`].
+fn jac_madd(
+    x1: FieldElement,
+    y1: FieldElement,
+    z1: FieldElement,
+    x2: FieldElement,
+    y2: FieldElement,
+) -> (FieldElement, FieldElement, FieldElement) {
+    let z1z1 = z1 * z1;
+    let u2 = x2 * z1z1;
+    let s2 = y2 * z1 * z1z1;
+    let h = (u2 - x1).normalize();
+    let r = (s2 - y1).normalize();
+    let hh = h * h;
+    let hhh = h * hh;
+    let x1hh = (x1 * hh).normalize();
+    // R² − HHH − 2·X1·HH as two subtractions of the normalized x1hh (see jac_double).
+    let x3 = (r * r - hhh - x1hh - x1hh).normalize();
+    let y3 = (r * (x1hh - x3) - y1 * hhh).normalize();
+    let z3 = h * z1;
+    (x3, y3, z3)
+}
+
+/// Replays the ECDAS double-and-add for `k·g` in Jacobian coordinates over
+/// k256's public field arithmetic, with one batched inversion for every point
+/// and another for the slope denominators. Produces the identical `StepPts`
+/// sequence as the BigUint reference replay (validated by the parity test in
+/// `tests::curve_tests`).
+///
+/// Perf note: k256's `ProjectivePoint::batch_normalize` measured ~5-6ms for
+/// the `2·len_k` points of one witness — no better than per-point `to_affine`
+/// — while this hand-rolled path runs the same replay in ~0.5ms.
 pub fn replay_double_and_add(k: &BigUint, g: &AffinePoint) -> (Vec<StepPts>, AffinePoint) {
     let sched = schedule(k);
     if sched.is_empty() {
         return (Vec::new(), g.clone()); // k == 1: result is g, no steps
     }
     let n = sched.len();
+    let gx = fe_from_biguint(&g.x);
+    let gy = fe_from_biguint(&g.y);
 
-    // 1. projective replay (no inversions): record a and r at every step.
-    let g_proj = ProjectivePoint::from(to_k256_affine(g));
-    let mut a_proj = g_proj;
-    let mut points = Vec::with_capacity(2 * n); // [a_0..a_{n-1}, r_0..r_{n-1}]
-    let mut r_projs = Vec::with_capacity(n);
+    // 1. Jacobian replay (no inversions): record the n+1 DISTINCT points of the
+    //    ladder — a_i = pts[i], r_i = pts[i+1]. (Pushing a_i and r_i separately
+    //    would hold 2n entries with n−1 exact duplicates: r_i is a_{i+1}.)
+    let mut pts: Vec<(FieldElement, FieldElement, FieldElement)> = Vec::with_capacity(n + 1);
+    let (mut ax, mut ay, mut az) = (gx, gy, FieldElement::ONE);
+    pts.push((ax, ay, az));
     for &(_, op, _) in &sched {
-        let r_proj = if op == 0 {
-            a_proj.double()
+        let (rx, ry, rz) = if op == 0 {
+            jac_double(ax, ay, az)
         } else {
-            a_proj + g_proj
+            jac_madd(ax, ay, az, gx, gy)
         };
-        points.push(a_proj);
-        r_projs.push(r_proj);
-        a_proj = r_proj;
+        pts.push((rx, ry, rz));
+        (ax, ay, az) = (rx, ry, rz);
     }
-    points.extend_from_slice(&r_projs);
 
-    // 2. one batch_normalize for every a and r.
-    let mut affine = vec![K256Affine::IDENTITY; points.len()];
-    ProjectivePoint::batch_normalize(&points, &mut affine);
-    let a_aff: Vec<AffinePoint> = affine[..n].iter().map(from_k256_affine).collect();
-    let r_aff: Vec<AffinePoint> = affine[n..].iter().map(from_k256_affine).collect();
+    // 2. one batched inversion for every z (Jacobian: affine = (x/z², y/z³)).
+    //    Affine coordinates are kept in BOTH forms: FieldElement for the slope
+    //    algebra in steps 3-4 (they are mul outputs, so the subtractions there
+    //    stay within negate(1)'s contract), BigUint for the StepPts the witness
+    //    consumes — each value is converted exactly once, not converted and then
+    //    re-parsed.
+    let zs: Vec<FieldElement> = pts.iter().map(|p| p.2).collect();
+    let zinvs = batch_invert(&zs);
+    let aff_fe: Vec<(FieldElement, FieldElement)> = pts
+        .iter()
+        .zip(&zinvs)
+        .map(|(&(x, y, _), zi)| {
+            let zi2 = zi * zi;
+            (x * zi2, y * zi2 * zi)
+        })
+        .collect();
+    let aff: Vec<AffinePoint> = aff_fe
+        .iter()
+        .map(|&(x, y)| AffinePoint {
+            x: biguint_from_fe(&x),
+            y: biguint_from_fe(&y),
+        })
+        .collect();
 
-    // 3. batch-invert all slope denominators (add: xG-xA, double: 2yA).
-    let gx_fe = fe_from_biguint(&g.x);
-    let gy_fe = fe_from_biguint(&g.y);
+    // 3. batch-invert all slope denominators (add: xG−xA, double: 2yA).
     let denoms: Vec<FieldElement> = (0..n)
         .map(|i| {
             if sched[i].1 == 1 {
-                gx_fe - fe_from_biguint(&a_aff[i].x)
+                gx - aff_fe[i].0
             } else {
-                let ya = fe_from_biguint(&a_aff[i].y);
+                let ya = aff_fe[i].1;
                 ya + ya
             }
         })
@@ -211,26 +286,23 @@ pub fn replay_double_and_add(k: &BigUint, g: &AffinePoint) -> (Vec<StepPts>, Aff
     let steps: Vec<StepPts> = (0..n)
         .map(|i| {
             let num = if sched[i].1 == 1 {
-                gy_fe - fe_from_biguint(&a_aff[i].y)
+                gy - aff_fe[i].1
             } else {
-                let x2 = {
-                    let xa = fe_from_biguint(&a_aff[i].x);
-                    xa * xa
-                };
+                let x2 = aff_fe[i].0 * aff_fe[i].0;
                 x2 + x2 + x2 // 3 xA^2
             };
             StepPts {
-                a: a_aff[i].clone(),
+                a: aff[i].clone(),
                 g: g.clone(),
                 round: sched[i].0,
                 op: sched[i].1,
                 next_op: sched[i].2,
-                r: r_aff[i].clone(),
+                r: aff[i + 1].clone(),
                 lambda: biguint_from_fe(&(num * inv_denoms[i])),
             }
         })
         .collect();
 
-    let result = r_aff[n - 1].clone();
+    let result = aff[n].clone();
     (steps, result)
 }

crypto/ecsm/src/tests/curve_tests.rs

@@ -59,6 +59,38 @@ fn k256_replay_matches_reference() {
     }
 }
 
+/// Same parity sweep with a non-generator base point: production feeds the
+/// replay guest-supplied points (e.g. the recovered R in ecrecover), and every
+/// other test uses G.
+#[test]
+fn k256_replay_matches_reference_non_generator_base() {
+    let g = generator();
+    let base_x = scalar_mul_affine_x(&BigUint::from(5u64), &g);
+    let base = AffinePoint {
+        y: recover_y_canonical(&base_x).expect("base on curve"),
+        x: base_x,
+    };
+    let mut scalars: Vec<BigUint> = (1u64..40).map(BigUint::from).collect();
+    for &kv in &[0xFFu64, 0xABCD, 1 << 20, 123_456_789, u64::MAX] {
+        scalars.push(BigUint::from(kv));
+    }
+    scalars.push(&n() / BigUint::from(2u8));
+    scalars.push(&n() - BigUint::from(1u8));
+
+    for k in scalars {
+        let (steps, result) = replay_double_and_add(&k, &base);
+        let (steps_ref, result_ref) = replay_double_and_add_reference(&k, &base);
+        assert_eq!(
+            result, result_ref,
+            "final point mismatch for k = {k} (non-G base)"
+        );
+        assert_eq!(
+            steps, steps_ref,
+            "step list mismatch for k = {k} (non-G base)"
+        );
+    }
+}
+
 /// The executor's fast path (`scalar_mul_affine_x`) and the prover's replay must agree
 /// on `x(k·G)`: the executor writes it to guest memory and the prover proves it, so any
 /// divergence would make a correct execution unprovable. They run through two distinct

crypto/ecsm/src/witness.rs

@@ -17,7 +17,10 @@
 //! integer recurrence here; the prover converts the resulting integers to field elements.
 
 use num_bigint::{BigInt, BigUint};
+use num_integer::Integer;
 use num_traits::{Signed, Zero};
+#[cfg(feature = "parallel")]
+use rayon::prelude::*;
 
 use crate::curve::{StepPts, replay_double_and_add};
 use crate::{B, EcsmError, P_BYTES, R_BYTES, n, p, prepare, to_le_32};
@@ -256,11 +259,12 @@ fn to_le_33(relation: &str, v: &BigUint) -> [u8; 33] {
 /// `r + numerator / p`, where `numerator` must be divisible by `p`. Asserts divisibility
 /// and that the result is non-negative (guaranteed by the spec quotient ranges).
 fn shifted_quotient(relation: &str, numerator: &BigInt, p_big: &BigInt, r_big: &BigInt) -> BigUint {
+    let (q, rem) = numerator.div_rem(p_big);
     assert!(
-        (numerator % p_big).is_zero(),
+        rem.is_zero(),
         "ECSM witness {relation}: numerator not divisible by p"
     );
-    let q = r_big + numerator / p_big;
+    let q = r_big + q;
     assert!(
         !q.is_negative(),
         "ECSM witness {relation}: quotient unexpectedly negative"
@@ -325,6 +329,14 @@ pub fn compute_witness(k_le: &[u8; 32], xg_le: &[u8; 32]) -> Result<EcsmWitness,
     let y_r = to_le_32(&result.y);
     let x_r_sub_p = to_le_32(&((&two_256 + &result.x) - p()));
 
+    // Steps are independent witnesses (each builds its own λ/quotient/carry data
+    // from one StepPts), so they parallelize freely when rayon is available.
+    #[cfg(feature = "parallel")]
+    let steps = steps_pts
+        .par_iter()
+        .map(|s| build_step(s, &p_big, &r_big, &r_ext, &pp))
+        .collect();
+    #[cfg(not(feature = "parallel"))]
     let steps = steps_pts
         .iter()
         .map(|s| build_step(s, &p_big, &r_big, &r_ext, &pp))

prover/Cargo.toml

@@ -6,7 +6,7 @@ license.workspace = true
 
 [features]
 default = ["parallel"]
-parallel = ["stark/parallel", "math/parallel", "crypto/parallel", "dep:rayon"]
+parallel = ["stark/parallel", "math/parallel", "crypto/parallel", "ecsm/parallel", "dep:rayon"]
 cuda = ["stark/cuda"]
 test-cuda-faults = ["cuda", "stark/test-cuda-faults"]
 debug-checks = ["stark/debug-checks"]