Include the FRI prover's PoW witness in the transcript

We don't think this is required for soundness, but just to remove any doubt.

Old protocol:
```
...
P sends final_poly
V samples random r
P sends pow_witness (not in transcript)
V computes pow_response = H(r, pow_witness)
V asserts pow_response has N leading 0s
...
```

New protocol:
```
...
P sends final_poly
P sends pow_witness
V samples random pow_response
V asserts pow_response has N leading 0s
...
```

plonky2/src/fri/challenges.rs

@@ -49,16 +49,8 @@ impl<F: RichField, H: Hasher<F>> Challenger<F, H> {
 
         self.observe_extension_elements(&final_poly.coeffs);
 
-        let fri_pow_response = C::InnerHasher::hash_no_pad(
+        self.observe_element(pow_witness);
-            &self
+        let fri_pow_response = self.get_challenge();
-                .get_hash()
-                .elements
-                .iter()
-                .copied()
-                .chain(Some(pow_witness))
-                .collect::<Vec<_>>(),
-        )
-        .elements[0];
 
         let fri_query_indices = (0..num_fri_queries)
             .map(|_| self.get_challenge().to_canonical_u64() as usize % lde_size)
@@ -105,16 +97,8 @@ impl<F: RichField + Extendable<D>, H: AlgebraicHasher<F>, const D: usize>
 
         self.observe_extension_elements(&final_poly.0);
 
-        let pow_inputs = self
+        self.observe_element(pow_witness);
-            .get_hash(builder)
+        let fri_pow_response = self.get_challenge(builder);
-            .elements
-            .iter()
-            .copied()
-            .chain(Some(pow_witness))
-            .collect();
-        let fri_pow_response = builder
-            .hash_n_to_hash_no_pad::<C::InnerHasher>(pow_inputs)
-            .elements[0];
 
         let fri_query_indices = (0..num_fri_queries)
             .map(|_| self.get_challenge(builder))

plonky2/src/fri/prover.rs

@@ -1,4 +1,3 @@
-use itertools::Itertools;
 use maybe_rayon::*;
 use plonky2_field::extension::{flatten, unflatten, Extendable};
 use plonky2_field::polynomial::{PolynomialCoeffs, PolynomialValues};
@@ -6,7 +5,8 @@ use plonky2_util::reverse_index_bits_in_place;
 
 use crate::fri::proof::{FriInitialTreeProof, FriProof, FriQueryRound, FriQueryStep};
 use crate::fri::{FriConfig, FriParams};
-use crate::hash::hash_types::{HashOut, RichField};
+use crate::hash::hash_types::RichField;
+use crate::hash::hashing::{PlonkyPermutation, SPONGE_RATE};
 use crate::hash::merkle_tree::MerkleTree;
 use crate::iop::challenger::Challenger;
 use crate::plonk::config::{GenericConfig, Hasher};
@@ -44,11 +44,10 @@ where
     );
 
     // PoW phase
-    let current_hash = challenger.get_hash();
     let pow_witness = timed!(
         timing,
         "find proof-of-work witness",
-        fri_proof_of_work::<F, C, D>(current_hash, &fri_params.config)
+        fri_proof_of_work::<F, C, D>(challenger, &fri_params.config)
     );
 
     // Query phase
@@ -114,28 +113,56 @@ where
     (trees, coeffs)
 }
 
+/// Performs the proof-of-work (a.k.a. grinding) step of the FRI protocol. Returns the PoW witness.
 fn fri_proof_of_work<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>(
-    current_hash: HashOut<F>,
+    challenger: &mut Challenger<F, C::Hasher>,
     config: &FriConfig,
 ) -> F {
-    (0..=F::NEG_ONE.to_canonical_u64())
+    // let pow_seed = challenger.get_hash();
+    let min_leading_zeros = config.proof_of_work_bits + (64 - F::order().bits()) as u32;
+
+    // The easiest implementation would be repeatedly clone our Challenger. With each clone, we'd
+    // observe an incrementing PoW witness, then get the PoW response. If it contained sufficient
+    // leading zeros, we'd end the search, and store this clone as our new challenger.
+    //
+    // However, performance is critical here. We want to avoid cloning Challenger, particularly
+    // since it stores vectors, which means allocations. We'd like a more compact state to clone.
+    //
+    // We know that a duplex will be performed right after we send the PoW witness, so we can ignore
+    // any output_buffer, which will be invalidated. We also know input_buffer.len() < SPONGE_WIDTH,
+    // an invariant of Challenger.
+    //
+    // We separate the duplex operation into two steps, one which can be performed now, and the
+    // other which depends on the PoW witness candidate. The first step is the overwrite our sponge
+    // state with any inputs (excluding the PoW witness candidate). The second step is to overwrite
+    // one more element of our sponge state with the candidate, then apply the permutation,
+    // obtaining our duplex's post-state which contains the PoW response.
+    let mut duplex_intermediate_state = challenger.sponge_state;
+    let witness_input_pos = challenger.input_buffer.len();
+    for (i, input) in challenger.input_buffer.iter().enumerate() {
+        duplex_intermediate_state[i] = *input;
+    }
+
+    let pow_witness = (0..=F::NEG_ONE.to_canonical_u64())
         .into_par_iter()
-        .find_any(|&i| {
+        .find_any(|&candidate| {
-            C::InnerHasher::hash_no_pad(
+            let mut duplex_state = duplex_intermediate_state;
-                &current_hash
+            duplex_state[witness_input_pos] = F::from_canonical_u64(candidate);
-                    .elements
+            duplex_state =
-                    .iter()
+                <<C as GenericConfig<D>>::Hasher as Hasher<F>>::Permutation::permute(duplex_state);
-                    .copied()
+            let pow_response = duplex_state[SPONGE_RATE - 1];
-                    .chain(Some(F::from_canonical_u64(i)))
+            let leading_zeros = pow_response.to_canonical_u64().leading_zeros();
-                    .collect_vec(),
+            leading_zeros >= min_leading_zeros
-            )
-            .elements[0]
-                .to_canonical_u64()
-                .leading_zeros()
-                >= config.proof_of_work_bits + (64 - F::order().bits()) as u32
         })
         .map(F::from_canonical_u64)
-        .expect("Proof of work failed. This is highly unlikely!")
+        .expect("Proof of work failed. This is highly unlikely!");
+
+    // Recompute pow_response using our normal Challenger code, and make sure it matches.
+    challenger.observe_element(pow_witness);
+    let pow_response = challenger.get_challenge();
+    let leading_zeros = pow_response.to_canonical_u64().leading_zeros();
+    assert!(leading_zeros >= min_leading_zeros);
+    pow_witness
 }
 
 fn fri_prover_query_rounds<

plonky2/src/iop/challenger.rs

@@ -15,8 +15,8 @@ use crate::plonk::config::{AlgebraicHasher, GenericHashOut, Hasher};
 /// Observes prover messages, and generates challenges by hashing the transcript, a la Fiat-Shamir.
 #[derive(Clone)]
 pub struct Challenger<F: RichField, H: Hasher<F>> {
-    sponge_state: [F; SPONGE_WIDTH],
+    pub(crate) sponge_state: [F; SPONGE_WIDTH],
-    input_buffer: Vec<F>,
+    pub(crate) input_buffer: Vec<F>,
     output_buffer: Vec<F>,
     _phantom: PhantomData<H>,
 }