Skip to content

perf(math): Add #[inline] attributes to hot field arithmetic paths #1128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
diegokingston merged 2 commits into from
Feb 3, 2026
Merged

perf(math): Add #[inline] attributes to hot field arithmetic paths #1128

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ed0c12d
perf(math): Add #[inline] attributes to hot field arithmetic paths
diegokingston Feb 3, 2026
94c4c0e
Merge branch 'main' into perf/inline-attributes
diegokingston Feb 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions crates/math/src/field/element.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,7 @@ where
F::BaseType: Clone,
F: IsField,
{
#[inline]
fn from(value: &F::BaseType) -> Self {
Self {
value: F::from_base_type(value.clone()),
Expand All @@ -132,6 +133,7 @@ impl<F> From<u64> for FieldElement<F>
where
F: IsField,
{
#[inline]
fn from(value: u64) -> Self {
Self {
value: F::from_u64(value),
Expand Down Expand Up @@ -159,6 +161,7 @@ where
F::BaseType: Clone,
F: IsField,
{
#[inline(always)]
pub fn from_raw(value: F::BaseType) -> Self {
Self { value }
}
Expand All @@ -173,6 +176,7 @@ impl<F> PartialEq<FieldElement<F>> for FieldElement<F>
where
F: IsField,
{
#[inline]
fn eq(&self, other: &FieldElement<F>) -> bool {
F::eq(&self.value, &other.value)
}
Expand Down Expand Up @@ -399,6 +403,7 @@ where
{
type Output = Result<FieldElement<L>, FieldError>;

#[inline]
fn div(self, rhs: &FieldElement<L>) -> Self::Output {
let value = <F as IsSubFieldOf<L>>::div(&self.value, &rhs.value)?;
Ok(FieldElement::<L> { value })
Expand All @@ -412,6 +417,7 @@ where
{
type Output = Result<FieldElement<L>, FieldError>;

#[inline]
fn div(self, rhs: FieldElement<L>) -> Self::Output {
&self / &rhs
}
Expand All @@ -424,6 +430,7 @@ where
{
type Output = Result<FieldElement<L>, FieldError>;

#[inline]
fn div(self, rhs: &FieldElement<L>) -> Self::Output {
&self / rhs
}
Expand All @@ -436,6 +443,7 @@ where
{
type Output = Result<FieldElement<L>, FieldError>;

#[inline]
fn div(self, rhs: FieldElement<L>) -> Self::Output {
self / &rhs
}
Expand Down Expand Up @@ -472,6 +480,7 @@ impl<F> Default for FieldElement<F>
where
F: IsField,
{
#[inline]
fn default() -> Self {
Self { value: F::zero() }
}
Expand Down Expand Up @@ -545,6 +554,7 @@ where
}

/// Returns the raw base type
#[inline(always)]
pub fn to_raw(self) -> F::BaseType {
self.value
}
Expand Down
19 changes: 19 additions & 0 deletions crates/math/src/field/fields/mersenne31/field.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ use core::fmt::{self, Display};
pub struct Mersenne31Field;

impl Mersenne31Field {
#[inline(always)]
fn weak_reduce(n: u32) -> u32 {
// To reduce 'n' to 31 bits we clear its MSB, then add it back in its reduced form.
let msb = n & (1 << 31);
Expand All @@ -27,6 +28,7 @@ impl Mersenne31Field {
res + msb_reduced
}

#[inline(always)]
fn as_representative(n: &u32) -> u32 {
if *n == MERSENNE_31_PRIME_FIELD_ORDER {
0
Expand All @@ -44,12 +46,14 @@ impl Mersenne31Field {
}

/// Computes a * 2^k, with 0 < k < 31
#[inline(always)]
pub fn mul_power_two(a: u32, k: u32) -> u32 {
let msb = (a & (u32::MAX << (31 - k))) >> (31 - k); // The k + 1 msf shifted right .
let lsb = (a & (u32::MAX >> (k + 1))) << k; // The 31 - k lsb shifted left.
Self::weak_reduce(msb + lsb)
}

#[inline]
pub fn pow_2(a: &u32, order: u32) -> u32 {
let mut res = *a;
(0..order).for_each(|_| res = Self::square(&res));
Expand All @@ -58,6 +62,7 @@ impl Mersenne31Field {

/// TODO: See if we can optimize this function.
/// Computes 2a^2 - 1
#[inline(always)]
pub fn two_square_minus_one(a: &u32) -> u32 {
if *a == 0 {
MERSENNE_31_PRIME_FIELD_ORDER - 1
Expand All @@ -76,6 +81,7 @@ impl IsField for Mersenne31Field {
type BaseType = u32;

/// Returns the sum of `a` and `b`.
#[inline(always)]
fn add(a: &u32, b: &u32) -> u32 {
// We are using that if a and b are field elements of Mersenne31, then
// a + b has at most 32 bits, so we can use the weak_reduce function to take mudulus p.
Expand All @@ -84,21 +90,25 @@ impl IsField for Mersenne31Field {

/// Returns the multiplication of `a` and `b`.
// Note: for powers of 2 we can perform bit shifting this would involve overriding the trait implementation
#[inline(always)]
fn mul(a: &u32, b: &u32) -> u32 {
Self::from_u64(u64::from(*a) * u64::from(*b))
}

#[inline(always)]
fn sub(a: &u32, b: &u32) -> u32 {
Self::weak_reduce(a + MERSENNE_31_PRIME_FIELD_ORDER - b)
}

/// Returns the additive inverse of `a`.
#[inline(always)]
fn neg(a: &u32) -> u32 {
// NOTE: MODULUS known to have 31 bit clear
MERSENNE_31_PRIME_FIELD_ORDER - a
}

/// Returns the multiplicative inverse of `a`.
#[inline]
fn inv(x: &u32) -> Result<u32, FieldError> {
if *x == Self::zero() || *x == MERSENNE_31_PRIME_FIELD_ORDER {
return Err(FieldError::InvZeroError);
Expand All @@ -117,36 +127,44 @@ impl IsField for Mersenne31Field {
}

/// Returns the division of `a` and `b`.
#[inline]
fn div(a: &u32, b: &u32) -> Result<u32, FieldError> {
let b_inv = Self::inv(b).map_err(|_| FieldError::DivisionByZero)?;
Ok(Self::mul(a, &b_inv))
}

/// Returns a boolean indicating whether `a` and `b` are equal or not.
#[inline(always)]
fn eq(a: &u32, b: &u32) -> bool {
Self::as_representative(a) == Self::representative(b)
}

/// Returns the additive neutral element.
#[inline(always)]
fn zero() -> u32 {
0u32
}

/// Returns the multiplicative neutral element.
#[inline(always)]
fn one() -> u32 {
1u32
}

/// Returns the element `x * 1` where 1 is the multiplicative neutral element.
#[inline(always)]
fn from_u64(x: u64) -> u32 {
(((((x >> 31) + x + 1) >> 31) + x) & (MERSENNE_31_PRIME_FIELD_ORDER as u64)) as u32
}

/// Takes as input an element of BaseType and returns the internal representation
/// of that element in the field.
#[inline(always)]
fn from_base_type(x: u32) -> u32 {
Self::weak_reduce(x)
}

#[inline(always)]
fn double(a: &u32) -> u32 {
Self::weak_reduce(a << 1)
}
Expand All @@ -157,6 +175,7 @@ impl IsPrimeField for Mersenne31Field {

// Since our invariant guarantees that `value` fits in 31 bits, there is only one possible value
// `value` that is not canonical, namely 2^31 - 1 = p = 0.
#[inline(always)]
fn representative(x: &u32) -> u32 {
debug_assert!((x >> 31) == 0);
Self::as_representative(x)
Copy link

@github-actions github-actions bot Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correctness

  • The mul_power_two implementation appears to be incorrectly handling the shifting of bits, leading to possible incorrect multiplication results when the bits carried exceed the maximum allowed.
  • The function pow_2 should ensure that squaring does not overflow, as continued squaring can exceed the 31-bit expectation.

Security

  • Functions like inv should ensure no secret-dependent branching for side-channel prevention; similar care is needed with functions like mul, sub, and neg for consistent timing.
  • The as_representative, neg, and inv methods may benefit from checking their handling of inputs like zero and the prime field order.

Performance

  • Excessive #[inline(always)] attributes may lead to code bloat; use them judiciously, especially for complex logic.
  • Review the mul conversion from u32 to u64 and back to ensure minimal overhead and evaluate if it's performant enough.

Bugs & Errors

  • The sub function uses a subtraction which could underflow if a < b.
  • Potential overflows in arithmetic functions considering operations return values larger than 31 bits which may lead to crashes.

Code Simplicity

  • The mul_power_two implementation lacks clarity and would benefit from comments or a simplified approach.
  • Multiple uses of weak_reduce imply a potential for consolidation or more straightforward flow; repeated usage may cover duplicated logic.

Additional Notes

  • Ensure rigorous testing for edge cases, such as successive operations leading to zero or the prime field order, especially in reduced field context.
  • Consider adding comments or documentation to clarify assumptions, such as when bitwise operations are safe and why certain methods may not need full reductions.

Due to these concerns in correctness, security, performance, and potential bugs, this code is not ready to merge. Further refinement and testing are required to ensure all conditions are safely managed and optimized.

Expand Down
7 changes: 7 additions & 0 deletions crates/math/src/field/fields/u64_goldilocks_field.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,7 @@ impl IsField for Goldilocks64Field {
}

/// Returns the multiplicative inverse of `a` using optimized addition chain.
#[inline]
fn inv(a: &u64) -> Result<u64, FieldError> {
let canonical = canonicalize(*a);
if canonical == 0 {
Expand All @@ -141,6 +142,7 @@ impl IsField for Goldilocks64Field {
}

/// Returns the division of `a` and `b`.
#[inline]
fn div(a: &u64, b: &u64) -> Result<u64, FieldError> {
let b_inv = <Self as IsField>::inv(b)?;
Ok(<Self as IsField>::mul(a, &b_inv))
Expand Down Expand Up @@ -415,6 +417,7 @@ impl IsField for Degree2GoldilocksExtensionField {

/// Returns the multiplicative inverse of `a`:
/// (a0 + a1*w)^-1 = (a0 - a1*w) / (a0^2 - 7*a1^2)
#[inline]
fn inv(a: &Self::BaseType) -> Result<Self::BaseType, FieldError> {
let a0_sq = a[0].square();
let a1_sq = a[1].square();
Expand All @@ -424,19 +427,23 @@ impl IsField for Degree2GoldilocksExtensionField {
Ok([a[0] * norm_inv, -a[1] * norm_inv])
}

#[inline]
fn div(a: &Self::BaseType, b: &Self::BaseType) -> Result<Self::BaseType, FieldError> {
let b_inv = Self::inv(b)?;
Ok(<Self as IsField>::mul(a, &b_inv))
}

#[inline(always)]
fn eq(a: &Self::BaseType, b: &Self::BaseType) -> bool {
a[0] == b[0] && a[1] == b[1]
}

#[inline(always)]
fn zero() -> Self::BaseType {
[FpE::zero(), FpE::zero()]
}

#[inline(always)]
fn one() -> Self::BaseType {
[FpE::one(), FpE::zero()]
}
Copy link

@github-actions github-actions bot Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correctness:

  • inv function: The handling of zero in inv seems correct, as it returns an error when the input is zero, which is an expected behavior to prevent division by zero.

Security:

  • There is no indication in the provided snippet whether the operations are constant-time. Ensure that modular arithmetic and field operations are constant-time to prevent timing side-channel attacks.
  • Make sure that sensitive data is properly zeroized after use. I couldn't see if and where sensitive data is zeroized from the snippet.
  • There's no detail about randomness in this snippet, so it should be secured elsewhere if used.
  • Ensure hash function domain separation where applicable.

Performance:

  • Inlining: Adding #[inline] and #[inline(always)] hints is good for performance, but be cautious as #[inline(always)] can sometimes lead to code bloat, though this seems appropriate here.
  • Multiplicative Inverse: The implementation uses addition chains for inverse calculation, which is often performance-optimal.

Bugs & Errors:

  • Error handling: The usage of Result types for inversion and division functions is appropriate to handle errors such as division by zero.
  • Potential Panics/Unwraps: There are no obvious uses of unwrap from the snippet provided. Ensure any other parts of the code do not invoke unwrap without checks.
  • Memory Safety: Without the broader context, there's nothing that appears unsafe here, but ensure all buffer/array accesses are bounds-checked.

Code Simplicity:

  • The code is concise and focuses on using existing utility functions like inv and mul, which promotes readability and maintainability.
  • Using descriptive names like norm_inv, a0_sq, and a1_sq aids in understanding the algorithm.

Recommendation:

  • Ensure you perform all cryptographic operations, especially those involving secrets, in constant time.
  • Double-check that all potential error conditions are appropriately handled, especially when the code relies on external inputs or parameters.
  • Make sure sensitive data is zeroed out after use, especially in areas handling private keys or other secrets.

Expand Down
Loading