Merge branch 'main' into perf/batch-byte-range-checks

Author: diegokingston

crypto/crypto/src/merkle_tree/backends/field_element_vector.rs

@@ -71,6 +71,23 @@ impl<F, D: Digest, const NUM_BYTES: usize> Default for FieldElementVectorBackend
     }
 }
 
+impl<F, D: Digest, const NUM_BYTES: usize> FieldElementVectorBackend<F, D, NUM_BYTES>
+where
+    [u8; NUM_BYTES]: From<Output<D>>,
+{
+    /// 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<F, D: Digest, const NUM_BYTES: usize> IsMerkleTreeBackend
     for FieldElementVectorBackend<F, D, NUM_BYTES>
 where

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<u8> {
         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<u8> {
         let mut bytes = ByteConversion::to_bytes_be(&self[2]);
@@ -470,6 +474,17 @@ impl Fp3E {
 // =====================================================
 
 impl ByteConversion for FieldElement<Degree3GoldilocksExtensionField> {
+    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<u8> {
         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::<Degree3GoldilocksExtensionField>::zero(),
+            FieldElement::<Degree3GoldilocksExtensionField>::one(),
+            FieldElement::<Degree3GoldilocksExtensionField>::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());
+        }
+    }
+}

crypto/math/src/field/goldilocks.rs

@@ -434,6 +434,14 @@ impl GoldilocksElement {
 // =====================================================
 
 impl ByteConversion for FieldElement<GoldilocksField> {
+    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<u8> {
         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::<GoldilocksField>::from(0u64),
+            FieldElement::<GoldilocksField>::from(1u64),
+            FieldElement::<GoldilocksField>::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::<GoldilocksField>::from_raw(GOLDILOCKS_PRIME + 5);
+        let mut buf = [0u8; 8];
+        elem.write_bytes_be(&mut buf);
+        assert_eq!(&buf[..], elem.as_bytes().as_slice());
+    }
+}

crypto/math/src/field/test_fields/u32_test_field.rs

@@ -11,6 +11,8 @@ use crate::traits::ByteConversion;
 pub struct U32Field<const MODULUS: u32>;
 
 impl ByteConversion for u32 {
+    const BYTE_LEN: usize = 4;
+
     #[cfg(feature = "alloc")]
     fn to_bytes_be(&self) -> alloc::vec::Vec<u8> {
         unimplemented!()

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<u8>;
 
@@ -22,6 +25,14 @@ pub trait ByteConversion {
     fn from_bytes_le(bytes: &[u8]) -> Result<Self, ByteConversionError>
     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<u8> {
         self.to_be_bytes().to_vec()

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<Field, FieldExtension, PI> for Prover<Field, FieldExtension, PI>
+where
+    FieldElement<Field>: math::traits::ByteConversion,
+    FieldElement<FieldExtension>: math::traits::ByteConversion,
 {
 }
 
@@ -352,22 +354,10 @@ pub trait IsStarkProver<
     Field: IsSubFieldOf<FieldExtension> + IsFFTField + Send + Sync,
     FieldExtension: Send + Sync + IsField,
     PI,
->
+> where
+    FieldElement<Field>: math::traits::ByteConversion,
+    FieldElement<FieldExtension>: math::traits::ByteConversion,
 {
-    /// Returns the Merkle tree and the commitment to the vectors `vectors`.
-    fn batch_commit_extension(
-        vectors: &[Vec<FieldElement<FieldExtension>>],
-    ) -> Option<(BatchedMerkleTree<FieldExtension>, Commitment)>
-    where
-        FieldElement<Field>: AsBytes + Sync + Send,
-        FieldElement<FieldExtension>: 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<FieldElement<E>>],
     ) -> Option<(BatchedMerkleTree<E>, Commitment)>
     where
-        FieldElement<E>: AsBytes + Sync + Send,
+        FieldElement<E>: 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 = <FieldElement<E> 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<Commitment> = iter
             .map(|row_idx| {
                 let br_idx = reverse_index(row_idx, num_rows as u64);
-                let row: Vec<FieldElement<E>> = (0..num_cols)
-                    .map(|col_idx| columns[col_idx][br_idx].clone())
-                    .collect();
-                BatchedMerkleTreeBackend::<E>::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::<E>::hash_bytes(&buf)
             })
             .collect();
 
@@ -690,8 +693,10 @@ pub trait IsStarkProver<
     ) -> Option<(BatchedMerkleTree<FieldExtension>, Commitment)>
     where
         FieldElement<Field>: AsBytes + Sync + Send,
-        FieldElement<FieldExtension>: AsBytes + Sync + Send,
+        FieldElement<FieldExtension>: 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 = <FieldElement<FieldExtension> 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::<FieldExtension>::hash_data(&leaf)
+                BatchedMerkleTreeBackend::<FieldExtension>::hash_bytes(&buf)
             })
             .collect();