Stark: Small refactor of StarkTranscript (lambdaclass#574)

* small refactor of starkprovertranscript * move keccak_hash to StarkProverTranscript impl
PatStiles · Sep 26, 2023 · 361e8ea · 361e8ea
1 parent fe829b0
commit 361e8ea
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 49 deletions.
diff --git a/provers/stark/src/prover.rs b/provers/stark/src/prover.rs
@@ -14,7 +14,7 @@ use rayon::prelude::{IndexedParallelIterator, IntoParallelRefIterator, ParallelI
 
 #[cfg(debug_assertions)]
 use crate::debug::validate_trace;
-use crate::transcript::{sample_z_ood, IsStarkTranscript};
+use crate::transcript::IsStarkTranscript;
 
 use super::config::{BatchedMerkleTree, Commitment};
 use super::constraints::evaluator::ConstraintEvaluator;
@@ -658,10 +658,9 @@ where
     let timer3 = Instant::now();
 
     // <<<< Receive challenge: z
-    let z = sample_z_ood(
+    let z = transcript.sample_z_ood(
         &domain.lde_roots_of_unity_coset,
         &domain.trace_roots_of_unity,
-        &mut transcript,
     );
 
     let round_3_result = round_3_evaluate_polynomials_in_out_of_domain_element(

diff --git a/provers/stark/src/transcript.rs b/provers/stark/src/transcript.rs
@@ -2,7 +2,7 @@ use lambdaworks_math::{
     field::{
         element::FieldElement,
         fields::fft_friendly::stark_252_prime_field::Stark252PrimeField,
-        traits::{IsFFTField, IsField, IsPrimeField},
+        traits::{IsFFTField, IsField},
     },
     traits::ByteConversion,
     unsigned_integer::element::U256,
@@ -15,46 +15,65 @@ pub trait IsStarkTranscript<F: IsField> {
     fn state(&self) -> [u8; 32];
     fn sample_field_element(&mut self) -> FieldElement<F>;
     fn sample_u64(&mut self, upper_bound: u64) -> u64;
+    fn sample_z_ood(
+        &mut self,
+        lde_roots_of_unity_coset: &[FieldElement<F>],
+        trace_roots_of_unity: &[FieldElement<F>],
+    ) -> FieldElement<F>
+    where
+        FieldElement<F>: ByteConversion,
+    {
+        loop {
+            let value: FieldElement<F> = self.sample_field_element();
+            if !lde_roots_of_unity_coset.iter().any(|x| x == &value)
+                && !trace_roots_of_unity.iter().any(|x| x == &value)
+            {
+                return value;
+            }
+        }
+    }
 }
 
-fn keccak_hash(data: &[u8]) -> Keccak256 {
-    let mut hasher = Keccak256::new();
-    hasher.update(data);
-    hasher
-}
-
-const MODULUS_MAX_MULTIPLE: U256 =
-    U256::from_hex_unchecked("f80000000000020f00000000000000000000000000000000000000000000001f");
-const R_INV: U256 =
-    U256::from_hex_unchecked("0x40000000000001100000000000012100000000000000000000000000000000");
-
 pub struct StoneProverTranscript {
-    hash: Keccak256,
+    state: [u8; 32],
     seed_increment: U256,
     counter: u32,
     spare_bytes: Vec<u8>,
 }
 
 impl StoneProverTranscript {
+    const MODULUS_MAX_MULTIPLE: U256 = U256::from_hex_unchecked(
+        "f80000000000020f00000000000000000000000000000000000000000000001f",
+    );
+    const R_INV: U256 = U256::from_hex_unchecked(
+        "0x40000000000001100000000000012100000000000000000000000000000000",
+    );
     pub fn new(public_input_data: &[u8]) -> Self {
-        let hash = keccak_hash(public_input_data);
         StoneProverTranscript {
-            hash,
+            state: Self::keccak_hash(public_input_data),
             seed_increment: U256::from_hex_unchecked("1"),
             counter: 0,
             spare_bytes: vec![],
         }
     }
 
+    fn keccak_hash(data: &[u8]) -> [u8; 32] {
+        let mut hasher = Keccak256::new();
+        hasher.update(data);
+        let mut result_hash = [0_u8; 32];
+        result_hash.copy_from_slice(&hasher.finalize_reset());
+        result_hash
+    }
+
     pub fn sample_block(&mut self, used_bytes: usize) -> Vec<u8> {
-        let mut first_part: Vec<u8> = self.hash.clone().finalize().to_vec();
+        let mut first_part: Vec<u8> = self.state.to_vec();
         let mut counter_bytes: Vec<u8> = vec![0; 28]
             .into_iter()
             .chain(self.counter.to_be_bytes().to_vec())
             .collect();
         self.counter += 1;
         first_part.append(&mut counter_bytes);
-        let block = keccak_hash(&first_part).finalize().to_vec();
+        let block = Self::keccak_hash(&first_part);
         self.spare_bytes.extend(&block[used_bytes..]);
         block[..used_bytes].to_vec()
     }
@@ -100,28 +119,26 @@ impl IsStarkTranscript<Stark252PrimeField> for StoneProverTranscript {
 
     fn append_bytes(&mut self, new_bytes: &[u8]) {
         let mut result_hash = [0_u8; 32];
-        result_hash.copy_from_slice(&self.hash.clone().finalize_reset());
+        result_hash.copy_from_slice(&self.state);
         result_hash.reverse();
 
-        let digest = U256::from_bytes_be(&self.hash.clone().finalize()).unwrap();
+        let digest = U256::from_bytes_be(&self.state).unwrap();
         let new_seed = (digest + self.seed_increment).to_bytes_be();
-        self.hash = keccak_hash(&[&new_seed, new_bytes].concat());
+        self.state = Self::keccak_hash(&[&new_seed, new_bytes].concat());
         self.counter = 0;
         self.spare_bytes.clear();
     }
 
     fn state(&self) -> [u8; 32] {
-        let mut state = [0u8; 32];
-        state.copy_from_slice(&self.hash.clone().finalize());
-        state
+        self.state
     }
 
     fn sample_field_element(&mut self) -> FieldElement<Stark252PrimeField> {
         let mut result = self.sample_big_int();
-        while result >= MODULUS_MAX_MULTIPLE {
+        while result >= Self::MODULUS_MAX_MULTIPLE {
             result = self.sample_big_int();
         }
-        FieldElement::new(result) * FieldElement::new(R_INV)
+        FieldElement::new(result) * FieldElement::new(Self::R_INV)
     }
 
     fn sample_u64(&mut self, upper_bound: u64) -> u64 {
@@ -133,24 +150,6 @@ impl IsStarkTranscript<Stark252PrimeField> for StoneProverTranscript {
     }
 }
 
-pub fn sample_z_ood<F: IsPrimeField>(
-    lde_roots_of_unity_coset: &[FieldElement<F>],
-    trace_roots_of_unity: &[FieldElement<F>],
-    transcript: &mut impl IsStarkTranscript<F>,
-) -> FieldElement<F>
-where
-    FieldElement<F>: ByteConversion,
-{
-    loop {
-        let value: FieldElement<F> = transcript.sample_field_element();
-        if !lde_roots_of_unity_coset.iter().any(|x| x == &value)
-            && !trace_roots_of_unity.iter().any(|x| x == &value)
-        {
-            return value;
-        }
-    }
-}
-
 pub fn batch_sample_challenges<F: IsFFTField>(
     size: usize,
     transcript: &mut impl IsStarkTranscript<F>,

diff --git a/provers/stark/src/verifier.rs b/provers/stark/src/verifier.rs
@@ -22,7 +22,7 @@ use super::{
     grinding::hash_transcript_with_int_and_get_leading_zeros,
     proof::{options::ProofOptions, stark::StarkProof},
     traits::AIR,
-    transcript::{batch_sample_challenges, sample_z_ood},
+    transcript::batch_sample_challenges,
 };
 
 struct Challenges<F, A>
@@ -90,10 +90,9 @@ where
     // ===================================
 
     // >>>> Send challenge: z
-    let z = sample_z_ood(
+    let z = transcript.sample_z_ood(
         &domain.lde_roots_of_unity_coset,
         &domain.trace_roots_of_unity,
-        transcript,
     );
 
     // <<<< Receive value: H₁(z²)