feat(zk): implement poseidon transcript

See [PoseidonRead](https://github.com/kroma-network/halo2/blob/7d0a36990452c8e7ebd600de258420781a9b7917/halo2_proofs/src/transcript/poseidon.rs#L12-L100), [PoseidonWrite](https://github.com/kroma-network/halo2/blob/7d0a36990452c8e7ebd600de258420781a9b7917/halo2_proofs/src/transcript/poseidon.rs#L102-L174), [Transcript](https://github.com/kroma-network/halo2/blob/7d0a36990452c8e7ebd600de258420781a9b7917/halo2_proofs/src/transcript.rs#L43-L63), [TranscriptRead](https://github.com/kroma-network/halo2/blob/7d0a36990452c8e7ebd600de258420781a9b7917/halo2_proofs/src/transcript.rs#L65-L73) and [TranscriptWirte](https://github.com/kroma-network/halo2/blob/7d0a36990452c8e7ebd600de258420781a9b7917/halo2_proofs/src/transcript.rs#L75-L83).
kroma-network · Oct 25, 2023 · 7f28f5e · 7f28f5e
1 parent d1bdca4
commit 7f28f5e
Show file tree

Hide file tree

Showing 7 changed files with 396 additions and 0 deletions.
diff --git a/tachyon/crypto/hashes/sponge/poseidon/BUILD.bazel b/tachyon/crypto/hashes/sponge/poseidon/BUILD.bazel
@@ -12,6 +12,15 @@ tachyon_cc_library(
     ],
 )
 
+tachyon_cc_library(
+    name = "halo2_poseidon",
+    hdrs = ["halo2_poseidon.h"],
+    deps = [
+        ":poseidon",
+        "@local_config_gmp//:gmp",
+    ],
+)
+
 tachyon_cc_library(
     name = "poseidon_config",
     hdrs = ["poseidon_config.h"],

diff --git a/tachyon/crypto/hashes/sponge/poseidon/halo2_poseidon.h b/tachyon/crypto/hashes/sponge/poseidon/halo2_poseidon.h
@@ -0,0 +1,59 @@
+// Copyright 2020-2022 The Electric Coin Company
+// Copyright 2022 The Halo2 developers
+// Use of this source code is governed by a MIT/Apache-2.0 style license that
+// can be found in the LICENSE-MIT.halo2 and the LICENCE-APACHE.halo2
+// file.
+
+#ifndef TACHYON_CRYPTO_HASHES_SPONGE_POSEIDON_HALO2_POSEIDON_H_
+#define TACHYON_CRYPTO_HASHES_SPONGE_POSEIDON_HALO2_POSEIDON_H_
+
+#include <vector>
+
+#include "third_party/gmp/include/gmpxx.h"
+
+#include "tachyon/crypto/hashes/sponge/poseidon/poseidon.h"
+
+namespace tachyon::crypto {
+
+template <typename PrimeFieldTy>
+struct Halo2PoseidonSponge : public PoseidonSponge<PrimeFieldTy> {
+  using F = PrimeFieldTy;
+
+  Halo2PoseidonSponge() = default;
+  explicit Halo2PoseidonSponge(const PoseidonConfig<F>& config)
+      : PoseidonSponge<PrimeFieldTy>(config) {
+    this->state.elements[0] = F::FromMpzClass(mpz_class(1) << 64);
+  }
+
+  // FieldBasedCryptographicSponge methods
+  std::vector<F> SqueezeNativeFieldElements(size_t num_elements) {
+    std::vector<F> ret = base::CreateVector(num_elements, F::Zero());
+    size_t squeeze_index = this->state.mode.next_index;
+    if (squeeze_index == this->config.rate) {
+      squeeze_index = 0;
+    }
+    this->state[squeeze_index + 1] = F::One();
+    switch (this->state.mode.type) {
+      case DuplexSpongeMode::Type::kAbsorbing: {
+        this->Permute();
+        this->SqueezeInternal(0, &ret);
+        return ret;
+      }
+      case DuplexSpongeMode::Type::kSqueezing: {
+        size_t squeeze_index = this->state.mode.next_index;
+        if (squeeze_index == this->config.rate) {
+          this->Permute();
+          squeeze_index = 0;
+        }
+        this->SqueezeInternal(squeeze_index, &ret);
+        return ret;
+      }
+    }
+    NOTREACHED();
+    return {};
+  }
+};
+
+}  // namespace tachyon::crypto
+
+#endif  // TACHYON_CRYPTO_HASHES_SPONGE_POSEIDON_HALO2_POSEIDON_H_
diff --git a/tachyon/crypto/hashes/sponge/poseidon/poseidon.h b/tachyon/crypto/hashes/sponge/poseidon/poseidon.h
@@ -242,6 +242,8 @@ struct PoseidonSponge
   }
 
   // FieldBasedCryptographicSponge methods
+  // NOTE(TomTaehoonKim): If you ever update this, please update
+  // |Halo2PoseidonSponge| for consistency.
   std::vector<F> SqueezeNativeFieldElements(size_t num_elements) {
     std::vector<F> ret =
         base::CreateVector(num_elements, []() { return F::Zero(); });

diff --git a/tachyon/zk/transcript/BUILD.bazel b/tachyon/zk/transcript/BUILD.bazel
@@ -0,0 +1,30 @@
+load("//bazel:tachyon_cc.bzl", "tachyon_cc_library", "tachyon_cc_unittest")
+
+package(default_visibility = ["//visibility:public"])
+
+tachyon_cc_library(
+    name = "transcript",
+    hdrs = ["transcript.h"],
+    deps = ["//tachyon/math/base:big_int"],
+)
+
+tachyon_cc_library(
+    name = "poseidon_transcript",
+    hdrs = ["poseidon_transcript.h"],
+    deps = [
+        ":transcript",
+        "//tachyon/base/buffer:vector_buffer",
+        "//tachyon/crypto/hashes/sponge/poseidon:halo2_poseidon",
+    ],
+)
+
+tachyon_cc_unittest(
+    name = "transcript_unittests",
+    srcs = [
+        "poseidon_transcript_unittest.cc",
+    ],
+    deps = [
+        ":poseidon_transcript",
+        "//tachyon/math/elliptic_curves/bn/bn254:g1",
+    ],
+)
diff --git a/tachyon/zk/transcript/poseidon_transcript.h b/tachyon/zk/transcript/poseidon_transcript.h
@@ -0,0 +1,129 @@
+// Copyright 2020-2022 The Electric Coin Company
+// Copyright 2022 The Halo2 developers
+// Use of this source code is governed by a MIT/Apache-2.0 style license that
+// can be found in the LICENSE-MIT.halo2 and the LICENCE-APACHE.halo2
+// file.
+
+#ifndef TACHYON_ZK_TRANSCRIPT_POSEIDON_TRANSCRIPT_H_
+#define TACHYON_ZK_TRANSCRIPT_POSEIDON_TRANSCRIPT_H_
+
+#include <utility>
+#include <vector>
+
+#include "tachyon/base/buffer/buffer.h"
+#include "tachyon/base/buffer/vector_buffer.h"
+#include "tachyon/crypto/hashes/sponge/poseidon/halo2_poseidon.h"
+#include "tachyon/zk/transcript/transcript.h"
+
+namespace tachyon::zk {
+
+template <typename AffinePointTy>
+class PoseidonRead : public TranscriptRead<AffinePointTy> {
+ public:
+  using ScalarField = typename AffinePointTy::ScalarField;
+  using Buffer = typename base::Buffer;
+
+  constexpr PoseidonRead() = default;
+  // Initialize a transcript given an input buffer.
+  explicit PoseidonRead(Buffer read_buf)
+      : state_(
+            crypto::PoseidonConfig<ScalarField>::CreateCustom(8, 5, 8, 63, 0)),
+        buffer_(std::move(read_buf)) {}
+
+  constexpr Buffer& buffer() { return buffer_; }
+
+  // TranscriptRead methods
+  constexpr bool ReadPoint(AffinePointTy* point_out) override {
+    if (buffer_.Read(point_out)) {
+      return WriteToTranscript(*point_out);
+    }
+    return false;
+  }
+
+  constexpr bool ReadScalar(ScalarField* scalar_out) override {
+    if (buffer_.Read(scalar_out)) {
+      return WriteToTranscript(*scalar_out);
+    }
+    return false;
+  }
+
+  constexpr Challenge255<AffinePointTy> SqueezeChallenge() override {
+    return Challenge255<AffinePointTy>(state_.SqueezeNativeFieldElements(1)[0]);
+  }
+
+  constexpr bool WriteToTranscript(const AffinePointTy& point) override {
+    std::vector<ScalarField> coords = {
+        ScalarField::FromBigInt(point.x().ToBigInt() %
+                                ScalarField::Config::kModulus),
+        ScalarField::FromBigInt(point.y().ToBigInt() %
+                                ScalarField::Config::kModulus)};
+    return state_.Absorb(coords);
+  }
+
+  constexpr bool WriteToTranscript(const ScalarField& scalar) override {
+    return state_.Absorb(scalar);
+  }
+
+ private:
+  crypto::Halo2PoseidonSponge<ScalarField> state_;
+  Buffer buffer_;
+};
+
+template <typename AffinePointTy>
+class PoseidonWrite : public TranscriptWrite<AffinePointTy> {
+ public:
+  using ScalarField = typename AffinePointTy::ScalarField;
+  using VectorBuffer = base::VectorBuffer;
+
+  constexpr PoseidonWrite() = default;
+  // Initialize a transcript given an output buffer.
+  explicit PoseidonWrite(VectorBuffer write_buf)
+      : state_(
+            crypto::PoseidonConfig<ScalarField>::CreateCustom(8, 5, 8, 63, 0)),
+        buffer_(std::move(write_buf)) {}
+
+  constexpr VectorBuffer& buffer() { return buffer_; }
+
+  // Conclude the interaction and return the output buffer (writer).
+  constexpr const VectorBuffer& Finalize() const { return buffer_; }
+
+  // TranscriptWrite methods
+  constexpr bool WriteToProof(const AffinePointTy& point) override {
+    if (WriteToTranscript(point)) {
+      return buffer_.Write(point);
+    }
+    return false;
+  }
+
+  constexpr bool WriteToProof(const ScalarField& scalar) override {
+    if (WriteToTranscript(scalar)) {
+      return buffer_.Write(scalar);
+    }
+    return false;
+  }
+
+  constexpr Challenge255<AffinePointTy> SqueezeChallenge() override {
+    return Challenge255<AffinePointTy>(state_.SqueezeNativeFieldElements(1)[0]);
+  }
+
+  constexpr bool WriteToTranscript(const AffinePointTy& point) override {
+    std::vector<ScalarField> coords = {
+        ScalarField::FromBigInt(point.x().ToBigInt() %
+                                ScalarField::Config::kModulus),
+        ScalarField::FromBigInt(point.y().ToBigInt() %
+                                ScalarField::Config::kModulus)};
+    return state_.Absorb(coords);
+  }
+
+  constexpr bool WriteToTranscript(const ScalarField& scalar) override {
+    return state_.Absorb(scalar);
+  }
+
+ private:
+  crypto::Halo2PoseidonSponge<ScalarField> state_;
+  VectorBuffer buffer_;
+};
+
+}  // namespace tachyon::zk
+
+#endif  // TACHYON_ZK_TRANSCRIPT_POSEIDON_TRANSCRIPT_H_
diff --git a/tachyon/zk/transcript/poseidon_transcript_unittest.cc b/tachyon/zk/transcript/poseidon_transcript_unittest.cc
@@ -0,0 +1,81 @@
+// Copyright 2020-2022 The Electric Coin Company
+// Copyright 2022 The Halo2 developers
+// Use of this source code is governed by a MIT/Apache-2.0 style license that
+// can be found in the LICENSE-MIT.halo2 and the LICENCE-APACHE.halo2
+// file.
+
+#include "tachyon/zk/transcript/poseidon_transcript.h"
+
+#include <utility>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "tachyon/math/elliptic_curves/bn/bn254/g1.h"
+
+namespace tachyon::zk {
+
+namespace {
+
+class PoseidonTranscriptTest : public testing::Test {
+ public:
+  static void SetUpTestSuite() { math::bn254::G1Curve::Init(); }
+};
+
+}  // namespace
+
+TEST_F(PoseidonTranscriptTest, WritePoint) {
+  using AffinePoint = math::bn254::G1AffinePoint;
+
+  base::VectorBuffer write_buf;
+  PoseidonWrite<AffinePoint> writer(std::move(write_buf));
+  AffinePoint expected = AffinePoint::Random();
+  ASSERT_TRUE(writer.WriteToProof(expected));
+
+  base::Buffer read_buf(writer.buffer().buffer(), writer.buffer().buffer_len());
+  PoseidonRead<AffinePoint> reader(std::move(read_buf));
+  AffinePoint actual;
+  ASSERT_TRUE(reader.ReadPoint(&actual));
+
+  EXPECT_EQ(expected, actual);
+}
+
+TEST_F(PoseidonTranscriptTest, WriteScalar) {
+  using AffinePoint = math::bn254::G1AffinePoint;
+  using ScalarField = AffinePoint::ScalarField;
+
+  base::VectorBuffer write_buf;
+  PoseidonWrite<AffinePoint> writer(std::move(write_buf));
+  ScalarField expected = ScalarField::Random();
+  ASSERT_TRUE(writer.WriteToProof(expected));
+
+  base::Buffer read_buf(writer.buffer().buffer(), writer.buffer().buffer_len());
+  PoseidonRead<AffinePoint> reader(std::move(read_buf));
+  ScalarField actual;
+  ASSERT_TRUE(reader.ReadScalar(&actual));
+
+  EXPECT_EQ(expected, actual);
+}
+
+TEST_F(PoseidonTranscriptTest, SqueezeChallenge) {
+  using AffinePoint = math::bn254::G1AffinePoint;
+  using ScalarField = AffinePoint::ScalarField;
+
+  base::VectorBuffer write_buf;
+  PoseidonWrite<AffinePoint> writer(std::move(write_buf));
+  AffinePoint generator = AffinePoint::Generator();
+  ASSERT_TRUE(writer.WriteToProof(generator));
+
+  std::vector<uint8_t> expected_bytes = {25,  86,  205, 219, 59,  135, 187, 231,
+                                         192, 54,  23,  138, 114, 176, 9,   157,
+                                         1,   97,  110, 174, 67,  9,   89,  85,
+                                         126, 129, 216, 121, 53,  99,  227, 26};
+  ScalarField expected =
+      ScalarField::FromBigInt(math::BigInt<4>::FromBytesLE(expected_bytes));
+
+  ScalarField actual = writer.SqueezeChallenge().challenge_as_scalar();
+
+  EXPECT_EQ(expected, actual);
+}
+
+}  // namespace tachyon::zk