improve challenge generation

.gas-snapshot

@@ -1,2 +1,2 @@
-LightClient_newFinalizedState_Test:testCorrectUpdateBench() (gas: 656742)
-PlonkVerifier_verify_Test:test_verify_succeeds() (gas: 491667)
+LightClient_newFinalizedState_Test:testCorrectUpdateBench() (gas: 620982)
+PlonkVerifier_verify_Test:test_verify_succeeds() (gas: 455622)

contracts/rust/adapter/src/jellyfish.rs

@@ -13,6 +13,7 @@ use ethers::{
     utils::hex::ToHex,
 };
 use jf_pcs::prelude::Commitment;
+use jf_plonk::constants::KECCAK256_STATE_SIZE;
 use jf_plonk::proof_system::structs::{OpenKey, Proof, ProofEvaluations, VerifyingKey};
 use jf_plonk::testing_apis::Challenges;
 use jf_plonk::transcript::SolidityTranscript;
@@ -95,6 +96,7 @@ pub fn open_key() -> OpenKey<Bn254> {
 /// an intermediate representation of the transcript parsed from abi.encode(transcript) from Solidity.
 #[derive(Clone, EthAbiType, EthAbiCodec)]
 pub struct ParsedTranscript {
+    pub(crate) state: H256,
     pub(crate) transcript: Bytes,
 }
 
@@ -108,16 +110,19 @@ impl FromStr for ParsedTranscript {
 
 impl From<SolidityTranscript> for ParsedTranscript {
     fn from(t: SolidityTranscript) -> Self {
-        let transcript = t.internal();
+        let (state, transcript) = t.internal();
         Self {
+            state: H256::from_slice(&state),
             transcript: transcript.into(),
         }
     }
 }
 
 impl From<ParsedTranscript> for SolidityTranscript {
     fn from(t: ParsedTranscript) -> Self {
-        Self::from_internal(t.transcript.to_vec())
+        let mut state = [0u8; KECCAK256_STATE_SIZE];
+        state.copy_from_slice(&t.state.to_fixed_bytes());
+        Self::from_internal(state, t.transcript.to_vec())
     }
 }
 

contracts/rust/diff-test/src/main.rs

@@ -172,7 +172,7 @@ fn main() {
             let field = u256_to_field::<Fr>(cli.args[1].parse::<U256>().unwrap());
 
             let mut t: SolidityTranscript = t_parsed.into();
-            t.append_challenge::<Bn254>(&[], &field).unwrap();
+            t.append_field_elem::<Bn254>(&[], &field).unwrap();
             let res: ParsedTranscript = t.into();
             println!("{}", (res,).encode_hex());
         }
@@ -198,7 +198,7 @@ fn main() {
             let t_parsed = cli.args[0].parse::<ParsedTranscript>().unwrap();
 
             let mut t: SolidityTranscript = t_parsed.into();
-            let chal = t.get_and_append_challenge::<Bn254>(&[]).unwrap();
+            let chal = t.get_challenge::<Bn254>(&[]).unwrap();
 
             let updated_t: ParsedTranscript = t.into();
             let res = (updated_t, field_to_u256(chal));
@@ -233,7 +233,7 @@ fn main() {
             let proof: Proof<Bn254> = proof_parsed.clone().into();
 
             let mut t: SolidityTranscript = t_parsed.into();
-            <SolidityTranscript as PlonkTranscript<Fr>>::append_proof_evaluations::<Bn254>(
+            <SolidityTranscript as PlonkTranscript<Fq>>::append_proof_evaluations::<Bn254>(
                 &mut t,
                 &proof.poly_evals,
             )

contracts/src/libraries/PlonkVerifier.sol

@@ -193,13 +193,13 @@ library PlonkVerifier {
             transcript.transcript, proof.wire3.x, proof.wire3.y, proof.wire4.x, proof.wire4.y
         );
 
-        res.beta = transcript.getAndAppendChallenge();
-        res.gamma = transcript.getAndAppendChallenge();
+        res.beta = transcript.getChallenge();
+        res.gamma = transcript.getChallenge();
 
         transcript.transcript =
             abi.encodePacked(transcript.transcript, proof.prodPerm.x, proof.prodPerm.y);
 
-        res.alpha = transcript.getAndAppendChallenge();
+        res.alpha = transcript.getChallenge();
 
         transcript.transcript = abi.encodePacked(
             transcript.transcript,
@@ -214,7 +214,7 @@ library PlonkVerifier {
             transcript.transcript, proof.split3.x, proof.split3.y, proof.split4.x, proof.split4.y
         );
 
-        res.zeta = transcript.getAndAppendChallenge();
+        res.zeta = transcript.getChallenge();
 
         // Append proof evaluations
         transcript.transcript = abi.encodePacked(
@@ -235,13 +235,13 @@ library PlonkVerifier {
             proof.prodPermZetaOmegaEval
         );
 
-        res.v = transcript.getAndAppendChallenge();
+        res.v = transcript.getChallenge();
 
         transcript.transcript = abi.encodePacked(
             transcript.transcript, proof.zeta.x, proof.zeta.y, proof.zetaOmega.x, proof.zetaOmega.y
         );
 
-        res.u = transcript.getAndAppendChallenge();
+        res.u = transcript.getChallenge();
 
         assembly {
             let alpha := mload(res)

contracts/src/libraries/Transcript.sol

@@ -7,6 +7,7 @@ import { IPlonkVerifier } from "../interfaces/IPlonkVerifier.sol";
 
 library Transcript {
     struct TranscriptData {
+        bytes32 state;
         bytes transcript;
     }
 
@@ -40,29 +41,39 @@ library Transcript {
         self.transcript = abi.encodePacked(self.transcript, comm.x, comm.y);
     }
 
-    function getAndAppendChallenge(TranscriptData memory self)
-        internal
-        pure
-        returns (uint256 ret)
-    {
-        bytes memory transcript = self.transcript;
+    // 1. state = hash(state | transcript)
+    // 2. transcript = Vec::new()
+    // 3. challenge = bytes_to_field(state)
+    //
+    // Note: every challenge generation is implicitly domain-separated, thus it's safe
+    // to call it multiple times in a row (e.g. multiple challenges in a single round)
+    // to get multiple different challenges.
+    function getChallenge(TranscriptData memory self) internal pure returns (uint256 ret) {
+        // memory layout:
+        // offset 0x00: state
+        // offset 0x20: pointer to transcript (e.g. point to 0x1a0)
+        // (e.g.0x1a0): length of transcript
+        //      0x1c0 : transcript bytes
         uint256 p = BN254.R_MOD;
 
+        // Instead of using free memory for scratch pad, we do the following trick:
+        // 1. overwrite offset 0x40 with current state
+        // 2. compute keccak of "state | transcript" (continuous in memory)
+        // 3. store the hash to offset 0x00 and update the length of transcript to 0
+        // 4. compute challenge from current state
         assembly {
-            let len := mload(transcript)
-            let newLen := add(len, 32)
-            let dataPtr := add(transcript, 0x20)
-
-            // uint256(keccak256(transcript)) % BN254.R_MOD;
-            ret := mod(keccak256(dataPtr, len), p)
-
-            // same as self.transcript = abi.encodePacked(self.transcript, ret);
-            mstore(transcript, newLen)
-            mstore(add(dataPtr, len), ret)
-
-            // update free memory pointer since we extend the dynamic array
-            // to prevent potential overwrite
-            mstore(0x40, add(mload(0x40), 32))
+            let lenPtr := mload(add(self, 0x20))
+            // step 1
+            let len := mload(lenPtr)
+            mstore(lenPtr, mload(self))
+
+            // step 2
+            let hash := keccak256(lenPtr, add(len, 0x20))
+            // step 3
+            mstore(self, hash)
+            mstore(lenPtr, 0)
+            // step 4
+            ret := mod(hash, p)
         }
     }
 
@@ -80,7 +91,8 @@ library Transcript {
             self.transcript,
             uint32(sizeInBits),
             uint64(verifyingKey.domainSize),
-            uint64(verifyingKey.numInputs)
+            uint64(verifyingKey.numInputs),
+            bytes12(0) // padding to align with word size
         );
 
         // G2 point from KZG SRS

contracts/test/Transcript.t.sol

@@ -107,11 +107,11 @@ contract Transcript_appendGroupElement_Test is Test {
     }
 }
 
-contract Transcript_getAndAppendChallenge_Test is Test {
+contract Transcript_getChallenge_Test is Test {
     using T for T.TranscriptData;
 
-    /// @dev Test if `getAndAppendChallenge` matches that of Jellyfish
-    function testFuzz_getAndAppendChallenge_matches(T.TranscriptData memory transcript) external {
+    /// @dev Test if `getChallenge` matches that of Jellyfish
+    function testFuzz_getChallenge_matches(T.TranscriptData memory transcript) external {
         string[] memory cmds = new string[](3);
         cmds[0] = "diff-test";
         cmds[1] = "transcript-get-chal";
@@ -121,7 +121,7 @@ contract Transcript_getAndAppendChallenge_Test is Test {
         (T.TranscriptData memory updated, uint256 chal) =
             abi.decode(result, (T.TranscriptData, uint256));
 
-        uint256 challenge = transcript.getAndAppendChallenge();
+        uint256 challenge = transcript.getChallenge();
 
         assertEq(updated.transcript, transcript.transcript);
         assertEq(chal, challenge);