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/math/src/field/extensions_goldilocks.rs b/crypto/math/src/field/extensions_goldilocks.rs index 44e0d3f85..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]); @@ -554,3 +569,27 @@ impl HasDefaultTranscript for Degree3GoldilocksExtensionField { fn mul_by_7(a: &FpE) -> FpE { FpE::from_raw(mul_by_7_raw(*a.value())) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::traits::ByteConversion; + + #[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()); + } + } +} diff --git a/crypto/math/src/field/goldilocks.rs b/crypto/math/src/field/goldilocks.rs index 42cd889e8..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() @@ -543,3 +551,32 @@ impl HasDefaultTranscript for GoldilocksField { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::traits::ByteConversion; + + #[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/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 ed4d1589b..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 @@ -47,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 41d27fc66..8e59807c1 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; @@ -63,6 +62,9 @@ impl< FieldExtension: Send + Sync + IsField, PI, > IsStarkProver for Prover +where + FieldElement: math::traits::ByteConversion, + FieldElement: math::traits::ByteConversion, { } @@ -352,22 +354,10 @@ pub trait IsStarkProver< Field: IsSubFieldOf + IsFFTField + Send + Sync, FieldExtension: Send + Sync + IsField, PI, -> +> where + FieldElement: math::traits::ByteConversion, + FieldElement: math::traits::ByteConversion, { - /// 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. /// @@ -379,28 +369,41 @@ pub trait IsStarkProver< columns: &[Vec>], ) -> Option<(BatchedMerkleTree, Commitment)> where - FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send + math::traits::ByteConversion, E: IsField, { + use math::traits::ByteConversion; + 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 ByteConversion>::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"))] let iter = 0..num_rows; + // 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); - 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]); + } + BatchedMerkleTreeBackend::::hash_bytes(&buf) }) .collect(); @@ -690,8 +693,10 @@ pub trait IsStarkProver< ) -> Option<(BatchedMerkleTree, Commitment)> where FieldElement: AsBytes + Sync + Send, - FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send + math::traits::ByteConversion, { + use math::traits::ByteConversion; + let num_parts = lde_composition_poly_parts_evaluations.len(); if num_parts == 0 { return None; @@ -702,18 +707,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. - // Each leaf = row_pair[2*i] ++ row_pair[2*i+1] after bit-reverse. - // Build the merged leaf data and hash in one pass. + 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. #[cfg(feature = "parallel")] let iter = (0..num_leaves).into_par_iter(); #[cfg(not(feature = "parallel"))] @@ -723,14 +725,18 @@ 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) + BatchedMerkleTreeBackend::::hash_bytes(&buf) }) .collect();