feat: 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 · 8825d63 · 8825d63
1 parent bbdeb05
commit 8825d63
Show file tree

Hide file tree

Showing 10 changed files with 437 additions and 0 deletions.
diff --git a/tachyon/crypto/hashes/BUILD.bazel b/tachyon/crypto/hashes/BUILD.bazel
@@ -9,6 +9,7 @@ tachyon_cc_library(
         "//tachyon/base/buffer",
         "//tachyon/math/base:big_int",
         "//tachyon/math/finite_fields:finite_field_traits",
+        "@com_google_absl//absl/container:inlined_vector",
     ],
 )
 

diff --git a/tachyon/crypto/hashes/prime_field_serializable.h b/tachyon/crypto/hashes/prime_field_serializable.h
@@ -6,9 +6,11 @@
 #ifndef TACHYON_CRYPTO_HASHES_PRIME_FIELD_SERIALIZABLE_H_
 #define TACHYON_CRYPTO_HASHES_PRIME_FIELD_SERIALIZABLE_H_
 
+#include <array>
 #include <type_traits>
 #include <vector>
 
+#include "absl/container/inlined_vector.h"
 #include "absl/types/span.h"
 
 #include "tachyon/math/base/big_int.h"
@@ -69,6 +71,28 @@ class PrimeFieldSerializable<std::vector<T>> {
   }
 };
 
+template <typename T, size_t N>
+class PrimeFieldSerializable<absl::InlinedVector<T, N>> {
+ public:
+  template <typename PrimeFieldTy>
+  constexpr static bool ToPrimeField(const absl::InlinedVector<T, N>& values,
+                                     std::vector<PrimeFieldTy>* fields) {
+    return PrimeFieldSerializable<T>::BatchToPrimeField(
+        absl::MakeConstSpan(values), fields);
+  }
+};
+
+template <typename T, size_t N>
+class PrimeFieldSerializable<std::array<T, N>> {
+ public:
+  template <typename PrimeFieldTy>
+  constexpr static bool ToPrimeField(const std::array<T, N>& values,
+                                     std::vector<PrimeFieldTy>* fields) {
+    return PrimeFieldSerializable<T>::BatchToPrimeField(
+        absl::MakeConstSpan(values), fields);
+  }
+};
+
 template <typename T, typename PrimeFieldTy>
 constexpr bool SerializeToFieldElements(const T& value,
                                         std::vector<PrimeFieldTy>* fields) {

diff --git a/tachyon/crypto/hashes/prime_field_serializable_unittest.cc b/tachyon/crypto/hashes/prime_field_serializable_unittest.cc
@@ -23,6 +23,36 @@ TEST_F(PrimeFieldSerializableTest, SerializeSingleValueToField) {
   EXPECT_EQ(fields, expected);
 }
 
+TEST_F(PrimeFieldSerializableTest, SerializeVectorToField) {
+  std::vector<int64_t> values = {1, 2, 3, 4, 5};
+  std::vector<math::GF7> fields;
+  ASSERT_TRUE(SerializeToFieldElements(values, &fields));
+  std::vector<math::GF7> expected = {
+      math::GF7(1), math::GF7(2), math::GF7(3), math::GF7(4), math::GF7(5),
+  };
+  EXPECT_EQ(fields, expected);
+}
+
+TEST_F(PrimeFieldSerializableTest, SerializeArrayToField) {
+  std::array<int64_t, 5> values = {1, 2, 3, 4, 5};
+  std::vector<math::GF7> fields;
+  ASSERT_TRUE(SerializeToFieldElements(values, &fields));
+  std::vector<math::GF7> expected = {
+      math::GF7(1), math::GF7(2), math::GF7(3), math::GF7(4), math::GF7(5),
+  };
+  EXPECT_EQ(fields, expected);
+}
+
+TEST_F(PrimeFieldSerializableTest, SerializeInlinedVectorToField) {
+  absl::InlinedVector<int64_t, 5> values = {1, 2, 3, 4, 5};
+  std::vector<math::GF7> fields;
+  ASSERT_TRUE(SerializeToFieldElements(values, &fields));
+  std::vector<math::GF7> expected = {
+      math::GF7(1), math::GF7(2), math::GF7(3), math::GF7(4), math::GF7(5),
+  };
+  EXPECT_EQ(fields, expected);
+}
+
 TEST_F(PrimeFieldSerializableTest, SerializeBatchToField) {
   std::vector<int64_t> values = {1, 2, 3, 4, 5};
   std::vector<math::GF7> fields;

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,112 @@
+// 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 <array>
+#include <utility>
+
+#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 CurveConfig = typename AffinePointTy::Curve::Config;
+
+  PoseidonRead() = default;
+  // Initialize a transcript given an input buffer.
+  explicit PoseidonRead(base::Buffer read_buf)
+      : state_(
+            crypto::PoseidonConfig<ScalarField>::CreateCustom(8, 5, 8, 63, 0)),
+        buffer_(std::move(read_buf)) {}
+
+  base::Buffer& buffer() { return buffer_; }
+  const base::Buffer& buffer() const { return buffer_; }
+
+  // Transcript methods
+  Challenge255<AffinePointTy> SqueezeChallenge() override {
+    return Challenge255<AffinePointTy>(state_.SqueezeNativeFieldElements(1)[0]);
+  }
+
+  bool WriteToTranscript(const AffinePointTy& point) override {
+    std::array<ScalarField, 2> coords = {CurveConfig::BaseToScalar(point.x()),
+                                         CurveConfig::BaseToScalar(point.y())};
+    return state_.Absorb(coords);
+  }
+
+  bool WriteToTranscript(const ScalarField& scalar) override {
+    return state_.Absorb(scalar);
+  }
+
+  // TranscriptRead methods
+  bool ReadPoint(AffinePointTy* point) override {
+    return buffer_.Read(point) && WriteToTranscript(*point);
+  }
+
+  bool ReadScalar(ScalarField* scalar) override {
+    return buffer_.Read(scalar) && WriteToTranscript(*scalar);
+  }
+
+ private:
+  crypto::Halo2PoseidonSponge<ScalarField> state_;
+  base::Buffer buffer_;
+};
+
+template <typename AffinePointTy>
+class PoseidonWrite : public TranscriptWrite<AffinePointTy> {
+ public:
+  using ScalarField = typename AffinePointTy::ScalarField;
+  using CurveConfig = typename AffinePointTy::Curve::Config;
+
+  PoseidonWrite() = default;
+  // Initialize a transcript given an output buffer.
+  explicit PoseidonWrite(base::VectorBuffer write_buf)
+      : state_(
+            crypto::PoseidonConfig<ScalarField>::CreateCustom(8, 5, 8, 63, 0)),
+        buffer_(std::move(write_buf)) {}
+
+  base::VectorBuffer& buffer() { return buffer_; }
+  const base::VectorBuffer& buffer() const { return buffer_; }
+
+  // Transcript methods
+  Challenge255<AffinePointTy> SqueezeChallenge() override {
+    return Challenge255<AffinePointTy>(state_.SqueezeNativeFieldElements(1)[0]);
+  }
+
+  bool WriteToTranscript(const AffinePointTy& point) override {
+    std::array<ScalarField, 2> coords = {CurveConfig::BaseToScalar(point.x()),
+                                         CurveConfig::BaseToScalar(point.y())};
+    return state_.Absorb(coords);
+  }
+
+  bool WriteToTranscript(const ScalarField& scalar) override {
+    return state_.Absorb(scalar);
+  }
+
+  // TranscriptWrite methods
+  bool WriteToProof(const AffinePointTy& point) override {
+    return WriteToTranscript(point) && buffer_.Write(point);
+  }
+
+  bool WriteToProof(const ScalarField& scalar) override {
+    return WriteToTranscript(scalar) && buffer_.Write(scalar);
+  }
+
+ private:
+  crypto::Halo2PoseidonSponge<ScalarField> state_;
+  base::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().ChallengeAsScalar();
+
+  EXPECT_EQ(expected, actual);
+}
+
+}  // namespace tachyon::zk