From 2ac0e33985be71b07b1dca35a48268d71449a6a3 Mon Sep 17 00:00:00 2001
From: Hossein Moghaddas <autquis@gmail.com>
Date: Wed, 15 Nov 2023 17:04:47 +0100
Subject: [PATCH] Update Hyrax with the shared state

---
 poly-commit/src/hyrax/data_structures.rs | 10 ++++++----
 poly-commit/src/hyrax/mod.rs             |  8 ++++----
 poly-commit/src/utils.rs                 |  5 +++--
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/poly-commit/src/hyrax/data_structures.rs b/poly-commit/src/hyrax/data_structures.rs
index da311c29..aa58b7cf 100644
--- a/poly-commit/src/hyrax/data_structures.rs
+++ b/poly-commit/src/hyrax/data_structures.rs
@@ -3,7 +3,10 @@ use ark_ff::PrimeField;
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use ark_std::{rand::RngCore, vec::Vec};
 
-use crate::{PCCommitment, PCCommitmentState, PCCommitterKey, PCUniversalParams, PCVerifierKey};
+use crate::{
+    utils::Matrix, PCCommitment, PCCommitmentState, PCCommitterKey, PCUniversalParams,
+    PCVerifierKey,
+};
 
 /// `UniversalParams` amounts to a Pederson commitment key of sufficient length
 #[derive(Derivative, CanonicalSerialize, CanonicalDeserialize)]
@@ -87,9 +90,8 @@ pub struct HyraxCommitmentState<F>
 where
     F: PrimeField,
 {
-    /// blah blah blah blah
-    /// blah blah blah blah
-    pub randomness: HyraxRandomness<F>,
+    pub(crate) randomness: HyraxRandomness<F>,
+    pub(crate) mat: Matrix<F>,
 }
 
 /// A vector of scalars, each of which multiplies the distinguished group
diff --git a/poly-commit/src/hyrax/mod.rs b/poly-commit/src/hyrax/mod.rs
index efc2e002..e169b91c 100644
--- a/poly-commit/src/hyrax/mod.rs
+++ b/poly-commit/src/hyrax/mod.rs
@@ -272,6 +272,7 @@ impl<G: AffineRepr, P: MultilinearExtension<G::ScalarField>>
             coms.push(l_comm);
             states.push(HyraxCommitmentState {
                 randomness: com_rands,
+                mat: Matrix::new_from_rows(m),
             });
         }
 
@@ -307,7 +308,7 @@ impl<G: AffineRepr, P: MultilinearExtension<G::ScalarField>>
             G::ScalarField,
             PoseidonSponge<G::ScalarField>,
         >,
-        rands: impl IntoIterator<Item = &'a Self::CommitmentState>,
+        states: impl IntoIterator<Item = &'a Self::CommitmentState>,
         rng: Option<&mut dyn RngCore>,
     ) -> Result<Self::Proof, Self::Error>
     where
@@ -343,7 +344,7 @@ impl<G: AffineRepr, P: MultilinearExtension<G::ScalarField>>
 
         for (l_poly, (l_com, state)) in labeled_polynomials
             .into_iter()
-            .zip(commitments.into_iter().zip(rands.into_iter()))
+            .zip(commitments.into_iter().zip(states.into_iter()))
         {
             let label = l_poly.label();
             if label != l_com.label() {
@@ -376,8 +377,7 @@ impl<G: AffineRepr, P: MultilinearExtension<G::ScalarField>>
             transcript.append_serializable_element(b"point", point)?;
 
             // Commiting to the matrix formed by the polynomial coefficients
-            let t_aux = flat_to_matrix_column_major(&poly.to_evaluations(), dim, dim);
-            let t = Matrix::new_from_rows(t_aux);
+            let t = &state.mat;
 
             let lt = t.row_mul(&l);
 
diff --git a/poly-commit/src/utils.rs b/poly-commit/src/utils.rs
index 7c4a0575..4a2ffeb4 100644
--- a/poly-commit/src/utils.rs
+++ b/poly-commit/src/utils.rs
@@ -7,7 +7,7 @@ use rayon::{
 };
 
 use ark_ff::{Field, PrimeField};
-use ark_serialize::CanonicalSerialize;
+use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use ark_std::vec::Vec;
 use merlin::Transcript;
 
@@ -31,7 +31,8 @@ pub(crate) fn ceil_div(x: usize, y: usize) -> usize {
     (x + y - 1) / y
 }
 
-#[derive(Debug)]
+#[derive(Derivative, CanonicalSerialize, CanonicalDeserialize)]
+#[derivative(Default(bound = ""), Clone(bound = ""), Debug(bound = ""))]
 pub(crate) struct Matrix<F: Field> {
     pub(crate) n: usize,
     pub(crate) m: usize,