Merge pull request #448 from kroma-network/feat/implement-plonky3-cha…

…llenges

feat(zk): implement plonky3 challenges

tachyon/math/finite_fields/generator/prime_field_generator/small_prime_field_config.h.tpl

@@ -3,6 +3,7 @@
 
 #include "tachyon/export.h"
 #include "tachyon/build/build_config.h"
+#include "tachyon/math/base/big_int.h"
 
 namespace %{namespace} {
 
@@ -63,6 +64,10 @@ class TACHYON_EXPORT %{class}Config {
     return (uint64_t{v} << 32) % kModulus;
   }
 
+  constexpr static uint32_t ToMontgomery(uint64_t v) {
+    return static_cast<uint32_t>(((tachyon::math::BigInt<2>(v) << 32) % tachyon::math::BigInt<2>(kModulus))[0]);
+  }
+
   constexpr static uint32_t FromMontgomery(uint64_t v) {
     constexpr uint64_t kMask = (uint64_t{1} << 32) - 1;
     uint64_t t = (v * kInverse32) & kMask;

tachyon/zk/air/plonky3/base/BUILD.bazel

@@ -0,0 +1,13 @@
+load("//bazel:tachyon_cc.bzl", "tachyon_cc_library")
+
+package(default_visibility = ["//visibility:public"])
+
+tachyon_cc_library(
+    name = "multi_field32_conversions",
+    hdrs = ["multi_field32_conversions.h"],
+    deps = [
+        "//tachyon/base/containers:adapters",
+        "//tachyon/build:build_config",
+        "@com_google_absl//absl/types:span",
+    ],
+)

tachyon/zk/air/plonky3/base/multi_field32_conversions.h

@@ -0,0 +1,73 @@
+// Copyright (c) 2022 The Plonky3 Authors
+// Use of this source code is governed by a MIT/Apache-2.0 style license that
+// can be found in the LICENSE-MIT.plonky3 and the LICENCE-APACHE.plonky3
+// file.
+
+#ifndef TACHYON_ZK_AIR_PLONKY3_BASE_MULTI_FIELD32_CONVERSIONS_H_
+#define TACHYON_ZK_AIR_PLONKY3_BASE_MULTI_FIELD32_CONVERSIONS_H_
+
+#include <stdint.h>
+
+#include <array>
+#include <limits>
+
+#include "absl/types/span.h"
+
+#include "tachyon/base/containers/adapters.h"
+#include "tachyon/build/build_config.h"
+
+namespace tachyon::zk::air::plonky3 {
+
+template <typename BigF, typename SmallF>
+BigF Reduce(absl::Span<const SmallF> values) {
+  static_assert(SmallF::Config::kModulusBits <= 32);
+  static_assert(BigF::Config::kModulusBits > 64);
+
+  using BigInt = typename BigF::BigIntTy;
+  CHECK_LT(values.size(), BigInt::kLimbNums * 2);
+
+  BigInt ret;
+  for (size_t i = 0; i < values.size(); i += 2) {
+    uint32_t value = values[i].value();
+    if constexpr (SmallF::Config::kUseMontgomery) {
+      ret[i >> 1] = SmallF::Config::FromMontgomery(value);
+    } else {
+      ret[i >> 1] = value;
+    }
+    if (i < values.size() - 1) {
+      uint64_t value2 = values[i + 1].value();
+      if constexpr (SmallF::Config::kUseMontgomery) {
+        ret[i >> 1] += uint64_t{SmallF::Config::FromMontgomery(value2)} << 32;
+      } else {
+        ret[i >> 1] += value2 << 32;
+      }
+    }
+  }
+  return BigF(ret % BigF::Config::kModulus);
+}
+
+template <typename SmallF, typename BigF,
+          size_t N = BigF::Config::kModulusBits / 64>
+std::array<SmallF, N> Split(const BigF& value) {
+  static_assert(SmallF::Config::kModulusBits <= 32);
+  static_assert(BigF::Config::kModulusBits > 64);
+  static_assert(ARCH_CPU_LITTLE_ENDIAN);
+
+  using BigInt = typename BigF::BigIntTy;
+  std::array<SmallF, N> ret;
+  BigInt value_bigint = value.ToBigInt();
+  for (size_t i = 0; i < N; ++i) {
+    uint64_t digit = value_bigint[0] & std::numeric_limits<uint64_t>::max();
+    if constexpr (SmallF::Config::kUseMontgomery) {
+      ret[i] = SmallF::FromMontgomery(SmallF::Config::ToMontgomery(digit));
+    } else {
+      ret[i] = SmallF(digit);
+    }
+    value_bigint >>= 64;
+  }
+  return ret;
+}
+
+}  // namespace tachyon::zk::air::plonky3
+
+#endif  // TACHYON_ZK_AIR_PLONKY3_BASE_MULTI_FIELD32_CONVERSIONS_H_

tachyon/zk/air/plonky3/challenger/BUILD.bazel

@@ -0,0 +1,71 @@
+load("//bazel:tachyon_cc.bzl", "tachyon_cc_library", "tachyon_cc_unittest")
+
+package(default_visibility = ["//visibility:public"])
+
+tachyon_cc_library(
+    name = "challenger",
+    hdrs = ["challenger.h"],
+    deps = [
+        ":challenger_traits_forward",
+        "//tachyon/base:logging",
+        "//tachyon/base:openmp_util",
+        "//tachyon/base:range",
+        "//tachyon/base/containers:container_util",
+    ],
+)
+
+tachyon_cc_library(
+    name = "challenger_traits_forward",
+    hdrs = ["challenger_traits_forward.h"],
+)
+
+tachyon_cc_library(
+    name = "duplex_challenger",
+    hdrs = ["duplex_challenger.h"],
+    deps = [
+        ":challenger",
+        "//tachyon/crypto/hashes/sponge:sponge_state",
+        "@com_google_absl//absl/container:inlined_vector",
+    ],
+)
+
+tachyon_cc_library(
+    name = "hash_challenger",
+    hdrs = ["hash_challenger.h"],
+    deps = [
+        ":challenger",
+        "//tachyon/base/containers:container_util",
+    ],
+)
+
+tachyon_cc_library(
+    name = "multi_field32_challenger",
+    hdrs = ["multi_field32_challenger.h"],
+    deps = [
+        ":challenger",
+        "//tachyon/base/containers:container_util",
+        "//tachyon/crypto/hashes/sponge:sponge_state",
+        "//tachyon/zk/air/plonky3/base:multi_field32_conversions",
+        "@com_google_absl//absl/container:inlined_vector",
+    ],
+)
+
+tachyon_cc_unittest(
+    name = "challenger_unittests",
+    srcs = [
+        "duplex_challenger_unittest.cc",
+        "hash_challenger_unittest.cc",
+        "multi_field32_challenger_unittest.cc",
+    ],
+    deps = [
+        ":duplex_challenger",
+        ":hash_challenger",
+        ":multi_field32_challenger",
+        "//tachyon/crypto/hashes/sponge:padding_free_sponge",
+        "//tachyon/crypto/hashes/sponge/poseidon2",
+        "//tachyon/crypto/hashes/sponge/poseidon2:poseidon2_plonky3_external_matrix",
+        "//tachyon/math/elliptic_curves/bn/bn254:poseidon2",
+        "//tachyon/math/finite_fields/baby_bear:poseidon2",
+        "//tachyon/math/finite_fields/test:finite_field_test",
+    ],
+)

tachyon/zk/air/plonky3/challenger/challenger.h

@@ -0,0 +1,126 @@
+// Copyright (c) 2022 The Plonky3 Authors
+// Use of this source code is governed by a MIT/Apache-2.0 style license that
+// can be found in the LICENSE-MIT.plonky3 and the LICENCE-APACHE.plonky3
+// file.
+
+#ifndef TACHYON_ZK_AIR_PLONKY3_CHALLENGER_CHALLENGER_H_
+#define TACHYON_ZK_AIR_PLONKY3_CHALLENGER_CHALLENGER_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <algorithm>
+#include <array>
+#include <limits>
+#include <vector>
+
+#include "tachyon/base/containers/container_util.h"
+#include "tachyon/base/logging.h"
+#include "tachyon/base/openmp_util.h"
+#include "tachyon/base/range.h"
+#include "tachyon/zk/air/plonky3/challenger/challenger_traits_forward.h"
+
+namespace tachyon::zk::air::plonky3 {
+
+template <typename Derived>
+class Challenger {
+ public:
+  using Field = typename ChallengerTraits<Derived>::Field;
+
+  template <typename T>
+  void Observe(const T& value) {
+    Derived* derived = static_cast<Derived*>(this);
+    derived->DoObserve(value);
+  }
+
+  template <typename Container>
+  void ObserveContainer(const Container& container) {
+    Derived* derived = static_cast<Derived*>(this);
+    for (const auto& value : container) {
+      derived->DoObserve(value);
+    }
+  }
+
+  template <typename Container>
+  void ObserveContainer2D(const Container& container_2d) {
+    Derived* derived = static_cast<Derived*>(this);
+    for (const auto& container : container_2d) {
+      for (const auto& value : container) {
+        derived->DoObserve(value);
+      }
+    }
+  }
+
+  Field Sample() {
+    Derived* derived = static_cast<Derived*>(this);
+    return derived->DoSample();
+  }
+
+  template <size_t N>
+  std::array<Field, N> SampleArray() {
+    return base::CreateArray(N, [this]() { return Sample(); });
+  }
+
+  template <typename ExtField>
+  ExtField SampleExtElement() {
+    constexpr size_t N = ExtField::kDegreeOverBasePrimeField;
+    using F = typename ExtField::BasePrimeField;
+    static_assert(std::is_same_v<F, Field>);
+    std::array<F, N> prime_fields = SampleArray<N>();
+    return ExtField::FromBasePrimeFields(prime_fields);
+  }
+
+  uint32_t SampleBits(uint32_t bits) {
+    static_assert(Field::Config::kModulusBits <= 32);
+    DCHECK_LT(bits, sizeof(uint32_t) * 8);
+    DCHECK_LT(uint32_t{1} << bits, Field::Config::kModulus);
+    Field rand_f = Sample();
+    uint32_t rand_size;
+    if constexpr (Field::Config::kUseMontgomery) {
+      rand_size = Field::Config::FromMontgomery(rand_f.value());
+    } else {
+      rand_size = rand_f.value();
+    }
+    return rand_size & ((uint32_t{1} << bits) - 1);
+  }
+
+  Field Grind(uint32_t bits) {
+    return Grind(bits, base::Range<uint32_t>::Until(Field::Config::kModulus));
+  }
+
+  Field Grind(uint32_t bits, base::Range<uint32_t> range) {
+#if defined(TACHYON_HAS_OPENMP)
+    uint32_t thread_nums = static_cast<uint32_t>(omp_get_max_threads());
+#else
+    uint32_t thread_nums = 1;
+#endif
+    uint32_t chunk_size = range.GetSize() / thread_nums;
+    std::vector<uint32_t> ret(thread_nums,
+                              std::numeric_limits<uint32_t>::max());
+    OPENMP_PARALLEL_FOR(uint32_t i = 0; i < thread_nums; ++i) {
+      uint32_t start = range.from + i * chunk_size;
+      uint32_t end = start + std::min(range.to - start, chunk_size);
+      for (uint32_t j = start; j < end; ++j) {
+        Derived derived = *static_cast<Derived*>(this);
+        if (derived.CheckWitness(bits, Field(j))) {
+          ret[i] = j;
+          break;
+        }
+      }
+    }
+    auto it = std::find_if(ret.begin(), ret.end(), [](uint32_t v) {
+      return v != std::numeric_limits<uint32_t>::max();
+    });
+    CHECK(it != ret.end());
+    return Field(*it);
+  }
+
+  bool CheckWitness(uint32_t bits, const Field& witness) {
+    Observe(witness);
+    return SampleBits(bits) == 0;
+  }
+};
+
+}  // namespace tachyon::zk::air::plonky3
+
+#endif  // TACHYON_ZK_AIR_PLONKY3_CHALLENGER_CHALLENGER_H_

tachyon/zk/air/plonky3/challenger/challenger_traits_forward.h

@@ -0,0 +1,16 @@
+// Copyright (c) 2022 The Plonky3 Authors
+// Use of this source code is governed by a MIT/Apache-2.0 style license that
+// can be found in the LICENSE-MIT.plonky3 and the LICENCE-APACHE.plonky3
+// file.
+
+#ifndef TACHYON_ZK_AIR_PLONKY3_CHALLENGER_CHALLENGER_TRAITS_FORWARD_H_
+#define TACHYON_ZK_AIR_PLONKY3_CHALLENGER_CHALLENGER_TRAITS_FORWARD_H_
+
+namespace tachyon::zk::air::plonky3 {
+
+template <typename T>
+struct ChallengerTraits;
+
+}  // namespace tachyon::zk::air::plonky3
+
+#endif  // TACHYON_ZK_AIR_PLONKY3_CHALLENGER_CHALLENGER_TRAITS_FORWARD_H_