Batch Inversion in Rust

Inversion is by far the most expensive basic operation in a prime field, and provers invert all the time: the denominators of a quotient polynomial over an evaluation domain, the slopes in affine point addition, the weights of a barycentric evaluation. Montgomery’s trick replaces $n$ inversions by one inversion and about $3n$ multiplications.

The trick

Write $p_i = a_0 a_1 \cdots a_{i-1}$ for the prefix products, with $p_0 = 1$ . Invert only the last one, $p_n^{-1}$ , then walk backwards. At every step two multiplications recover one inverse and shorten the running inverse by one factor:

a_i^{-1} = p_{i+1}^{-1} \cdot p_i, \qquad p_i^{-1} = p_{i+1}^{-1} \cdot a_i .

All the algorithm needs from the field is a multiplication, a zero test and a single inversion:

src/field.rs

use core::ops::{Mul, MulAssign};

/// The handful of operations batch inversion needs from a field element.
pub trait Field: Copy + Mul<Output = Self> + MulAssign {
    const ONE: Self;

    fn is_zero(&self) -> bool;

    /// Returns `None` for zero.
    fn invert(&self) -> Option<Self>;
}

The code

Zeros have no inverse, so they are skipped in both passes and left untouched. The highlighted lines are the single inversion and the two multiplications per element.

src/batch_inverse.rs

use crate::field::Field;

/// Inverts every non-zero element of `values` in place; zeros stay zero.
/// Costs one field inversion and about 3n multiplications.
pub fn batch_inverse<F: Field>(values: &mut [F]) {
    // Forward pass: prefix[i] = product of the non-zero values before i.
    let mut prefix = Vec::with_capacity(values.len());
    let mut acc = F::ONE;
    for v in values.iter() {
        prefix.push(acc);
        if !v.is_zero() {
            acc *= *v;
        }
    }

    // The only inversion. `acc` is a product of non-zero elements.
    let mut inv = acc.invert().expect("product of non-zero elements");

    // Backward pass: `inv` is the inverse of the product of values[..=i].
    for (v, p) in values.iter_mut().zip(prefix).rev() {
        if v.is_zero() {
            continue;
        }
        let next = inv * *v; // drop v from the running inverse
        *v = inv * p; // v⁻¹ = (a₀ ⋯ aᵢ)⁻¹ · (a₀ ⋯ aᵢ₋₁)
        inv = next;
    }
}

The function allocates one scratch vector of $n$ elements. When memory matters more than simplicity, the prefix products can be written into the output buffer of the caller instead, and the slice can be processed in chunks that fit the cache:

cargo test batch_inverse
cargo bench --bench inversion -- --sample-size 50

One caveat for cryptographic code: the branch on is_zero makes the running time depend on which elements are zero. That is fine for public data such as evaluation domains, and it is not fine for secrets.

The trick

The code

Comments