[PlonK] Optimize PlonK setup and prover with O(n log n) FFT / coset NTT polynomial evaluations

[satran004]

Background & Objective

In PlonKSetupBLS381.java and PlonKProverBLS381.java, polynomial evaluations and coset evaluations currently use quadratic loops in certain hotspots (cosetEval, evalOmegaPow).

This $O(n^2)$ behavior throttles PlonK proving for larger circuit sizes ($n > 2^{14}$).

Scope & Tasks

• Replace naive evaluation loops with radix-2 Fast Fourier Transforms (FFT / IFFT) and coset NTTs over the BLS12-381 scalar field.
• Implement fast domain evaluation structures for powers of roots of unity.
• Validate polynomial arithmetic against existing PlonK test vectors.
• Benchmark setup and prover throughput on larger PlonK circuits.

Risk Classification

• R1: Proof system performance & algorithmic scalability.
• Reference: docs/zeroj-comprehensive-review-and-readiness-gemini.md (Section 5.2).

[satran004]

Right symptom, wrong diagnosis — and the real fix is about ten lines, not an FFT implementation.

The proposed remedy is already in the repo

Task 1 says "Replace naive evaluation loops with radix-2 Fast Fourier Transforms (FFT / IFFT) and coset NTTs." Those exist and are already used throughout: FieldFFTBLS381.fft/ifft runs across PlonKSetupBLS381.java:118-125, and cosetEval itself ends in an FFT (PlonKProverBLS381.java:467-473). Anyone picking this up as written would implement an FFT that is already there.

Where the Θ(n²) actually is

powFr(base, exp) is a linear loop:

// PlonKProverBLS381.java:554
private static MontFr381 powFr(MontFr381 base, int exp) {
    if (exp == 0) return MontFr381.ONE;
    MontFr381 r = base;
    for (int i = 1; i < exp; i++) r = r.mul(base);   // O(exp)
    return r;
}

and cosetEval calls it inside a loop over i:

// PlonKProverBLS381.java:467-473
for (int i = 0; i < n; i++)
    padded[i] = i < coeffs.length ? coeffs[i].mul(powFr(shift, i)) : MontFr381.ZERO;

That is Θ(n²) field multiplications, and cosetEval is called ~15 times per proof (:234-259).

The same pattern is in setup — evalOmegaPow (PlonKSetupBLS381.java:300-304) is the identical linear loop, called as evalOmegaPow(omega, i) in the domainSize loop (:102) and three more times per gate via decodeSigmaTarget (:153). So PlonK setup is Θ(n²) too, which is the harder wall.

Minor, same family: zetaN at PlonKProverBLS381.java:394-395 is an O(n) loop where square-and-multiply is O(log n).

The fix

A running-product accumulator — which cosetIFFT, twelve lines below cosetEval in the same file, already does correctly:

// PlonKProverBLS381.java:476-491 — the correct pattern, already present
MontFr381 power = MontFr381.ONE;
for (int i = 0; i < coeffs.length; i++) {
    coeffs[i] = coeffs[i].mul(power);
    power = power.mul(shiftInv);
}

Apply that shape in cosetEval, precompute one omega-power table in setup, and switch zetaN/powFr to square-and-multiply. Orders of magnitude at n ≥ 2¹⁴, for roughly ten lines.

Suggested retitle

"Fix Θ(n²) power-of-root loops in PlonK setup and prover (powFr / evalOmegaPow)". Existing PlonK test vectors are the validation gate — the outputs must be unchanged, this is pure algorithmic substitution. R1 is right, but this is the best effort-to-value ratio in the batch.