From 07de24c63ae24797751d73c9721785a898eaa63f Mon Sep 17 00:00:00 2001 From: diegokingston Date: Thu, 16 Apr 2026 16:22:58 -0300 Subject: [PATCH 1/5] perf: eliminate per-element Vec allocation in Merkle tree hashing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add WriteBytes trait for zero-allocation byte serialization and use it in the prover's Merkle commitment hot paths. Problem: as_bytes() returns Vec (heap allocation) for every field element. For CPU table commitment (2^20 rows × 74 cols), that's ~77M allocations of 8-byte Vecs. Benchmarks show this as a 1.3-1.7x slowdown. Solution: - WriteBytes trait: write_bytes_be(&self, buf: &mut [u8]) writes to caller-provided buffer, zero heap allocations per element. - Implemented for Goldilocks (8 bytes) and cubic extension (24 bytes). - commit_columns_bit_reversed: single buffer per row, write all columns directly, hash once. Eliminates both the row Vec and per-element Vec. - commit_composition_polynomial: same approach for composition pairs. The AsBytes trait and FieldElementVectorBackend are unchanged — only the prover hot paths are optimized. Verifier and generic code still use as_bytes() for backward compatibility. Benchmarked in lambdaworks (~/dev/lambdaworks bench/merkle-comparison): - 16 cols: 1.3-1.5x speedup, beats Plonky3 by 1.3-1.4x - 64 cols: 1.6-1.8x speedup, beats Plonky3 by 1.5-1.6x --- .../math/src/field/extensions_goldilocks.rs | 11 ++++ crypto/math/src/field/goldilocks.rs | 8 +++ crypto/math/src/traits.rs | 12 ++++ crypto/stark/src/prover.rs | 59 +++++++++++++------ 4 files changed, 71 insertions(+), 19 deletions(-) diff --git a/crypto/math/src/field/extensions_goldilocks.rs b/crypto/math/src/field/extensions_goldilocks.rs index 44e0d3f85..26f921a06 100644 --- a/crypto/math/src/field/extensions_goldilocks.rs +++ b/crypto/math/src/field/extensions_goldilocks.rs @@ -524,6 +524,17 @@ impl AsBytes for FieldElement { } } +impl crate::traits::WriteBytes for FieldElement { + const BYTE_LEN: usize = 24; // 3 × 8 bytes + #[inline(always)] + fn write_bytes_be(&self, buf: &mut [u8]) { + let components = self.value(); + components[0].write_bytes_be(&mut buf[0..8]); + components[1].write_bytes_be(&mut buf[8..16]); + components[2].write_bytes_be(&mut buf[16..24]); + } +} + impl HasDefaultTranscript for Degree3GoldilocksExtensionField { fn get_random_field_element_from_rng(rng: &mut impl rand::Rng) -> FieldElement { let mut sample = [0u8; 8]; diff --git a/crypto/math/src/field/goldilocks.rs b/crypto/math/src/field/goldilocks.rs index 42cd889e8..37f532ae5 100644 --- a/crypto/math/src/field/goldilocks.rs +++ b/crypto/math/src/field/goldilocks.rs @@ -482,6 +482,14 @@ impl AsBytes for FieldElement { } } +impl crate::traits::WriteBytes for FieldElement { + const BYTE_LEN: usize = 8; + #[inline(always)] + fn write_bytes_be(&self, buf: &mut [u8]) { + buf[..8].copy_from_slice(&self.canonical_u64().to_be_bytes()); + } +} + // Implement IsPrimeField for the native Goldilocks use crate::errors::CreationError; use crate::field::traits::IsPrimeField; diff --git a/crypto/math/src/traits.rs b/crypto/math/src/traits.rs index ed4d1589b..525b50c81 100644 --- a/crypto/math/src/traits.rs +++ b/crypto/math/src/traits.rs @@ -32,6 +32,18 @@ pub trait AsBytes { fn as_bytes(&self) -> alloc::vec::Vec; } +/// Zero-allocation byte serialization for Merkle tree hashing. +/// +/// Unlike `AsBytes::as_bytes()` which returns `Vec` (heap allocation per call), +/// this trait writes bytes directly into a caller-provided buffer. For Merkle trees +/// with millions of field elements, this eliminates millions of 8-byte allocations. +pub trait WriteBytes { + /// Byte length of this element's big-endian representation. + const BYTE_LEN: usize; + /// Write big-endian bytes into `buf[..BYTE_LEN]`. + fn write_bytes_be(&self, buf: &mut [u8]); +} + #[cfg(feature = "alloc")] impl AsBytes for u32 { fn as_bytes(&self) -> alloc::vec::Vec { diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 41d27fc66..f1d12d9f3 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -4,7 +4,6 @@ use std::sync::Arc; use std::time::Instant; use crypto::fiat_shamir::is_transcript::IsStarkTranscript; -use crypto::merkle_tree::traits::IsMerkleTreeBackend; use math::fft::cpu::bit_reversing::{in_place_bit_reverse_permute, reverse_index}; use math::fft::cpu::bowers_fft::LayerTwiddles; use math::fft::errors::FFTError; @@ -32,7 +31,7 @@ use crate::proof::stark::{DeepPolynomialOpenings, PolynomialOpenings}; use crate::table::Table; use crate::trace::LDETraceTable; -use super::config::{BatchedMerkleTree, BatchedMerkleTreeBackend, Commitment}; +use super::config::{BatchedMerkleTree, Commitment}; use super::constraints::evaluator::ConstraintEvaluator; use super::domain::Domain; use super::fri::fri_decommit::FriDecommitment; @@ -63,6 +62,9 @@ impl< FieldExtension: Send + Sync + IsField, PI, > IsStarkProver for Prover +where + FieldElement: math::traits::WriteBytes, + FieldElement: math::traits::WriteBytes, { } @@ -352,7 +354,9 @@ pub trait IsStarkProver< Field: IsSubFieldOf + IsFFTField + Send + Sync, FieldExtension: Send + Sync + IsField, PI, -> +> where + FieldElement: math::traits::WriteBytes, + FieldElement: math::traits::WriteBytes, { /// Returns the Merkle tree and the commitment to the vectors `vectors`. fn batch_commit_extension( @@ -379,28 +383,39 @@ pub trait IsStarkProver< columns: &[Vec>], ) -> Option<(BatchedMerkleTree, Commitment)> where - FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send + math::traits::WriteBytes, E: IsField, { + use math::traits::WriteBytes; + use sha3::Digest; + if columns.is_empty() || columns[0].is_empty() { return None; } let num_rows = columns[0].len(); let num_cols = columns.len(); + let byte_len = as WriteBytes>::BYTE_LEN; #[cfg(feature = "parallel")] let iter = (0..num_rows).into_par_iter(); #[cfg(not(feature = "parallel"))] let iter = 0..num_rows; + // Zero per-element heap allocations: write bytes directly from columns + // into a per-row buffer, then hash once. let hashed_leaves: Vec = iter .map(|row_idx| { let br_idx = reverse_index(row_idx, num_rows as u64); - let row: Vec> = (0..num_cols) - .map(|col_idx| columns[col_idx][br_idx].clone()) - .collect(); - BatchedMerkleTreeBackend::::hash_data(&row) + let total_bytes = num_cols * byte_len; + let mut buf = vec![0u8; total_bytes]; + for col_idx in 0..num_cols { + columns[col_idx][br_idx] + .write_bytes_be(&mut buf[col_idx * byte_len..(col_idx + 1) * byte_len]); + } + let mut hasher = sha3::Keccak256::new(); + hasher.update(&buf); + hasher.finalize().into() }) .collect(); @@ -690,8 +705,11 @@ pub trait IsStarkProver< ) -> Option<(BatchedMerkleTree, Commitment)> where FieldElement: AsBytes + Sync + Send, - FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send + math::traits::WriteBytes, { + use math::traits::WriteBytes; + use sha3::Digest; + let num_parts = lde_composition_poly_parts_evaluations.len(); if num_parts == 0 { return None; @@ -702,18 +720,15 @@ pub trait IsStarkProver< } let num_leaves = num_rows / 2; - debug_assert!( - num_rows.is_power_of_two(), - "num_rows must be a power of two for reverse_index to be correct" - ); debug_assert!( num_rows.is_power_of_two(), "num_rows must be a power of two for reverse_index" ); - // Skip the transpose + merge by computing leaf data inline. + let byte_len = as WriteBytes>::BYTE_LEN; + + // Zero per-element allocations: write bytes directly from column data. // Each leaf = row_pair[2*i] ++ row_pair[2*i+1] after bit-reverse. - // Build the merged leaf data and hash in one pass. #[cfg(feature = "parallel")] let iter = (0..num_leaves).into_par_iter(); #[cfg(not(feature = "parallel"))] @@ -723,14 +738,20 @@ pub trait IsStarkProver< .map(|leaf_idx| { let br_0 = reverse_index(2 * leaf_idx, num_rows as u64); let br_1 = reverse_index(2 * leaf_idx + 1, num_rows as u64); - let mut leaf = Vec::with_capacity(2 * num_parts); + let total_bytes = 2 * num_parts * byte_len; + let mut buf = vec![0u8; total_bytes]; + let mut offset = 0; for part in lde_composition_poly_parts_evaluations.iter() { - leaf.push(part[br_0].clone()); + part[br_0].write_bytes_be(&mut buf[offset..offset + byte_len]); + offset += byte_len; } for part in lde_composition_poly_parts_evaluations.iter() { - leaf.push(part[br_1].clone()); + part[br_1].write_bytes_be(&mut buf[offset..offset + byte_len]); + offset += byte_len; } - BatchedMerkleTreeBackend::::hash_data(&leaf) + let mut hasher = sha3::Keccak256::new(); + hasher.update(&buf); + hasher.finalize().into() }) .collect(); From e46bf85b4311dbc4ae3ed42f309fe3beae1bd0b4 Mon Sep 17 00:00:00 2001 From: Nicole Date: Fri, 17 Apr 2026 10:59:52 -0300 Subject: [PATCH 2/5] add debug_asserts, fix comments, remove dead code, add tests --- .../math/src/field/extensions_goldilocks.rs | 25 ++++++++++++++++ crypto/math/src/field/goldilocks.rs | 30 +++++++++++++++++++ crypto/stark/src/prover.rs | 22 +++----------- 3 files changed, 59 insertions(+), 18 deletions(-) diff --git a/crypto/math/src/field/extensions_goldilocks.rs b/crypto/math/src/field/extensions_goldilocks.rs index 26f921a06..40f933c2a 100644 --- a/crypto/math/src/field/extensions_goldilocks.rs +++ b/crypto/math/src/field/extensions_goldilocks.rs @@ -528,6 +528,7 @@ impl crate::traits::WriteBytes for FieldElement const BYTE_LEN: usize = 24; // 3 × 8 bytes #[inline(always)] fn write_bytes_be(&self, buf: &mut [u8]) { + debug_assert!(buf.len() >= 24); let components = self.value(); components[0].write_bytes_be(&mut buf[0..8]); components[1].write_bytes_be(&mut buf[8..16]); @@ -555,6 +556,30 @@ impl HasDefaultTranscript for Degree3GoldilocksExtensionField { } } +#[cfg(test)] +mod tests { + use super::*; + use crate::traits::WriteBytes; + + #[test] + fn write_bytes_be_matches_as_bytes() { + let cases = [ + FieldElement::::zero(), + FieldElement::::one(), + FieldElement::::new([ + FpE::from(1u64), + FpE::from(2u64), + FpE::from(3u64), + ]), + ]; + for elem in &cases { + let mut buf = [0u8; 24]; + elem.write_bytes_be(&mut buf); + assert_eq!(&buf[..], elem.as_bytes().as_slice()); + } + } +} + // ===================================================== // HELPER FUNCTIONS // ===================================================== diff --git a/crypto/math/src/field/goldilocks.rs b/crypto/math/src/field/goldilocks.rs index 37f532ae5..2aa0865c6 100644 --- a/crypto/math/src/field/goldilocks.rs +++ b/crypto/math/src/field/goldilocks.rs @@ -486,6 +486,7 @@ impl crate::traits::WriteBytes for FieldElement { const BYTE_LEN: usize = 8; #[inline(always)] fn write_bytes_be(&self, buf: &mut [u8]) { + debug_assert!(buf.len() >= 8); buf[..8].copy_from_slice(&self.canonical_u64().to_be_bytes()); } } @@ -551,3 +552,32 @@ impl HasDefaultTranscript for GoldilocksField { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::traits::WriteBytes; + + #[test] + fn write_bytes_be_matches_as_bytes() { + let cases = [ + FieldElement::::from(0u64), + FieldElement::::from(1u64), + FieldElement::::from(GOLDILOCKS_PRIME - 1), + ]; + for elem in &cases { + let mut buf = [0u8; 8]; + elem.write_bytes_be(&mut buf); + assert_eq!(&buf[..], elem.as_bytes().as_slice()); + } + } + + #[test] + fn write_bytes_be_matches_as_bytes_noncanonical() { + // Values stored as-is via from_raw are non-canonical (>= p) until serialized. + let elem = FieldElement::::from_raw(GOLDILOCKS_PRIME + 5); + let mut buf = [0u8; 8]; + elem.write_bytes_be(&mut buf); + assert_eq!(&buf[..], elem.as_bytes().as_slice()); + } +} diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index f1d12d9f3..9a6604d3b 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -358,20 +358,6 @@ pub trait IsStarkProver< FieldElement: math::traits::WriteBytes, FieldElement: math::traits::WriteBytes, { - /// Returns the Merkle tree and the commitment to the vectors `vectors`. - fn batch_commit_extension( - vectors: &[Vec>], - ) -> Option<(BatchedMerkleTree, Commitment)> - where - FieldElement: AsBytes + Sync + Send, - FieldElement: AsBytes + Sync + Send, - { - let tree = BatchedMerkleTree::build(vectors)?; - - let commitment = tree.root; - Some((tree, commitment)) - } - /// Builds a Merkle tree commitment from column-major LDE evaluations with /// bit-reverse permutation, without cloning the full evaluation matrix. /// @@ -402,8 +388,8 @@ pub trait IsStarkProver< #[cfg(not(feature = "parallel"))] let iter = 0..num_rows; - // Zero per-element heap allocations: write bytes directly from columns - // into a per-row buffer, then hash once. + // One allocation per row (was one per field element): write all columns + // into a single buffer, then hash once. let hashed_leaves: Vec = iter .map(|row_idx| { let br_idx = reverse_index(row_idx, num_rows as u64); @@ -727,8 +713,8 @@ pub trait IsStarkProver< let byte_len = as WriteBytes>::BYTE_LEN; - // Zero per-element allocations: write bytes directly from column data. - // Each leaf = row_pair[2*i] ++ row_pair[2*i+1] after bit-reverse. + // One allocation per leaf (was one per field element): write all parts + // into a single buffer. Each leaf = row_pair[2*i] ++ row_pair[2*i+1] after bit-reverse. #[cfg(feature = "parallel")] let iter = (0..num_leaves).into_par_iter(); #[cfg(not(feature = "parallel"))] From 33b6e500e277c589f67e6e03c430367224f8c670 Mon Sep 17 00:00:00 2001 From: Nicole Date: Fri, 17 Apr 2026 11:05:03 -0300 Subject: [PATCH 3/5] fix lint --- .../math/src/field/extensions_goldilocks.rs | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/crypto/math/src/field/extensions_goldilocks.rs b/crypto/math/src/field/extensions_goldilocks.rs index 40f933c2a..e7bc8ed19 100644 --- a/crypto/math/src/field/extensions_goldilocks.rs +++ b/crypto/math/src/field/extensions_goldilocks.rs @@ -556,6 +556,17 @@ impl HasDefaultTranscript for Degree3GoldilocksExtensionField { } } +// ===================================================== +// HELPER FUNCTIONS +// ===================================================== + +/// Multiply a field element by 7 (the quadratic non-residue). +/// Wraps the raw u64 implementation for use with FieldElement types. +#[inline(always)] +fn mul_by_7(a: &FpE) -> FpE { + FpE::from_raw(mul_by_7_raw(*a.value())) +} + #[cfg(test)] mod tests { use super::*; @@ -579,14 +590,3 @@ mod tests { } } } - -// ===================================================== -// HELPER FUNCTIONS -// ===================================================== - -/// Multiply a field element by 7 (the quadratic non-residue). -/// Wraps the raw u64 implementation for use with FieldElement types. -#[inline(always)] -fn mul_by_7(a: &FpE) -> FpE { - FpE::from_raw(mul_by_7_raw(*a.value())) -} From a41bdf8e91139656f4232bdedf9ad5451e9a862b Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Fri, 17 Apr 2026 14:18:04 -0300 Subject: [PATCH 4/5] refactor: make prover leaf-hashing hash-agnostic via backend's hash_bytes (#506) Add hash_bytes() inherent method to FieldElementVectorBackend that surfaces the generic digest parameter D. The prover now calls BatchedMerkleTreeBackend::::hash_bytes() instead of hardcoding sha3::Keccak256, so changing the config.rs type alias automatically propagates the hash function to the commit hot paths. Also adds missing debug_assert for power-of-two in commit_columns_bit_reversed for consistency with commit_composition_polynomial. --- .../backends/field_element_vector.rs | 17 +++++++++++++++++ crypto/stark/src/prover.rs | 17 ++++++++--------- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/crypto/crypto/src/merkle_tree/backends/field_element_vector.rs b/crypto/crypto/src/merkle_tree/backends/field_element_vector.rs index a97c0ce37..1023b3a03 100644 --- a/crypto/crypto/src/merkle_tree/backends/field_element_vector.rs +++ b/crypto/crypto/src/merkle_tree/backends/field_element_vector.rs @@ -71,6 +71,23 @@ impl Default for FieldElementVectorBackend } } +impl FieldElementVectorBackend +where + [u8; NUM_BYTES]: From>, +{ + /// Hash raw bytes using the same digest (`D`) as this backend's leaf hashing. + /// Enables callers to pre-serialize field elements into a byte buffer and hash + /// once, avoiding per-element allocations while staying consistent with the + /// backend's hash function. + pub fn hash_bytes(data: &[u8]) -> [u8; NUM_BYTES] { + let mut hasher = D::new(); + hasher.update(data); + let mut result = [0u8; NUM_BYTES]; + result.copy_from_slice(&hasher.finalize()); + result + } +} + impl IsMerkleTreeBackend for FieldElementVectorBackend where diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 9a6604d3b..66e76ca1e 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -31,7 +31,7 @@ use crate::proof::stark::{DeepPolynomialOpenings, PolynomialOpenings}; use crate::table::Table; use crate::trace::LDETraceTable; -use super::config::{BatchedMerkleTree, Commitment}; +use super::config::{BatchedMerkleTree, BatchedMerkleTreeBackend, Commitment}; use super::constraints::evaluator::ConstraintEvaluator; use super::domain::Domain; use super::fri::fri_decommit::FriDecommitment; @@ -373,7 +373,6 @@ pub trait IsStarkProver< E: IsField, { use math::traits::WriteBytes; - use sha3::Digest; if columns.is_empty() || columns[0].is_empty() { return None; @@ -383,6 +382,11 @@ pub trait IsStarkProver< let num_cols = columns.len(); let byte_len = as WriteBytes>::BYTE_LEN; + debug_assert!( + num_rows.is_power_of_two(), + "num_rows must be a power of two for reverse_index" + ); + #[cfg(feature = "parallel")] let iter = (0..num_rows).into_par_iter(); #[cfg(not(feature = "parallel"))] @@ -399,9 +403,7 @@ pub trait IsStarkProver< columns[col_idx][br_idx] .write_bytes_be(&mut buf[col_idx * byte_len..(col_idx + 1) * byte_len]); } - let mut hasher = sha3::Keccak256::new(); - hasher.update(&buf); - hasher.finalize().into() + BatchedMerkleTreeBackend::::hash_bytes(&buf) }) .collect(); @@ -694,7 +696,6 @@ pub trait IsStarkProver< FieldElement: AsBytes + Sync + Send + math::traits::WriteBytes, { use math::traits::WriteBytes; - use sha3::Digest; let num_parts = lde_composition_poly_parts_evaluations.len(); if num_parts == 0 { @@ -735,9 +736,7 @@ pub trait IsStarkProver< part[br_1].write_bytes_be(&mut buf[offset..offset + byte_len]); offset += byte_len; } - let mut hasher = sha3::Keccak256::new(); - hasher.update(&buf); - hasher.finalize().into() + BatchedMerkleTreeBackend::::hash_bytes(&buf) }) .collect(); From 65c73aff50cb20e49d22478d29decf13d7f2cb11 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Fri, 17 Apr 2026 14:34:17 -0300 Subject: [PATCH 5/5] refactor: fold WriteBytes into ByteConversion trait Consolidate serialization traits by adding BYTE_LEN and write_bytes_be to ByteConversion instead of having a separate WriteBytes trait. Goldilocks and cubic extension override write_bytes_be with zero-alloc implementations; other types use the default which delegates to to_bytes_be. --- .../math/src/field/extensions_goldilocks.rs | 29 ++++++++++--------- crypto/math/src/field/goldilocks.rs | 19 ++++++------ .../src/field/test_fields/u32_test_field.rs | 2 ++ crypto/math/src/traits.rs | 27 ++++++++--------- crypto/stark/src/prover.rs | 20 ++++++------- 5 files changed, 51 insertions(+), 46 deletions(-) diff --git a/crypto/math/src/field/extensions_goldilocks.rs b/crypto/math/src/field/extensions_goldilocks.rs index e7bc8ed19..f12af9ce9 100644 --- a/crypto/math/src/field/extensions_goldilocks.rs +++ b/crypto/math/src/field/extensions_goldilocks.rs @@ -12,6 +12,8 @@ use crate::field::{ use crate::traits::{AsBytes, ByteConversion}; impl ByteConversion for [FpE; 2] { + const BYTE_LEN: usize = 16; + #[cfg(feature = "alloc")] fn to_bytes_be(&self) -> alloc::vec::Vec { unimplemented!() @@ -38,6 +40,8 @@ impl ByteConversion for [FpE; 2] { } impl ByteConversion for [FpE; 3] { + const BYTE_LEN: usize = 24; + #[cfg(feature = "alloc")] fn to_bytes_be(&self) -> alloc::vec::Vec { let mut bytes = ByteConversion::to_bytes_be(&self[2]); @@ -470,6 +474,17 @@ impl Fp3E { // ===================================================== impl ByteConversion for FieldElement { + const BYTE_LEN: usize = 24; + + #[inline(always)] + fn write_bytes_be(&self, buf: &mut [u8]) { + debug_assert!(buf.len() >= 24); + let components = self.value(); + components[0].write_bytes_be(&mut buf[0..8]); + components[1].write_bytes_be(&mut buf[8..16]); + components[2].write_bytes_be(&mut buf[16..24]); + } + #[cfg(feature = "alloc")] fn to_bytes_be(&self) -> alloc::vec::Vec { let mut byte_slice = ByteConversion::to_bytes_be(&self.value()[0]); @@ -524,18 +539,6 @@ impl AsBytes for FieldElement { } } -impl crate::traits::WriteBytes for FieldElement { - const BYTE_LEN: usize = 24; // 3 × 8 bytes - #[inline(always)] - fn write_bytes_be(&self, buf: &mut [u8]) { - debug_assert!(buf.len() >= 24); - let components = self.value(); - components[0].write_bytes_be(&mut buf[0..8]); - components[1].write_bytes_be(&mut buf[8..16]); - components[2].write_bytes_be(&mut buf[16..24]); - } -} - impl HasDefaultTranscript for Degree3GoldilocksExtensionField { fn get_random_field_element_from_rng(rng: &mut impl rand::Rng) -> FieldElement { let mut sample = [0u8; 8]; @@ -570,7 +573,7 @@ fn mul_by_7(a: &FpE) -> FpE { #[cfg(test)] mod tests { use super::*; - use crate::traits::WriteBytes; + use crate::traits::ByteConversion; #[test] fn write_bytes_be_matches_as_bytes() { diff --git a/crypto/math/src/field/goldilocks.rs b/crypto/math/src/field/goldilocks.rs index 2aa0865c6..9891071b3 100644 --- a/crypto/math/src/field/goldilocks.rs +++ b/crypto/math/src/field/goldilocks.rs @@ -434,6 +434,14 @@ impl GoldilocksElement { // ===================================================== impl ByteConversion for FieldElement { + const BYTE_LEN: usize = 8; + + #[inline(always)] + fn write_bytes_be(&self, buf: &mut [u8]) { + debug_assert!(buf.len() >= 8); + buf[..8].copy_from_slice(&self.canonical_u64().to_be_bytes()); + } + #[cfg(feature = "alloc")] fn to_bytes_be(&self) -> alloc::vec::Vec { self.canonical_u64().to_be_bytes().to_vec() @@ -482,15 +490,6 @@ impl AsBytes for FieldElement { } } -impl crate::traits::WriteBytes for FieldElement { - const BYTE_LEN: usize = 8; - #[inline(always)] - fn write_bytes_be(&self, buf: &mut [u8]) { - debug_assert!(buf.len() >= 8); - buf[..8].copy_from_slice(&self.canonical_u64().to_be_bytes()); - } -} - // Implement IsPrimeField for the native Goldilocks use crate::errors::CreationError; use crate::field::traits::IsPrimeField; @@ -556,7 +555,7 @@ impl HasDefaultTranscript for GoldilocksField { #[cfg(test)] mod tests { use super::*; - use crate::traits::WriteBytes; + use crate::traits::ByteConversion; #[test] fn write_bytes_be_matches_as_bytes() { diff --git a/crypto/math/src/field/test_fields/u32_test_field.rs b/crypto/math/src/field/test_fields/u32_test_field.rs index e030f81f8..428f7b7c8 100644 --- a/crypto/math/src/field/test_fields/u32_test_field.rs +++ b/crypto/math/src/field/test_fields/u32_test_field.rs @@ -11,6 +11,8 @@ use crate::traits::ByteConversion; pub struct U32Field; impl ByteConversion for u32 { + const BYTE_LEN: usize = 4; + #[cfg(feature = "alloc")] fn to_bytes_be(&self) -> alloc::vec::Vec { unimplemented!() diff --git a/crypto/math/src/traits.rs b/crypto/math/src/traits.rs index 525b50c81..b92441312 100644 --- a/crypto/math/src/traits.rs +++ b/crypto/math/src/traits.rs @@ -5,7 +5,10 @@ use crate::errors::ByteConversionError; /// for getting an element from its byte representation in big-endian or /// little-endian order. pub trait ByteConversion { - /// Returns the byte representation of the element in big-endian order.} + /// Byte length of the big-endian representation. + const BYTE_LEN: usize; + + /// Returns the byte representation of the element in big-endian order. #[cfg(feature = "alloc")] fn to_bytes_be(&self) -> alloc::vec::Vec; @@ -22,6 +25,14 @@ pub trait ByteConversion { fn from_bytes_le(bytes: &[u8]) -> Result where Self: Sized; + + /// Write big-endian bytes into `buf[..BYTE_LEN]`. + /// Override for zero-allocation performance in hot paths. + #[cfg(feature = "alloc")] + fn write_bytes_be(&self, buf: &mut [u8]) { + let bytes = self.to_bytes_be(); + buf[..bytes.len()].copy_from_slice(&bytes); + } } /// Serialize function without args @@ -32,18 +43,6 @@ pub trait AsBytes { fn as_bytes(&self) -> alloc::vec::Vec; } -/// Zero-allocation byte serialization for Merkle tree hashing. -/// -/// Unlike `AsBytes::as_bytes()` which returns `Vec` (heap allocation per call), -/// this trait writes bytes directly into a caller-provided buffer. For Merkle trees -/// with millions of field elements, this eliminates millions of 8-byte allocations. -pub trait WriteBytes { - /// Byte length of this element's big-endian representation. - const BYTE_LEN: usize; - /// Write big-endian bytes into `buf[..BYTE_LEN]`. - fn write_bytes_be(&self, buf: &mut [u8]); -} - #[cfg(feature = "alloc")] impl AsBytes for u32 { fn as_bytes(&self) -> alloc::vec::Vec { @@ -59,6 +58,8 @@ impl AsBytes for u64 { } impl ByteConversion for u64 { + const BYTE_LEN: usize = 8; + #[cfg(feature = "alloc")] fn to_bytes_be(&self) -> alloc::vec::Vec { self.to_be_bytes().to_vec() diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 66e76ca1e..8e59807c1 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -63,8 +63,8 @@ impl< PI, > IsStarkProver for Prover where - FieldElement: math::traits::WriteBytes, - FieldElement: math::traits::WriteBytes, + FieldElement: math::traits::ByteConversion, + FieldElement: math::traits::ByteConversion, { } @@ -355,8 +355,8 @@ pub trait IsStarkProver< FieldExtension: Send + Sync + IsField, PI, > where - FieldElement: math::traits::WriteBytes, - FieldElement: math::traits::WriteBytes, + FieldElement: math::traits::ByteConversion, + FieldElement: math::traits::ByteConversion, { /// Builds a Merkle tree commitment from column-major LDE evaluations with /// bit-reverse permutation, without cloning the full evaluation matrix. @@ -369,10 +369,10 @@ pub trait IsStarkProver< columns: &[Vec>], ) -> Option<(BatchedMerkleTree, Commitment)> where - FieldElement: AsBytes + Sync + Send + math::traits::WriteBytes, + FieldElement: AsBytes + Sync + Send + math::traits::ByteConversion, E: IsField, { - use math::traits::WriteBytes; + use math::traits::ByteConversion; if columns.is_empty() || columns[0].is_empty() { return None; @@ -380,7 +380,7 @@ pub trait IsStarkProver< let num_rows = columns[0].len(); let num_cols = columns.len(); - let byte_len = as WriteBytes>::BYTE_LEN; + let byte_len = as ByteConversion>::BYTE_LEN; debug_assert!( num_rows.is_power_of_two(), @@ -693,9 +693,9 @@ pub trait IsStarkProver< ) -> Option<(BatchedMerkleTree, Commitment)> where FieldElement: AsBytes + Sync + Send, - FieldElement: AsBytes + Sync + Send + math::traits::WriteBytes, + FieldElement: AsBytes + Sync + Send + math::traits::ByteConversion, { - use math::traits::WriteBytes; + use math::traits::ByteConversion; let num_parts = lde_composition_poly_parts_evaluations.len(); if num_parts == 0 { @@ -712,7 +712,7 @@ pub trait IsStarkProver< "num_rows must be a power of two for reverse_index" ); - let byte_len = as WriteBytes>::BYTE_LEN; + let byte_len = as ByteConversion>::BYTE_LEN; // One allocation per leaf (was one per field element): write all parts // into a single buffer. Each leaf = row_pair[2*i] ++ row_pair[2*i+1] after bit-reverse.