feat!: update c API of CosetLDEBatch() to accept pre-allocated memory

This change removes the need for the Rust side to call `std::mem::forget()`,
streamlining memory management.

BREAKING CHANGE: The function signature of
`tachyon_sp1_baby_bear_poseidon2_two_adic_fri_coset_lde_batch()` is
changed.

tachyon/c/crypto/commitments/fri/two_adic_fri_impl.h

@@ -22,39 +22,33 @@ class TwoAdicFRIImpl
  using Commitment = typename Base::Commitment;
  using ProverData = typename Base::ProverData;
  using Domain = typename Base::Domain;
- using OpeningPoints = typename Base::OpeningPoints;
- using OpenedValues = typename Base::OpenedValues;
- using FRIProof = typename Base::FRIProof;
 
  using Base::Base;
 
  void AllocateLDEs(size_t size) { ldes_.reserve(size); }
 
- template <typename Derived>
- absl::Span<F> CosetLDEBatch(Eigen::MatrixBase<Derived>& matrix, F shift) {
+ void CosetLDEBatch(Eigen::Map<tachyon::math::RowMajorMatrix<F>>& matrix,
+ F shift,
+ Eigen::Map<tachyon::math::RowMajorMatrix<F>>& lde) {
  Domain coset = this->GetNaturalDomainForDegree(matrix.rows());
- tachyon::math::RowMajorMatrix<F> lde(
- matrix.rows() << this->config_.log_blowup, matrix.cols());
  coset.domain()->CosetLDEBatch(matrix, this->config_.log_blowup, shift, lde,
  /*reverse_at_last=*/false);
- absl::Span<F> ret(lde.data(), lde.size());
- ldes_.push_back(std::move(lde));
- return ret;
+ ldes_.push_back(Eigen::Map<const tachyon::math::RowMajorMatrix<F>>(
+ lde.data(), lde.rows(), lde.cols()));
  }
 
  using Base::Commit;
 
  void Commit(Commitment* commitment, ProverData** prover_data_out,
  std::vector<std::unique_ptr<ProverData>>* prover_data_by_round) {
  std::unique_ptr<ProverData> prover_data(new ProverData);
- CHECK(this->mmcs_.CommitOwned(std::move(ldes_), commitment,
- prover_data.get()));
+ CHECK(this->mmcs_.Commit(std::move(ldes_), commitment, prover_data.get()));
  *prover_data_out = prover_data.get();
  prover_data_by_round->push_back(std::move(prover_data));
  }
 
  protected:
- std::vector<math::RowMajorMatrix<F>> ldes_;
+ std::vector<Eigen::Map<const tachyon::math::RowMajorMatrix<F>>> ldes_;
 };
 
 } // namespace tachyon::c::crypto

tachyon/c/zk/air/sp1/baby_bear_poseidon2_two_adic_fri.cc

@@ -99,17 +99,20 @@ void tachyon_sp1_baby_bear_poseidon2_two_adic_fri_allocate_ldes(
  c::base::native_cast(pcs)->AllocateLDEs(size);
 }
 
-tachyon_baby_bear* tachyon_sp1_baby_bear_poseidon2_two_adic_fri_coset_lde_batch(
+void tachyon_sp1_baby_bear_poseidon2_two_adic_fri_coset_lde_batch(
  tachyon_sp1_baby_bear_poseidon2_two_adic_fri* pcs,
  tachyon_baby_bear* values, size_t rows, size_t cols,
- tachyon_baby_bear shift, size_t* new_rows) {
+ tachyon_baby_bear* extended_values, tachyon_baby_bear shift) {
+ PCS* native_pcs = c::base::native_cast(pcs);
  Eigen::Map<math::RowMajorMatrix<F>> matrix(c::base::native_cast(values),
  static_cast<Eigen::Index>(rows),
  static_cast<Eigen::Index>(cols));
- absl::Span<F> ret = c::base::native_cast(pcs)->CosetLDEBatch(
- matrix, c::base::native_cast(shift));
- *new_rows = ret.size() / cols;
- return c::base::c_cast(ret.data());
+ Eigen::Map<math::RowMajorMatrix<F>> extended_matrix(
+ c::base::native_cast(extended_values),
+ static_cast<Eigen::Index>(rows) << (native_pcs->config().log_blowup),
+ static_cast<Eigen::Index>(cols));
+ native_pcs->CosetLDEBatch(matrix, c::base::native_cast(shift),
+ extended_matrix);
 }
 
 void tachyon_sp1_baby_bear_poseidon2_two_adic_fri_commit(

tachyon/c/zk/air/sp1/baby_bear_poseidon2_two_adic_fri.h

@@ -68,16 +68,15 @@ tachyon_sp1_baby_bear_poseidon2_two_adic_fri_allocate_ldes(
  * @param values A pointer to the data of the baby bear row major matrix.
  * @param rows The number of rows.
  * @param cols The number of columns.
+ * @param extended_values A pointer to the data of the extended baby bear row
+ * major matrix.
  * @param shift The shift value.
- * @param new_rows The number of rows of the baby bear row major matrix.
- * @return A pointer to the data of the newly created baby bear row major
- * matrix.
  */
-TACHYON_C_EXPORT tachyon_baby_bear*
+TACHYON_C_EXPORT void
 tachyon_sp1_baby_bear_poseidon2_two_adic_fri_coset_lde_batch(
  tachyon_sp1_baby_bear_poseidon2_two_adic_fri* pcs,
  tachyon_baby_bear* values, size_t rows, size_t cols,
- tachyon_baby_bear shift, size_t* new_rows);
+ tachyon_baby_bear* extended_values, tachyon_baby_bear shift);
 
 /**
  * @brief Commits to the mixed matrix created by

tachyon/c/zk/air/sp1/baby_bear_poseidon2_two_adic_fri_unittest.cc

@@ -27,12 +27,14 @@ using MMCS = c::zk::air::plonky3::baby_bear::MMCS;
 using Coset = c::zk::air::plonky3::baby_bear::Coset;
 using PCS = c::zk::air::plonky3::baby_bear::PCS;
 
+constexpr uint32_t kLogBlowup = 1;
 constexpr size_t kRounds = 1;
 
 class TwoAdicFRITest : public testing::Test {
  public:
  void SetUp() override {
- pcs_ = tachyon_sp1_baby_bear_poseidon2_two_adic_fri_create(1, 10, 8);
+ pcs_ =
+ tachyon_sp1_baby_bear_poseidon2_two_adic_fri_create(kLogBlowup, 10, 8);
  prover_data_vec_ =
  tachyon_sp1_baby_bear_poseidon2_field_merkle_tree_vec_create();
  opening_points_ =
@@ -74,6 +76,7 @@ class TwoAdicFRITest : public testing::Test {
 TEST_F(TwoAdicFRITest, APIs) {
  constexpr size_t kRowsForInput = 32;
  constexpr size_t kColsForInput = 5;
+ constexpr size_t kExtendedRowsForInput = kRowsForInput << kLogBlowup;
  constexpr size_t kRowsForOpening = 2;
  constexpr size_t kColsForOpening = 2;
 
@@ -90,14 +93,15 @@ TEST_F(TwoAdicFRITest, APIs) {
  []() { return F::Random(); });
  std::vector<F> matrix_data_clone = matrix_data;
 
- tachyon_sp1_baby_bear_poseidon2_two_adic_fri_allocate_ldes(pcs_, 1);
+ tachyon_sp1_baby_bear_poseidon2_two_adic_fri_allocate_ldes(pcs_, kLogBlowup);
+ std::vector<F> extended_matrix_data(kExtendedRowsForInput * kColsForInput);
  F shift = F::FromMontgomery(F::Config::kSubgroupGenerator) *
  coset.domain()->offset_inv();
- size_t new_rows;
- tachyon_baby_bear* new_matrix_data =
- tachyon_sp1_baby_bear_poseidon2_two_adic_fri_coset_lde_batch(
-  pcs_, c::base::c_cast(const_cast<F*>(matrix_data.data())),
-  kRowsForInput, kColsForInput, c::base::c_cast(shift), &new_rows);
+ tachyon_sp1_baby_bear_poseidon2_two_adic_fri_coset_lde_batch(
+  pcs_, c::base::c_cast(const_cast<F*>(matrix_data.data())), kRowsForInput,
+ kColsForInput,
+ c::base::c_cast(const_cast<F*>(extended_matrix_data.data())),
+ c::base::c_cast(shift));
  tachyon_baby_bear commitment[TACHYON_PLONKY3_BABY_BEAR_POSEIDON2_CHUNK];
  tachyon_sp1_baby_bear_poseidon2_field_merkle_tree* prover_data = nullptr;
  tachyon_sp1_baby_bear_poseidon2_two_adic_fri_commit(
@@ -118,7 +122,7 @@ TEST_F(TwoAdicFRITest, APIs) {
  }
 
  math::RowMajorMatrix<F> leaf = Eigen::Map<math::RowMajorMatrix<F>>(
- c::base::native_cast(new_matrix_data), new_rows, kColsForInput);
+ extended_matrix_data.data(), kExtendedRowsForInput, kColsForInput);
  EXPECT_EQ(leaf, native_prover_data.leaves()[0]);
 
  ExtF point = ExtF::Random();

tachyon/crypto/commitments/fri/two_adic_fri.h

@@ -66,6 +66,8 @@ class TwoAdicFRI {
  TwoAdicFRI(InputMMCS&& mmcs, FRIConfig<ChallengeMMCS>&& config)
  : mmcs_(std::move(mmcs)), config_(std::move(config)) {}
 
+ const FRIConfig<ChallengeMMCS>& config() const { return config_; }
+
  Domain GetNaturalDomainForDegree(size_t size) {
  uint32_t log_n = base::bits::CheckedLog2(size);
  return Domain(log_n, F::One());

tachyon/math/polynomials/univariate/radix2_evaluation_domain.h

@@ -118,10 +118,10 @@ class Radix2EvaluationDomain
  ReverseMatrixIndexBits(mat);
  }
 
- template <typename Derived, typename Derived2>
+ template <typename Derived>
  CONSTEXPR_IF_NOT_OPENMP void CosetLDEBatch(
  Eigen::MatrixBase<Derived>& mat, size_t added_bits, F shift,
- Eigen::MatrixBase<Derived2>& out, bool reverse_at_last = true) const {
+ Eigen::MatrixBase<Derived>& out, bool reverse_at_last = true) const {
  TRACE_EVENT("EvaluationDomain", "Radix2EvaluationDomain::CosetLDEBatch");
  if constexpr (F::Config::kModulusBits > 32) {
  NOTREACHED();
@@ -448,10 +448,10 @@ class Radix2EvaluationDomain
  // moving values from row |i| to row |i|^(|added_bits|). All new entries are
  // set to |F::Zero()|.
  // Note that it crashes if the |added_bits| is zero.
- template <typename Derived, typename Derived2>
+ template <typename Derived>
  CONSTEXPR_IF_NOT_OPENMP static void ExpandWithZeroPad(
  Eigen::MatrixBase<Derived>& mat, size_t added_bits,
- Eigen::MatrixBase<Derived2>& out) {
+ Eigen::MatrixBase<Derived>& out) {
  CHECK_GT(added_bits, size_t{0});
 
  Eigen::Index new_rows = mat.rows() << added_bits;

vendors/sp1/include/baby_bear_poseidon2_two_adic_fri_pcs.h

@@ -34,9 +34,9 @@ class TwoAdicFriPcs {
  ~TwoAdicFriPcs();
 
  void allocate_ldes(size_t size) const;
- rust::Slice<TachyonBabyBear> coset_lde_batch(
- rust::Slice<TachyonBabyBear> values, size_t cols,
- const TachyonBabyBear& shift) const;
+ void coset_lde_batch(rust::Slice<TachyonBabyBear> values, size_t cols,
+  rust::Slice<TachyonBabyBear> extended_values,
+  const TachyonBabyBear& shift) const;
  std::unique_ptr<ProverData> commit(
  const ProverDataVec& prover_data_vec) const;
  std::unique_ptr<OpeningProof> do_open(const ProverDataVec& prover_data_vec,

vendors/sp1/src/baby_bear_poseidon2.rs

@@ -157,8 +157,9 @@ pub mod ffi {
  &self,
  values: &mut [TachyonBabyBear],
  cols: usize,
+ extended_values: &mut [TachyonBabyBear],
  shift: &TachyonBabyBear,
- ) -> &mut [TachyonBabyBear];
+ );
  fn commit(&self, prover_data_vec: &ProverDataVec) -> UniquePtr<ProverData>;
  fn do_open(
  &self,
@@ -572,6 +573,7 @@ impl<Val> ProverDataVec<Val> {
 
 pub struct TwoAdicFriPcs<Val, Dft, InputMmcs, FriMmcs> {
  log_n: usize,
+ log_blowup: usize,
  inner: cxx::UniquePtr<ffi::TwoAdicFriPcs>,
  prover_data_vec: ProverDataVec<Val>,
  _marker: PhantomData<(Val, Dft, InputMmcs, FriMmcs)>,
@@ -593,6 +595,7 @@ where
  ) -> TwoAdicFriPcs<Val, Dft, InputMmcs, FriMmcs> {
  Self {
  log_n,
+ log_blowup: fri_config.log_blowup,
  inner: ffi::new_two_adic_fri_pcs(
  fri_config.log_blowup,
  fri_config.num_queries,
@@ -610,21 +613,16 @@ where
  pub fn coset_lde_batch(
  &self,
  evals: &mut RowMajorMatrix<Val>,
+ extended_evals: &mut RowMajorMatrix<Val>,
  shift: Val,
- ) -> RowMajorMatrix<Val> {
+ ) {
  unsafe {
- let data = self.inner.coset_lde_batch(
+ self.inner.coset_lde_batch(
  std::mem::transmute(evals.values.as_mut_slice()),
  evals.width,
+ std::mem::transmute(extended_evals.values.as_mut_slice()),
  std::mem::transmute(&shift),
  );
-
- let vec = Vec::from_raw_parts(
- std::mem::transmute(data.as_mut_ptr()),
- data.len(),
- data.len(),
- );
- RowMajorMatrix::new(vec, evals.width)
  }
  }
 
@@ -706,7 +704,9 @@ where
  for (domain, mut evals) in evaluations.into_iter() {
  assert_eq!(domain.size(), evals.height());
  let shift = Val::generator() / domain.shift;
- ldes.push(self.coset_lde_batch(&mut evals, shift));
+ let mut lde = RowMajorMatrix::default(evals.width(), evals.height() << self.log_blowup);
+ self.coset_lde_batch(&mut evals, &mut lde, shift);
+ ldes.push(lde);
  }
  let mut prover_data = self.do_commit();
  prover_data.ldes = ldes;

vendors/sp1/src/baby_bear_poseidon2_two_adic_fri_pcs.cc

@@ -19,17 +19,15 @@ void TwoAdicFriPcs::allocate_ldes(size_t size) const {
  const_cast<tachyon_sp1_baby_bear_poseidon2_two_adic_fri*>(pcs_), size);
 }
 
-rust::Slice<TachyonBabyBear> TwoAdicFriPcs::coset_lde_batch(
+void TwoAdicFriPcs::coset_lde_batch(
  rust::Slice<TachyonBabyBear> values, size_t cols,
+ rust::Slice<TachyonBabyBear> extended_values,
  const TachyonBabyBear& shift) const {
- size_t new_rows;
- tachyon_baby_bear* data =
- tachyon_sp1_baby_bear_poseidon2_two_adic_fri_coset_lde_batch(
- const_cast<tachyon_sp1_baby_bear_poseidon2_two_adic_fri*>(pcs_),
- reinterpret_cast<tachyon_baby_bear*>(values.data()),
- values.size() / cols, cols,
- reinterpret_cast<const tachyon_baby_bear&>(shift), &new_rows);
- return {reinterpret_cast<TachyonBabyBear*>(data), new_rows * cols};
+ tachyon_sp1_baby_bear_poseidon2_two_adic_fri_coset_lde_batch(
+ const_cast<tachyon_sp1_baby_bear_poseidon2_two_adic_fri*>(pcs_),
+ reinterpret_cast<tachyon_baby_bear*>(values.data()), values.size() / cols,
+ cols, reinterpret_cast<tachyon_baby_bear*>(extended_values.data()),
+ reinterpret_cast<const tachyon_baby_bear&>(shift));
 }
 
 std::unique_ptr<ProverData> TwoAdicFriPcs::commit(

vendors/sp1/src/two_adic_fri_pcs.rs

@@ -185,8 +185,5 @@ mod test {
  assert!(tachyon_pcs
  .verify(rounds, &tachyon_proof, &mut tachyon_challenger_for_verify)
  .is_ok());
-
- // TODO(chokobole): `std::mem::forget` was used to prevent it from double-free. We need to figure out a more elegant solution.
- std::mem::forget(tachyon_data_by_round);
  }
 }