refac: enable TranscriptRead|Writer to choose transcript writing

This is needed by improving PCS interface.
kroma-network · Jan 10, 2024 · 5efe19b · 5efe19b
1 parent b3005b8
commit 5efe19b
Show file tree

Hide file tree

Showing 10 changed files with 108 additions and 51 deletions.
diff --git a/tachyon/crypto/commitments/fri/fri.h b/tachyon/crypto/commitments/fri/fri.h
@@ -52,7 +52,8 @@ class FRI final
     Evals evals = domain_->FFT(poly);
     F root;
     if (!tree.Commit(evals.evaluations(), &root)) return false;
-    if (!writer->WriteToProof(root)) return false;
+    if (!writer->template WriteToProof</*NeedToWriteToTranscript=*/true>(root))
+      return false;
     const Poly* cur_poly = &poly;
 
     F beta;
@@ -67,7 +68,9 @@ class FRI final
                                                    hasher_);
         evals = sub_domains_[i - 1]->FFT(folded_poly);
         if (!tree.Commit(evals.evaluations(), &root)) return false;
-        if (!writer->WriteToProof(root)) return false;
+        if (!writer->template WriteToProof</*NeedToWriteToTranscript=*/true>(
+                root))
+          return false;
         cur_poly = &folded_poly;
       }
     }
@@ -76,7 +79,8 @@ class FRI final
     folded_poly = cur_poly->template Fold<false>(beta);
     const F* constant = folded_poly[0];
     root = constant ? *constant : F::Zero();
-    return writer->WriteToProof(root);
+    return writer->template WriteToProof</*NeedToWriteToTranscript=*/true>(
+        root);
   }
 
   [[nodiscard]] bool DoCreateOpeningProof(size_t index,
@@ -126,7 +130,9 @@ class FRI final
       BinaryMerkleTreeStorage<F>* layer = storage_->GetLayer(i);
       BinaryMerkleTree<F, F, MaxDegree + 1> tree(layer, hasher_);
 
-      if (!reader->ReadFromProof(&root)) return false;
+      if (!reader->template ReadFromProof</*NeedToWriteToTranscript=*/true>(
+              &root))
+        return false;
       if (!tree.VerifyOpeningProof(root, proof.evaluations[i], proof.paths[i]))
         return false;
 
@@ -185,7 +191,9 @@ class FRI final
     evaluation += evaluation_sym;
     evaluation *= two_inv;
 
-    if (!reader->ReadFromProof(&root)) return false;
+    if (!reader->template ReadFromProof</*NeedToWriteToTranscript=*/true>(
+            &root))
+      return false;
     if (root != evaluation) {
       LOG(ERROR) << "Root doesn't match with expected evaluation";
       return false;

diff --git a/tachyon/crypto/commitments/kzg/shplonk.h b/tachyon/crypto/commitments/kzg/shplonk.h
@@ -107,7 +107,8 @@ class SHPlonk final : public UnivariatePolynomialCommitmentScheme<
     Commitment h;
     if (!this->Commit(h_poly, &h)) return false;
 
-    if (!writer->WriteToProof(h)) return false;
+    if (!writer->template WriteToProof</*NeedToWriteToTranscript=*/true>(h))
+      return false;
     Field u = writer->SqueezeChallenge();
 
     // Create [L₀(X), L₁(X), L₂(X)].
@@ -190,7 +191,7 @@ class SHPlonk final : public UnivariatePolynomialCommitmentScheme<
     // Commit Q(X)
     Commitment q;
     if (!this->Commit(q_poly, &q)) return false;
-    return writer->WriteToProof(q);
+    return writer->template WriteToProof</*NeedToWriteToTranscript=*/true>(q);
   }
 
   template <typename Container>
@@ -203,12 +204,14 @@ class SHPlonk final : public UnivariatePolynomialCommitmentScheme<
     Field v = reader->SqueezeChallenge();
 
     Commitment h;
-    if (!reader->ReadFromProof(&h)) return false;
+    if (!reader->template ReadFromProof</*NeedToWriteToTranscript=*/true>(&h))
+      return false;
 
     Field u = reader->SqueezeChallenge();
 
     Commitment q;
-    if (!reader->ReadFromProof(&q)) return false;
+    if (!reader->template ReadFromProof</*NeedToWriteToTranscript=*/true>(&q))
+      return false;
 
     PolynomialOpeningGrouper<Poly, Commitment> grouper;
     grouper.GroupByPolyOracleAndPoints(poly_openings);

diff --git a/tachyon/crypto/transcripts/transcript.h b/tachyon/crypto/transcripts/transcript.h
@@ -35,13 +35,11 @@ class TranscriptImpl<_Commitment, false> {
   // Squeeze an encoded verifier challenge from the transcript.
   virtual Field SqueezeChallenge() = 0;
 
-  // Write a |commitment| to the transcript without writing it to the proof,
-  // treating it as a common input.
+  // Write a |commitment| to the transcript.
   [[nodiscard]] virtual bool WriteToTranscript(
       const Commitment& commitment) = 0;
 
-  // Write a |value| to the transcript without writing it to the proof,
-  // treating it as a common input.
+  // Write a |value| to the transcript.
   [[nodiscard]] virtual bool WriteToTranscript(const Field& value) = 0;
 
   TranscriptWriterImpl<Commitment, false>* ToWriter() {
@@ -63,8 +61,7 @@ class TranscriptImpl<_Field, true> {
   // Squeeze an encoded verifier challenge from the transcript.
   virtual Field SqueezeChallenge() = 0;
 
-  // Write a |value| to the transcript without writing it to the proof,
-  // treating it as a common input.
+  // Write a |value| to the transcript.
   [[nodiscard]] virtual bool WriteToTranscript(const Field& value) = 0;
 
   TranscriptWriterImpl<Field, true>* ToWriter() {
@@ -96,15 +93,28 @@ class TranscriptReaderImpl<Commitment, false> : public Transcript<Commitment> {
   const base::Buffer& buffer() const { return buffer_; }
 
   // Read a |commitment| from the proof. Note that it also writes the
-  // |commitment| to the transcript by calling |WriteToTranscript()| internally.
+  // |commitment| to the transcript by calling |WriteToTranscript()| internally
+  // when |NeedToWriteToTranscript| is set to true.
+  template <bool NeedToWriteToTranscript>
   [[nodiscard]] bool ReadFromProof(Commitment* commitment) {
-    return DoReadFromProof(commitment) && this->WriteToTranscript(*commitment);
+    if constexpr (NeedToWriteToTranscript) {
+      return DoReadFromProof(commitment) &&
+             this->WriteToTranscript(*commitment);
+    } else {
+      return DoReadFromProof(commitment);
+    }
   }
 
   // Read a |value| from the proof. Note that it also writes the
-  // |value| to the transcript by calling |WriteToTranscript()| internally.
+  // |value| to the transcript by calling |WriteToTranscript()| internally when
+  // |NeedToWriteToTranscript| is set to true.
+  template <bool NeedToWriteToTranscript>
   [[nodiscard]] bool ReadFromProof(Field* value) {
-    return DoReadFromProof(value) && this->WriteToTranscript(*value);
+    if constexpr (NeedToWriteToTranscript) {
+      return DoReadFromProof(value) && this->WriteToTranscript(*value);
+    } else {
+      return DoReadFromProof(value);
+    }
   }
 
  protected:
@@ -129,9 +139,15 @@ class TranscriptReaderImpl<Field, true> : public Transcript<Field> {
   const base::Buffer& buffer() const { return buffer_; }
 
   // Read a |value| from the proof. Note that it also writes the
-  // |value| to the transcript by calling |WriteToTranscript()| internally.
+  // |value| to the transcript by calling |WriteToTranscript()| internally when
+  // |NeedToWriteToTranscript| is set to true.
+  template <bool NeedToWriteToTranscript>
   [[nodiscard]] bool ReadFromProof(Field* value) {
-    return DoReadFromProof(value) && this->WriteToTranscript(*value);
+    if constexpr (NeedToWriteToTranscript) {
+      return DoReadFromProof(value) && this->WriteToTranscript(*value);
+    } else {
+      return DoReadFromProof(value);
+    }
   }
 
  protected:
@@ -164,15 +180,27 @@ class TranscriptWriterImpl<Commitment, false> : public Transcript<Commitment> {
   base::Uint8VectorBuffer&& TakeBuffer() && { return std::move(buffer_); }
 
   // Write a |commitment| to the proof. Note that it also writes the
-  // |commitment| to the transcript by calling |WriteToTranscript()| internally.
+  // |commitment| to the transcript by calling |WriteToTranscript()| internally
+  // when |NeedToWriteToTranscript| is set to true.
+  template <bool NeedToWriteToTranscript>
   [[nodiscard]] bool WriteToProof(const Commitment& commitment) {
-    return this->WriteToTranscript(commitment) && DoWriteToProof(commitment);
+    if constexpr (NeedToWriteToTranscript) {
+      return this->WriteToTranscript(commitment) && DoWriteToProof(commitment);
+    } else {
+      return DoWriteToProof(commitment);
+    }
   }
 
   // Write a |value| to the proof. Note that it also writes the
-  // |value| to the transcript by calling |WriteToTranscript()| internally.
+  // |value| to the transcript by calling |WriteToTranscript()| internally when
+  // |NeedToWriteToTranscript| is set to true.
+  template <bool NeedToWriteToTranscript>
   [[nodiscard]] bool WriteToProof(const Field& value) {
-    return this->WriteToTranscript(value) && DoWriteToProof(value);
+    if constexpr (NeedToWriteToTranscript) {
+      return this->WriteToTranscript(value) && DoWriteToProof(value);
+    } else {
+      return DoWriteToProof(value);
+    }
   }
 
  protected:
@@ -201,9 +229,15 @@ class TranscriptWriterImpl<Field, true> : public Transcript<Field> {
   base::Uint8VectorBuffer&& TakeBuffer() && { return std::move(buffer_); }
 
   // Write a |value| to the proof. Note that it also writes the
-  // |value| to the transcript by calling |WriteToTranscript()| internally.
+  // |value| to the transcript by calling |WriteToTranscript()| internally when
+  // |NeedToWriteToTranscript| is set to true.
+  template <bool NeedToWriteToTranscript>
   [[nodiscard]] bool WriteToProof(const Field& value) {
-    return this->WriteToTranscript(value) && DoWriteToProof(value);
+    if constexpr (NeedToWriteToTranscript) {
+      return this->WriteToTranscript(value) && DoWriteToProof(value);
+    } else {
+      return DoWriteToProof(value);
+    }
   }
 
  protected:

diff --git a/tachyon/zk/base/entities/prover_base.h b/tachyon/zk/base/entities/prover_base.h
@@ -44,22 +44,25 @@ class ProverBase : public Entity<PCS> {
   [[nodiscard]] bool Commit(const Poly& poly) {
     Commitment commitment;
     if (!this->pcs_.Commit(poly, &commitment)) return false;
-    return GetWriter()->WriteToProof(commitment);
+    return GetWriter()->template WriteToProof</*NeedToWriteToTranscript=*/true>(
+        commitment);
   }
 
   template <typename Container>
   [[nodiscard]] bool Commit(const Container& coeffs) {
     Commitment commitment;
     if (!this->pcs_.DoCommit(coeffs, &commitment)) return false;
-    return GetWriter()->WriteToProof(commitment);
+    return GetWriter()->template WriteToProof</*NeedToWriteToTranscript=*/true>(
+        commitment);
   }
 
   [[nodiscard]] bool CommitEvals(const Evals& evals) {
     if (evals.NumElements() != this->domain_->size()) return false;
 
     Commitment commitment;
     if (!this->pcs_.CommitLagrange(evals, &commitment)) return false;
-    return GetWriter()->WriteToProof(commitment);
+    return GetWriter()->template WriteToProof</*NeedToWriteToTranscript=*/true>(
+        commitment);
   }
 
   [[nodiscard]] bool CommitEvalsWithBlind(const Evals& evals,
@@ -71,7 +74,8 @@ class ProverBase : public Entity<PCS> {
 
   [[nodiscard]] bool Evaluate(const Poly& poly, const F& x) {
     F result = poly.Evaluate(x);
-    return GetWriter()->WriteToProof(result);
+    return GetWriter()->template WriteToProof</*NeedToWriteToTranscript=*/true>(
+        result);
   }
 
  protected:

diff --git a/tachyon/zk/plonk/halo2/blake2b_transcript_unittest.cc b/tachyon/zk/plonk/halo2/blake2b_transcript_unittest.cc
@@ -30,12 +30,12 @@ TEST_F(Blake2bTranscriptTest, WritePoint) {
   base::Uint8VectorBuffer write_buf;
   Blake2bWriter<G1AffinePoint> writer(std::move(write_buf));
   G1AffinePoint expected = G1AffinePoint::Random();
-  ASSERT_TRUE(writer.WriteToProof(expected));
+  ASSERT_TRUE(writer.WriteToProof</*NeedToWriteToTranscript=*/true>(expected));
 
   base::Buffer read_buf(writer.buffer().buffer(), writer.buffer().buffer_len());
   Blake2bReader<G1AffinePoint> reader(std::move(read_buf));
   G1AffinePoint actual;
-  ASSERT_TRUE(reader.ReadFromProof(&actual));
+  ASSERT_TRUE(reader.ReadFromProof</*NeedToWriteToTranscript=*/true>(&actual));
 
   EXPECT_EQ(expected, actual);
 }
@@ -44,12 +44,12 @@ TEST_F(Blake2bTranscriptTest, WriteScalar) {
   base::Uint8VectorBuffer write_buf;
   Blake2bWriter<G1AffinePoint> writer(std::move(write_buf));
   Fr expected = Fr::Random();
-  ASSERT_TRUE(writer.WriteToProof(expected));
+  ASSERT_TRUE(writer.WriteToProof</*NeedToWriteToTranscript=*/true>(expected));
 
   base::Buffer read_buf(writer.buffer().buffer(), writer.buffer().buffer_len());
   Blake2bReader<G1AffinePoint> reader(std::move(read_buf));
   Fr actual;
-  ASSERT_TRUE(reader.ReadFromProof(&actual));
+  ASSERT_TRUE(reader.ReadFromProof</*NeedToWriteToTranscript=*/true>(&actual));
 
   EXPECT_EQ(expected, actual);
 }
@@ -58,7 +58,7 @@ TEST_F(Blake2bTranscriptTest, SqueezeChallenge) {
   base::Uint8VectorBuffer write_buf;
   Blake2bWriter<G1AffinePoint> writer(std::move(write_buf));
   G1AffinePoint generator = G1AffinePoint::Generator();
-  ASSERT_TRUE(writer.WriteToProof(generator));
+  ASSERT_TRUE(writer.WriteToProof</*NeedToWriteToTranscript=*/true>(generator));
 
   std::vector<uint8_t> expected_bytes = {57, 2,   118, 182, 16,  184, 59,  179,
                                          70, 176, 223, 71,  62,  168, 222, 171,

diff --git a/tachyon/zk/plonk/halo2/poseidon_transcript_unittest.cc b/tachyon/zk/plonk/halo2/poseidon_transcript_unittest.cc
@@ -30,12 +30,12 @@ TEST_F(PoseidonTranscriptTest, WritePoint) {
   base::Uint8VectorBuffer write_buf;
   PoseidonWriter<G1AffinePoint> writer(std::move(write_buf));
   G1AffinePoint expected = G1AffinePoint::Random();
-  ASSERT_TRUE(writer.WriteToProof(expected));
+  ASSERT_TRUE(writer.WriteToProof</*NeedToWriteToTranscript=*/true>(expected));
 
   base::Buffer read_buf(writer.buffer().buffer(), writer.buffer().buffer_len());
   PoseidonReader<G1AffinePoint> reader(std::move(read_buf));
   G1AffinePoint actual;
-  ASSERT_TRUE(reader.ReadFromProof(&actual));
+  ASSERT_TRUE(reader.ReadFromProof</*NeedToWriteToTranscript=*/true>(&actual));
 
   EXPECT_EQ(expected, actual);
 }
@@ -44,12 +44,12 @@ TEST_F(PoseidonTranscriptTest, WriteScalar) {
   base::Uint8VectorBuffer write_buf;
   PoseidonWriter<G1AffinePoint> writer(std::move(write_buf));
   Fr expected = Fr::Random();
-  ASSERT_TRUE(writer.WriteToProof(expected));
+  ASSERT_TRUE(writer.WriteToProof</*NeedToWriteToTranscript=*/true>(expected));
 
   base::Buffer read_buf(writer.buffer().buffer(), writer.buffer().buffer_len());
   PoseidonReader<G1AffinePoint> reader(std::move(read_buf));
   Fr actual;
-  ASSERT_TRUE(reader.ReadFromProof(&actual));
+  ASSERT_TRUE(reader.ReadFromProof</*NeedToWriteToTranscript=*/true>(&actual));
 
   EXPECT_EQ(expected, actual);
 }
@@ -58,7 +58,7 @@ TEST_F(PoseidonTranscriptTest, SqueezeChallenge) {
   base::Uint8VectorBuffer write_buf;
   PoseidonWriter<G1AffinePoint> writer(std::move(write_buf));
   G1AffinePoint generator = G1AffinePoint::Generator();
-  ASSERT_TRUE(writer.WriteToProof(generator));
+  ASSERT_TRUE(writer.WriteToProof</*NeedToWriteToTranscript=*/true>(generator));
 
   std::vector<uint8_t> expected_bytes = {25,  86,  205, 219, 59,  135, 187, 231,
                                          192, 54,  23,  138, 114, 176, 9,   157,

diff --git a/tachyon/zk/plonk/halo2/proof_reader.h b/tachyon/zk/plonk/halo2/proof_reader.h
@@ -244,7 +244,8 @@ class ProofReader {
   template <typename T>
   T Read() {
     T value;
-    CHECK(transcript_->ReadFromProof(&value));
+    CHECK(transcript_->template ReadFromProof</*NeedToWriteToTranscript=*/true>(
+        &value));
     return value;
   }
 

diff --git a/tachyon/zk/plonk/halo2/sha256_transcript_unittest.cc b/tachyon/zk/plonk/halo2/sha256_transcript_unittest.cc
@@ -30,12 +30,12 @@ TEST_F(Sha256TranscriptTest, WritePoint) {
   base::Uint8VectorBuffer write_buf;
   Sha256Writer<G1AffinePoint> writer(std::move(write_buf));
   G1AffinePoint expected = G1AffinePoint::Random();
-  ASSERT_TRUE(writer.WriteToProof(expected));
+  ASSERT_TRUE(writer.WriteToProof</*NeedToWriteToTranscript=*/true>(expected));
 
   base::Buffer read_buf(writer.buffer().buffer(), writer.buffer().buffer_len());
   Sha256Reader<G1AffinePoint> reader(std::move(read_buf));
   G1AffinePoint actual;
-  ASSERT_TRUE(reader.ReadFromProof(&actual));
+  ASSERT_TRUE(reader.ReadFromProof</*NeedToWriteToTranscript=*/true>(&actual));
 
   EXPECT_EQ(expected, actual);
 }
@@ -44,12 +44,12 @@ TEST_F(Sha256TranscriptTest, WriteScalar) {
   base::Uint8VectorBuffer write_buf;
   Sha256Writer<G1AffinePoint> writer(std::move(write_buf));
   Fr expected = Fr::Random();
-  ASSERT_TRUE(writer.WriteToProof(expected));
+  ASSERT_TRUE(writer.WriteToProof</*NeedToWriteToTranscript=*/true>(expected));
 
   base::Buffer read_buf(writer.buffer().buffer(), writer.buffer().buffer_len());
   Sha256Reader<G1AffinePoint> reader(std::move(read_buf));
   Fr actual;
-  ASSERT_TRUE(reader.ReadFromProof(&actual));
+  ASSERT_TRUE(reader.ReadFromProof</*NeedToWriteToTranscript=*/true>(&actual));
 
   EXPECT_EQ(expected, actual);
 }
@@ -58,7 +58,7 @@ TEST_F(Sha256TranscriptTest, SqueezeChallenge) {
   base::Uint8VectorBuffer write_buf;
   Sha256Writer<G1AffinePoint> writer(std::move(write_buf));
   G1AffinePoint generator = G1AffinePoint::Generator();
-  ASSERT_TRUE(writer.WriteToProof(generator));
+  ASSERT_TRUE(writer.WriteToProof</*NeedToWriteToTranscript=*/true>(generator));
 
   std::vector<uint8_t> expected_bytes = {144, 70,  170, 43,  125, 191, 116, 100,
                                          115, 242, 37,  247, 43,  227, 23,  192,

diff --git a/tachyon/zk/plonk/vanishing/prover_vanishing_argument.h b/tachyon/zk/plonk/vanishing/prover_vanishing_argument.h
@@ -119,7 +119,9 @@ template <typename PCS, typename F, typename Commitment>
   VanishingCommitted<EntityTy::kProver, PCS> committed =
       std::move(constructed).TakeCommitted();
   F random_eval = committed.random_poly().Evaluate(x);
-  if (!writer->WriteToProof(random_eval)) return false;
+  if (!writer->template WriteToProof</*NeedToWriteToTranscript=*/true>(
+          random_eval))
+    return false;
 
   *evaluated_out = {std::move(h_poly), std::move(h_blind),
                     std::move(committed)};