Init commit!

tensor_pcs/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "tensor-pcs"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+rand = "0.8.5"
+serde = { version = "1.0.152", features = ["derive"] }
+merlin = "3.0.0"
+ecfft = { git = "https://github.com/DanTehrani/ecfft" }
+tiny-keccak = { version = "2.0.2", features = ["keccak"] }
+halo2curves = "0.1.0"
+
+[dev-dependencies]
+criterion = { version = "0.4", features = ["html_reports"] }
+
+[[bench]]
+name = "prove"
+harness = false

tensor_pcs/benches/prove.rs

@@ -0,0 +1,93 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use tensor_pcs::{
+    rs_config, FieldExt, SparseMLPoly, TensorMultilinearPCS, TensorRSMultilinearPCSConfig,
+    Transcript,
+};
+
+fn poly<F: FieldExt>(num_vars: usize) -> SparseMLPoly<F> {
+    let num_entries: usize = 2usize.pow(num_vars as u32);
+
+    let evals = (0..num_entries)
+        .map(|i| (i, F::from(i as u64)))
+        .collect::<Vec<(usize, F)>>();
+
+    let ml_poly = SparseMLPoly::new(evals, num_vars);
+    ml_poly
+}
+
+fn config_base<F: FieldExt>(ml_poly: &SparseMLPoly<F>) -> TensorRSMultilinearPCSConfig<F> {
+    let num_vars = ml_poly.num_vars;
+    let num_evals = 2usize.pow(num_vars as u32);
+    let num_rows = 2usize.pow((num_vars / 2) as u32);
+
+    let expansion_factor = 2;
+
+    TensorRSMultilinearPCSConfig::<F> {
+        expansion_factor,
+        domain_powers: None,
+        fft_domain: None,
+        ecfft_config: None,
+        l: 10,
+        num_entries: num_evals,
+        num_rows,
+    }
+}
+
+fn pcs_fft_bench(c: &mut Criterion) {
+    type F = halo2curves::pasta::Fp;
+
+    let num_vars = 13;
+    let ml_poly = poly(num_vars);
+    let open_at = (0..ml_poly.num_vars)
+        .map(|i| F::from(i as u64))
+        .collect::<Vec<F>>();
+
+    let mut config = config_base(&ml_poly);
+    config.fft_domain = Some(rs_config::smooth::gen_config::<F>(config.num_cols()));
+
+    let mut group = c.benchmark_group("pcs fft");
+    group.bench_function("prove", |b| {
+        b.iter(|| {
+            let pcs = TensorMultilinearPCS::<F>::new(config.clone());
+
+            let mut transcript = Transcript::new(b"bench");
+            let comm = pcs.commit(&black_box(ml_poly.clone()));
+            pcs.open(&comm, &ml_poly, &open_at, &mut transcript);
+        })
+    });
+}
+
+fn pcs_ecfft_bench(c: &mut Criterion) {
+    type F = halo2curves::secp256k1::Fp;
+
+    let num_vars = 13;
+    let ml_poly = poly(num_vars);
+    let open_at = (0..ml_poly.num_vars)
+        .map(|i| F::from(i as u64))
+        .collect::<Vec<F>>();
+
+    let mut config = config_base(&ml_poly);
+    config.ecfft_config = Some(rs_config::ecfft::gen_config::<F>(config.num_cols()));
+
+    let mut group = c.benchmark_group("pcs ecfft");
+    group.bench_function("prove", |b| {
+        b.iter(|| {
+            let pcs = TensorMultilinearPCS::<F>::new(config.clone());
+
+            let mut transcript = Transcript::new(b"bench");
+            let comm = pcs.commit(&black_box(ml_poly.clone()));
+            pcs.open(&comm, &ml_poly, &open_at, &mut transcript);
+        })
+    });
+}
+
+fn set_duration() -> Criterion {
+    Criterion::default().sample_size(10)
+}
+
+criterion_group! {
+    name = benches;
+    config = set_duration();
+    targets = pcs_fft_bench, pcs_ecfft_bench
+}
+criterion_main!(benches);

tensor_pcs/src/fft.rs

@@ -0,0 +1,118 @@
+use crate::FieldExt;
+use halo2curves::ff::Field;
+use std::vec;
+
+pub fn fft<F>(coeffs: &[F], domain: &[F]) -> Vec<F>
+where
+    F: FieldExt,
+{
+    debug_assert_eq!(coeffs.len(), domain.len());
+    if coeffs.len() == 1 {
+        return coeffs.to_vec();
+    }
+
+    // TODO: Just borrow the values
+    // Split into evens and odds
+    let L = coeffs
+        .iter()
+        .enumerate()
+        .filter(|(i, _)| i % 2 == 0)
+        .map(|(_, x)| *x)
+        .collect::<Vec<F>>();
+
+    let R = coeffs
+        .iter()
+        .enumerate()
+        .filter(|(i, _)| i % 2 == 1)
+        .map(|(_, x)| *x)
+        .collect::<Vec<F>>();
+
+    // Square the domain values
+    let domain_squared: Vec<F> = (0..(domain.len() / 2)).map(|i| domain[i * 2]).collect();
+
+    let fft_e = fft(&L, &domain_squared);
+    let fft_o = fft(&R, &domain_squared);
+
+    let mut evals_L = vec![];
+    let mut evals_R = vec![];
+    for i in 0..(coeffs.len() / 2) {
+        // We can use the previous evaluations to create a list of evaluations
+        // of the domain
+        evals_L.push(fft_e[i] + fft_o[i] * domain[i]);
+        evals_R.push(fft_e[i] - fft_o[i] * domain[i]);
+    }
+
+    evals_L.extend(evals_R);
+    return evals_L;
+}
+
+pub fn ifft<F: FieldExt + Field>(domain: &[F], evals: &[F]) -> Vec<F> {
+    let mut coeffs = vec![];
+    let len_mod_inv = F::from(domain.len() as u64).invert().unwrap();
+    let vals = fft(&evals, &domain);
+
+    coeffs.push(vals[0] * len_mod_inv);
+    for val in vals[1..].iter().rev() {
+        coeffs.push(*val * len_mod_inv);
+    }
+
+    coeffs
+}
+
+#[cfg(test)]
+mod tests {
+    use halo2curves::ff::PrimeField;
+    use halo2curves::pasta::Fp;
+
+    use super::*;
+    #[test]
+    fn test_fft_ifft() {
+        // f(x) = 1 + 2x + 3x^2 + 4x^3
+        let mut coeffs = vec![
+            Fp::from(1),
+            Fp::from(2),
+            Fp::from(3),
+            Fp::from(4),
+            Fp::from(5),
+            Fp::from(6),
+            Fp::from(7),
+            Fp::from(81),
+        ];
+
+        let mut domain = vec![];
+
+        let root_of_unity = Fp::ROOT_OF_UNITY;
+
+        let subgroup_order = (coeffs.len() * 2).next_power_of_two();
+
+        coeffs.resize(subgroup_order, Fp::ZERO);
+
+        // Generator for the subgroup with order _subgroup_order_ in the field
+        let generator = root_of_unity.pow(&[
+            2u32.pow(32 - ((subgroup_order as f64).log2() as u32)) as u64,
+            0,
+            0,
+            0,
+        ]);
+
+        for i in 0..(subgroup_order) {
+            domain.push(generator.pow(&[i as u64, 0, 0, 0]));
+        }
+
+        let mut expected_evals = vec![];
+
+        for w in &domain {
+            let mut eval = Fp::ZERO;
+            for (i, coeff) in (&coeffs).iter().enumerate() {
+                eval += *coeff * w.pow(&[i as u64, 0, 0, 0]);
+            }
+            expected_evals.push(eval);
+        }
+
+        let evals = fft(&coeffs, &domain);
+        debug_assert!(evals == expected_evals);
+
+        let recovered_coeffs = ifft(&domain, &evals);
+        debug_assert!(recovered_coeffs == coeffs);
+    }
+}

tensor_pcs/src/lib.rs

@@ -0,0 +1,19 @@
+mod fft;
+mod polynomial;
+pub mod rs_config;
+mod tensor_code;
+mod tensor_pcs;
+mod transcript;
+mod tree;
+mod utils;
+use halo2curves::ff::FromUniformBytes;
+
+pub trait FieldExt: FromUniformBytes<64, Repr = [u8; 32]> {}
+
+impl FieldExt for halo2curves::secp256k1::Fp {}
+impl FieldExt for halo2curves::pasta::Fp {}
+
+pub use polynomial::eq_poly::EqPoly;
+pub use polynomial::sparse_ml_poly::SparseMLPoly;
+pub use tensor_pcs::{TensorMLOpening, TensorMultilinearPCS, TensorRSMultilinearPCSConfig};
+pub use transcript::{AppendToTranscript, Transcript};

tensor_pcs/src/polynomial/eq_poly.rs

@@ -0,0 +1,68 @@
+use crate::FieldExt;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize)]
+pub struct EqPoly<F: FieldExt> {
+    t: Vec<F>,
+}
+
+impl<F: FieldExt> EqPoly<F> {
+    pub fn new(t: Vec<F>) -> Self {
+        Self { t }
+    }
+
+    pub fn eval(&self, x: &[F]) -> F {
+        let mut result = F::ONE;
+        let one = F::ONE;
+
+        for i in 0..x.len() {
+            result *= self.t[i] * x[i] + (one - self.t[i]) * (one - x[i]);
+        }
+        result
+    }
+
+    // Copied from microsoft/Spartan
+    pub fn evals(&self) -> Vec<F> {
+        let ell = self.t.len(); // 4
+
+        let mut evals: Vec<F> = vec![F::ONE; 2usize.pow(ell as u32)];
+        let mut size = 1;
+        for j in 0..ell {
+            // in each iteration, we double the size of chis
+            size *= 2; // 2 4 8 16
+            for i in (0..size).rev().step_by(2) {
+                // copy each element from the prior iteration twice
+                let scalar = evals[i / 2]; // i = 0, 2, 4, 7
+                evals[i] = scalar * self.t[j]; // (1 * t0)(1 * t1)
+                evals[i - 1] = scalar - evals[i]; // 1 - (1 * t0)(1 * t1)
+            }
+        }
+        evals
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::polynomial::sparse_ml_poly::SparseMLPoly;
+    use halo2curves::ff::Field;
+
+    type F = halo2curves::secp256k1::Fp;
+
+    pub fn dot_prod<F: FieldExt>(x: &[F], y: &[F]) -> F {
+        assert_eq!(x.len(), y.len());
+        let mut result = F::ZERO;
+        for i in 0..x.len() {
+            result += x[i] * y[i];
+        }
+        result
+    }
+
+    #[test]
+    fn test_eq_poly() {
+        let m = 4;
+        let t = (0..m).map(|i| F::from((i + 33) as u64)).collect::<Vec<F>>();
+        let eq_poly = EqPoly::new(t.clone());
+        eq_poly.evals();
+    }
+}

tensor_pcs/src/polynomial/mod.rs

@@ -0,0 +1,2 @@
+pub mod eq_poly;
+pub mod sparse_ml_poly;

tensor_pcs/src/polynomial/sparse_ml_poly.rs

@@ -0,0 +1,44 @@
+use crate::{EqPoly, FieldExt};
+
+#[derive(Clone, Debug)]
+pub struct SparseMLPoly<F> {
+    pub evals: Vec<(usize, F)>,
+    pub num_vars: usize,
+}
+
+impl<F: FieldExt> SparseMLPoly<F> {
+    pub fn new(evals: Vec<(usize, F)>, num_vars: usize) -> Self {
+        Self { evals, num_vars }
+    }
+
+    pub fn from_dense(dense_evals: Vec<F>) -> Self {
+        let sparse_evals = dense_evals
+            .iter()
+            .filter(|eval| **eval != F::ZERO)
+            .enumerate()
+            .map(|(i, eval)| (i, *eval))
+            .collect::<Vec<(usize, F)>>();
+        let num_vars = (dense_evals.len() as f64).log2() as usize;
+
+        Self {
+            evals: sparse_evals,
+            num_vars,
+        }
+    }
+
+    pub fn eval(&self, t: &[F]) -> F {
+        // Evaluate the multilinear extension of the polynomial `a`,
+        // over the boolean hypercube
+
+        let eq_poly = EqPoly::new(t.to_vec());
+        let eq_evals = eq_poly.evals();
+
+        let mut result = F::ZERO;
+
+        for eval in &self.evals {
+            result += eq_evals[eval.0] * eval.1;
+        }
+
+        result
+    }
+}

tensor_pcs/src/rs_config/ecfft.rs

@@ -0,0 +1,27 @@
+use crate::FieldExt;
+use ecfft::{prepare_domain, prepare_matrices, GoodCurve, Matrix2x2};
+
+#[derive(Clone, Debug)]
+pub struct ECFFTConfig<F: FieldExt> {
+    pub domain: Vec<Vec<F>>,
+    pub matrices: Vec<Vec<Matrix2x2<F>>>,
+    pub inverse_matrices: Vec<Vec<Matrix2x2<F>>>,
+}
+
+pub fn gen_config<F: FieldExt>(num_cols: usize) -> ECFFTConfig<F> {
+    assert!(num_cols.is_power_of_two());
+    let expansion_factor = 2;
+    let codeword_len = num_cols * expansion_factor;
+
+    let k = (codeword_len as f64).log2() as usize;
+
+    let good_curve = GoodCurve::find_k(k);
+    let domain = prepare_domain(good_curve);
+    let (matrices, inverse_matrices) = prepare_matrices(&domain);
+
+    ECFFTConfig {
+        domain,
+        matrices,
+        inverse_matrices,
+    }
+}

tensor_pcs/src/rs_config/mod.rs

@@ -0,0 +1,3 @@
+pub mod ecfft;
+pub mod naive;
+pub mod smooth;

tensor_pcs/src/rs_config/naive.rs

@@ -0,0 +1,21 @@
+use crate::FieldExt;
+
+pub fn gen_config<F: FieldExt>(num_cols: usize) -> Vec<Vec<F>> {
+    assert!(num_cols.is_power_of_two());
+    let expansion_factor = 2;
+    let codeword_len = num_cols * expansion_factor;
+    let domain = (0..codeword_len)
+        .map(|i| F::from((i + 3) as u64))
+        .collect::<Vec<F>>();
+
+    let mut domain_powers = Vec::with_capacity(codeword_len);
+    for eval_at in domain {
+        let mut powers_i = vec![F::ONE];
+        for j in 0..(num_cols - 1) {
+            powers_i.push(powers_i[j] * eval_at);
+        }
+        domain_powers.push(powers_i);
+    }
+
+    domain_powers
+}

tensor_pcs/src/rs_config/smooth.rs

@@ -0,0 +1,23 @@
+use crate::FieldExt;
+
+pub fn gen_config<F: FieldExt>(num_cols: usize) -> Vec<F> {
+    assert!(num_cols.is_power_of_two());
+    let expansion_factor = 2;
+    let codeword_len = num_cols * expansion_factor;
+
+    let domain_generator = F::ROOT_OF_UNITY.pow(&[
+        2u32.pow(32 - ((codeword_len as f64).log2() as u32)) as u64,
+        0,
+        0,
+        0,
+    ]);
+
+    // Compute the FFT domain
+    let mut fft_domain = Vec::with_capacity(codeword_len);
+    fft_domain.push(F::ONE);
+    for i in 0..(codeword_len - 1) {
+        fft_domain.push(fft_domain[i] * domain_generator);
+    }
+
+    fft_domain
+}

tensor_pcs/src/tensor_code.rs

@@ -0,0 +1,42 @@
+use crate::tree::CommittedMerkleTree;
+use crate::FieldExt;
+
+#[derive(Clone)]
+pub struct TensorCode<F>(pub Vec<Vec<F>>)
+where
+    F: FieldExt;
+
+impl<F: FieldExt> TensorCode<F> {
+    pub fn commit(&self, num_cols: usize, num_rows: usize) -> CommittedTensorCode<F> {
+        // Flatten the tensor codeword in column major order
+        let mut tensor_codeword = vec![];
+        for j in 0..(num_cols * 2) {
+            for i in 0..num_rows {
+                tensor_codeword.push(self.0[i][j])
+            }
+        }
+        // Merkle commit the codewords
+        let committed_tree = CommittedMerkleTree::from_leaves(tensor_codeword, num_cols * 2);
+
+        CommittedTensorCode {
+            committed_tree,
+            tensor_codeword: Self(self.0.clone()),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct CommittedTensorCode<F: FieldExt> {
+    pub committed_tree: CommittedMerkleTree<F>,
+    pub tensor_codeword: TensorCode<F>,
+}
+
+impl<F: FieldExt> CommittedTensorCode<F> {
+    pub fn query_column(&self, column: usize, num_cols: usize) -> Vec<F> {
+        let num_rows = self.tensor_codeword.0.len();
+
+        let leaves =
+            self.committed_tree.leaves[column * num_rows..((column + 1) * num_rows)].to_vec();
+        leaves
+    }
+}

tensor_pcs/src/tensor_pcs.rs

@@ -0,0 +1,446 @@
+use crate::rs_config::ecfft::ECFFTConfig;
+use crate::tree::BaseOpening;
+use crate::FieldExt;
+use ecfft::extend;
+use serde::{Deserialize, Serialize};
+
+use crate::fft::fft;
+use crate::polynomial::eq_poly::EqPoly;
+use crate::polynomial::sparse_ml_poly::SparseMLPoly;
+use crate::tensor_code::TensorCode;
+use crate::transcript::Transcript;
+use crate::utils::{dot_prod, hash_all, rlc_rows, sample_indices};
+
+use super::tensor_code::CommittedTensorCode;
+
+#[derive(Clone)]
+pub struct TensorRSMultilinearPCSConfig<F: FieldExt> {
+    pub expansion_factor: usize,
+    pub domain_powers: Option<Vec<Vec<F>>>,
+    pub fft_domain: Option<Vec<F>>,
+    pub ecfft_config: Option<ECFFTConfig<F>>,
+    pub l: usize,
+    pub num_entries: usize,
+    pub num_rows: usize,
+}
+
+impl<F: FieldExt> TensorRSMultilinearPCSConfig<F> {
+    pub fn num_cols(&self) -> usize {
+        self.num_entries / self.num_rows()
+    }
+
+    pub fn num_rows(&self) -> usize {
+        self.num_rows
+    }
+}
+
+#[derive(Clone)]
+pub struct TensorMultilinearPCS<F: FieldExt> {
+    config: TensorRSMultilinearPCSConfig<F>,
+}
+
+#[derive(Clone, Serialize, Deserialize)]
+pub struct TensorMLOpening<F: FieldExt> {
+    pub x: Vec<F>,
+    pub y: F,
+    pub base_opening: BaseOpening,
+    pub test_query_leaves: Vec<Vec<F>>,
+    pub eval_query_leaves: Vec<Vec<F>>,
+    u_hat_comm: [u8; 32],
+    pub test_u_prime: Vec<F>,
+    pub test_r_prime: Vec<F>,
+    pub eval_r_prime: Vec<F>,
+    pub eval_u_prime: Vec<F>,
+}
+
+impl<F: FieldExt> TensorMultilinearPCS<F> {
+    pub fn new(config: TensorRSMultilinearPCSConfig<F>) -> Self {
+        Self { config }
+    }
+
+    pub fn commit(&self, poly: &SparseMLPoly<F>) -> CommittedTensorCode<F> {
+        // Merkle commit to the evaluations of the polynomial
+        let tensor_code = self.encode_zk(poly);
+        let tree = tensor_code.commit(self.config.num_cols(), self.config.num_rows());
+        tree
+    }
+
+    pub fn open(
+        &self,
+        u_hat_comm: &CommittedTensorCode<F>,
+        poly: &SparseMLPoly<F>,
+        point: &[F],
+        transcript: &mut Transcript<F>,
+    ) -> TensorMLOpening<F> {
+        let num_cols = self.config.num_cols();
+        let num_rows = self.config.num_rows();
+        debug_assert_eq!(poly.num_vars, point.len());
+
+        transcript.append_bytes(&u_hat_comm.committed_tree.root());
+
+        // ########################################
+        // Testing phase
+        // Prove the consistency between the random linear combination of the evaluation tensor (u_prime)
+        // and the tensor codeword (u_hat)
+        // ########################################
+
+        // Derive the challenge vector;
+        let r_u = transcript.challenge_vec(num_rows);
+
+        let u = (0..num_rows)
+            .map(|i| {
+                poly.evals[(i * num_cols)..((i + 1) * num_cols)]
+                    .iter()
+                    .map(|entry| entry.1)
+                    .collect::<Vec<F>>()
+            })
+            .collect::<Vec<Vec<F>>>();
+
+        // Random linear combination of the rows of the polynomial in a tensor structure
+        let test_u_prime = rlc_rows(u.clone(), &r_u);
+
+        // Random linear combination of the blinder
+        let blinder = u_hat_comm
+            .tensor_codeword
+            .0
+            .iter()
+            .map(|row| row[(row.len() / 2)..].to_vec())
+            .collect::<Vec<Vec<F>>>();
+
+        debug_assert_eq!(blinder[0].len(), u_hat_comm.tensor_codeword.0[0].len() / 2);
+
+        let test_r_prime = rlc_rows(blinder.clone(), &r_u);
+
+        let num_indices = self.config.l;
+        let indices = sample_indices(num_indices, num_cols * 2, transcript);
+
+        let test_queries = self.test_phase(&indices, &u_hat_comm);
+
+        // ########################################
+        // Evaluation phase
+        // Prove the consistency
+        // ########################################
+
+        // Get the evaluation point
+        let mut point_rev = point.to_vec();
+        point_rev.reverse();
+
+        let log2_num_rows = (num_rows as f64).log2() as usize;
+        let q1 = EqPoly::new(point_rev[0..log2_num_rows].to_vec()).evals();
+
+        let eval_r_prime = rlc_rows(blinder, &q1);
+
+        let eval_u_prime = rlc_rows(u.clone(), &q1);
+
+        let eval_queries = self.test_phase(&indices, &u_hat_comm);
+
+        TensorMLOpening {
+            x: point.to_vec(),
+            y: poly.eval(&point_rev),
+            eval_query_leaves: eval_queries,
+            test_query_leaves: test_queries,
+            u_hat_comm: u_hat_comm.committed_tree.root(),
+            test_u_prime,
+            test_r_prime,
+            eval_r_prime,
+            eval_u_prime,
+            base_opening: BaseOpening {
+                hashes: u_hat_comm.committed_tree.column_roots.clone(),
+            },
+        }
+    }
+}
+
+impl<F: FieldExt> TensorMultilinearPCS<F> {
+    pub fn verify(
+        &self,
+        opening: &TensorMLOpening<F>,
+        commitment: &[u8; 32],
+        transcript: &mut Transcript<F>,
+    ) {
+        let num_rows = self.config.num_rows();
+        let num_cols = self.config.num_cols();
+
+        let u_hat_comm = opening.u_hat_comm;
+        transcript.append_bytes(&u_hat_comm);
+
+        assert_eq!(&u_hat_comm, commitment);
+
+        // Verify the base opening
+
+        let base_opening = &opening.base_opening;
+        base_opening.verify(u_hat_comm);
+
+        // ########################################
+        // Verify test phase
+        // ########################################
+
+        let r_u = transcript.challenge_vec(num_rows);
+
+        let test_u_prime_rs_codeword = self
+            .rs_encode(&opening.test_u_prime)
+            .iter()
+            .zip(opening.test_r_prime.iter())
+            .map(|(c, r)| *c + *r)
+            .collect::<Vec<F>>();
+
+        let num_indices = self.config.l;
+        let indices = sample_indices(num_indices, num_cols * 2, transcript);
+
+        debug_assert_eq!(indices.len(), opening.test_query_leaves.len());
+        for (expected_index, leaves) in indices.iter().zip(opening.test_query_leaves.iter()) {
+            // Verify that the hashes of the leaves equals the corresponding column root
+            let leaf_bytes = leaves
+                .iter()
+                .map(|x| x.to_repr())
+                .collect::<Vec<[u8; 32]>>();
+            let column_root = hash_all(&leaf_bytes);
+            let expected_column_root = base_opening.hashes[*expected_index];
+            assert_eq!(column_root, expected_column_root);
+
+            let mut sum = F::ZERO;
+            for (leaf, r_i) in leaves.iter().zip(r_u.iter()) {
+                sum += *r_i * *leaf;
+            }
+            assert_eq!(sum, test_u_prime_rs_codeword[*expected_index]);
+        }
+
+        // ########################################
+        // Verify evaluation phase
+        // ########################################
+
+        let mut x_rev = opening.x.clone();
+        x_rev.reverse();
+
+        let log2_num_rows = (num_rows as f64).log2() as usize;
+        let q1 = EqPoly::new(x_rev[0..log2_num_rows].to_vec()).evals();
+        let q2 = EqPoly::new(x_rev[log2_num_rows..].to_vec()).evals();
+
+        let eval_u_prime_rs_codeword = self
+            .rs_encode(&opening.eval_u_prime)
+            .iter()
+            .zip(opening.eval_r_prime.iter())
+            .map(|(c, r)| *c + *r)
+            .collect::<Vec<F>>();
+
+        debug_assert_eq!(q1.len(), opening.eval_query_leaves[0].len());
+        debug_assert_eq!(indices.len(), opening.test_query_leaves.len());
+        for (expected_index, leaves) in indices.iter().zip(opening.eval_query_leaves.iter()) {
+            // TODO: Don't need to check the leaves again?
+            // Verify that the hashes of the leaves equals the corresponding column root
+            let leaf_bytes = leaves
+                .iter()
+                .map(|x| x.to_repr())
+                .collect::<Vec<[u8; 32]>>();
+            let column_root = hash_all(&leaf_bytes);
+            let expected_column_root = base_opening.hashes[*expected_index];
+            assert_eq!(column_root, expected_column_root);
+
+            let mut sum = F::ZERO;
+            for (leaf, q1_i) in leaves.iter().zip(q1.iter()) {
+                sum += *q1_i * *leaf;
+            }
+            assert_eq!(sum, eval_u_prime_rs_codeword[*expected_index]);
+        }
+
+        let expected_eval = dot_prod(&opening.eval_u_prime, &q2);
+        assert_eq!(expected_eval, opening.y);
+    }
+
+    fn split_encode(&self, message: &[F]) -> Vec<F> {
+        let codeword = self.rs_encode(message);
+
+        let mut rng = rand::thread_rng();
+        let blinder = (0..codeword.len())
+            .map(|_| F::random(&mut rng))
+            .collect::<Vec<F>>();
+
+        let mut randomized_codeword = codeword
+            .iter()
+            .zip(blinder.clone().iter())
+            .map(|(c, b)| *b + *c)
+            .collect::<Vec<F>>();
+
+        randomized_codeword.extend_from_slice(&blinder);
+        debug_assert_eq!(randomized_codeword.len(), codeword.len() * 2);
+        randomized_codeword
+    }
+
+    fn rs_encode(&self, message: &[F]) -> Vec<F> {
+        let codeword = if self.config.fft_domain.is_some() {
+            let fft_domain = self.config.fft_domain.as_ref().unwrap();
+            let mut padded_coeffs = message.clone().to_vec();
+            padded_coeffs.resize(fft_domain.len(), F::ZERO);
+            let codeword = fft(&padded_coeffs, &fft_domain);
+
+            codeword
+        } else if self.config.ecfft_config.is_some() {
+            let ecfft_config = self.config.ecfft_config.as_ref().unwrap();
+            assert_eq!(
+                message.len() * self.config.expansion_factor,
+                ecfft_config.domain[0].len()
+            );
+            let extended_evals = extend(
+                message,
+                &ecfft_config.domain,
+                &ecfft_config.matrices,
+                &ecfft_config.inverse_matrices,
+                0,
+            );
+
+            let codeword = [message.to_vec(), extended_evals].concat();
+            codeword
+        } else {
+            let domain_powers = self.config.domain_powers.as_ref().unwrap();
+            assert_eq!(message.len(), domain_powers[0].len());
+            assert_eq!(
+                message.len() * self.config.expansion_factor,
+                domain_powers.len()
+            );
+
+            let codeword = domain_powers
+                .iter()
+                .map(|powers| {
+                    message
+                        .iter()
+                        .zip(powers.iter())
+                        .fold(F::ZERO, |acc, (m, p)| acc + *m * *p)
+                })
+                .collect::<Vec<F>>();
+
+            codeword
+        };
+
+        codeword
+    }
+
+    fn test_phase(&self, indices: &[usize], u_hat_comm: &CommittedTensorCode<F>) -> Vec<Vec<F>> {
+        let num_cols = self.config.num_cols() * 2;
+
+        // Query the columns of u_hat
+        let num_indices = self.config.l;
+
+        let u_hat_openings = indices
+            .iter()
+            .map(|index| u_hat_comm.query_column(*index, num_cols))
+            .collect::<Vec<Vec<F>>>();
+
+        debug_assert_eq!(u_hat_openings.len(), num_indices);
+
+        u_hat_openings
+    }
+
+    fn encode_zk(&self, poly: &SparseMLPoly<F>) -> TensorCode<F> {
+        let num_rows = self.config.num_rows();
+        let num_cols = self.config.num_cols();
+
+        let codewords = (0..num_rows)
+            .map(|i| {
+                poly.evals[i * num_cols..(i + 1) * num_cols]
+                    .iter()
+                    .map(|entry| entry.1)
+                    .collect::<Vec<F>>()
+            })
+            .map(|row| self.split_encode(&row))
+            .collect::<Vec<Vec<F>>>();
+
+        TensorCode(codewords)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::rs_config::{ecfft, naive, smooth};
+
+    const TEST_NUM_VARS: usize = 10;
+    const TEST_L: usize = 10;
+
+    fn test_poly<F: FieldExt>() -> SparseMLPoly<F> {
+        let num_entries: usize = 2usize.pow(TEST_NUM_VARS as u32);
+
+        let evals = (0..num_entries)
+            .map(|i| (i, F::from(i as u64)))
+            .collect::<Vec<(usize, F)>>();
+
+        let ml_poly = SparseMLPoly::new(evals, TEST_NUM_VARS);
+        ml_poly
+    }
+
+    fn prove_and_verify<F: FieldExt>(ml_poly: SparseMLPoly<F>, pcs: TensorMultilinearPCS<F>) {
+        let comm = pcs.commit(&ml_poly);
+
+        let open_at = (0..ml_poly.num_vars)
+            .map(|i| F::from(i as u64))
+            .collect::<Vec<F>>();
+
+        let mut prover_transcript = Transcript::<F>::new(b"test");
+        let opening = pcs.open(&comm, &ml_poly, &open_at, &mut prover_transcript);
+
+        let mut verifier_transcript = Transcript::<F>::new(b"test");
+        pcs.verify(
+            &opening,
+            &comm.committed_tree.root(),
+            &mut verifier_transcript,
+        );
+    }
+
+    fn config_base<F: FieldExt>(ml_poly: &SparseMLPoly<F>) -> TensorRSMultilinearPCSConfig<F> {
+        let num_vars = ml_poly.num_vars;
+        let num_evals = 2usize.pow(num_vars as u32);
+        let num_rows = 2usize.pow((num_vars / 2) as u32);
+
+        let expansion_factor = 2;
+
+        TensorRSMultilinearPCSConfig::<F> {
+            expansion_factor,
+            domain_powers: None,
+            fft_domain: None,
+            ecfft_config: None,
+            l: TEST_L,
+            num_entries: num_evals,
+            num_rows,
+        }
+    }
+
+    #[test]
+    fn test_tensor_pcs_fft() {
+        type F = halo2curves::pasta::Fp;
+        // FFT config
+        let ml_poly = test_poly();
+        let mut config = config_base(&ml_poly);
+        config.fft_domain = Some(smooth::gen_config(config.num_cols()));
+
+        // Test FFT PCS
+        let tensor_pcs_fft = TensorMultilinearPCS::<F>::new(config);
+        prove_and_verify(ml_poly, tensor_pcs_fft);
+    }
+
+    #[test]
+    fn test_tensor_pcs_ecfft() {
+        type F = halo2curves::secp256k1::Fp;
+        let ml_poly = test_poly();
+
+        let mut config = config_base(&ml_poly);
+        config.ecfft_config = Some(ecfft::gen_config(config.num_cols()));
+
+        // Test FFT PCS
+        let tensor_pcs_ecf = TensorMultilinearPCS::<F>::new(config);
+        prove_and_verify(ml_poly, tensor_pcs_ecf);
+    }
+
+    #[test]
+    fn test_tensor_pcs_naive() {
+        type F = halo2curves::secp256k1::Fp;
+        // FFT config
+        let ml_poly = test_poly();
+
+        // Naive config
+        let mut config = config_base(&ml_poly);
+        config.domain_powers = Some(naive::gen_config(config.num_cols()));
+
+        // Test FFT PCS
+        let tensor_pcs_naive = TensorMultilinearPCS::<F>::new(config);
+        prove_and_verify(ml_poly, tensor_pcs_naive);
+    }
+}

tensor_pcs/src/transcript.rs

@@ -0,0 +1,58 @@
+use crate::FieldExt;
+use halo2curves::ff::PrimeField;
+use merlin::Transcript as MerlinTranscript;
+use std::{io::Repeat, marker::PhantomData, panic::UnwindSafe};
+
+#[derive(Clone)]
+pub struct Transcript<F: FieldExt> {
+    transcript_inner: MerlinTranscript,
+    _marker: PhantomData<F>,
+}
+
+impl<F: FieldExt> Transcript<F> {
+    pub fn new(label: &'static [u8]) -> Self {
+        Self {
+            transcript_inner: MerlinTranscript::new(label),
+            _marker: PhantomData,
+        }
+    }
+
+    pub fn append_fe(&mut self, fe: &F) {
+        self.transcript_inner.append_message(b"", &fe.to_repr());
+    }
+
+    pub fn append_bytes(&mut self, bytes: &[u8]) {
+        self.transcript_inner.append_message(b"", bytes);
+    }
+
+    pub fn challenge_vec(&mut self, n: usize) -> Vec<F> {
+        (0..n)
+            .map(|_| {
+                let mut bytes = [0u8; 64];
+                self.transcript_inner.challenge_bytes(b"", &mut bytes);
+                F::from_uniform_bytes(&bytes)
+            })
+            .collect()
+    }
+
+    pub fn challenge_fe(&mut self) -> F {
+        // TODO: This is insecure
+        let mut bytes = [0u8; 32];
+        self.transcript_inner.challenge_bytes(b"", &mut bytes);
+        F::from_repr(bytes).unwrap()
+    }
+
+    pub fn challenge_bytes(&mut self, bytes: &mut [u8]) {
+        self.transcript_inner.challenge_bytes(b"", bytes);
+    }
+}
+
+pub trait AppendToTranscript<F: FieldExt> {
+    fn append_to_transcript(&self, transcript: &mut Transcript<F>);
+}
+
+impl<F: FieldExt> AppendToTranscript<F> for [u8; 32] {
+    fn append_to_transcript(&self, transcript: &mut Transcript<F>) {
+        transcript.append_bytes(self);
+    }
+}

tensor_pcs/src/tree.rs

@@ -0,0 +1,64 @@
+use core::num;
+use std::marker::PhantomData;
+
+use super::utils::hash_two;
+use crate::{utils::hash_all, FieldExt};
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone)]
+pub struct CommittedMerkleTree<F> {
+    pub column_roots: Vec<[u8; 32]>,
+    pub leaves: Vec<F>,
+    pub num_cols: usize,
+    pub root: [u8; 32],
+}
+
+impl<F: FieldExt> CommittedMerkleTree<F> {
+    pub fn from_leaves(leaves: Vec<F>, num_cols: usize) -> Self {
+        let n = leaves.len();
+        debug_assert!(n.is_power_of_two());
+        let num_rows = n / num_cols;
+        assert!(num_rows & 1 == 0); // Number of rows must be even
+
+        let leaf_bytes = leaves
+            .iter()
+            .map(|x| x.to_repr())
+            .collect::<Vec<[u8; 32]>>();
+
+        let mut column_roots = Vec::with_capacity(num_cols);
+        for col in 0..num_cols {
+            let column_leaves = leaf_bytes[col * num_rows..(col + 1) * num_rows].to_vec();
+            let column_root = hash_all(&column_leaves);
+            column_roots.push(column_root);
+        }
+
+        let root = hash_all(&column_roots);
+
+        Self {
+            column_roots,
+            leaves,
+            root,
+            num_cols,
+        }
+    }
+
+    pub fn root(&self) -> [u8; 32] {
+        self.root
+    }
+
+    pub fn leaves(&self) -> Vec<F> {
+        self.leaves.clone()
+    }
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug)]
+pub struct BaseOpening {
+    pub hashes: Vec<[u8; 32]>,
+}
+
+impl BaseOpening {
+    pub fn verify(&self, root: [u8; 32]) -> bool {
+        let r = hash_all(&self.hashes);
+        root == r
+    }
+}

tensor_pcs/src/utils.rs

@@ -0,0 +1,89 @@
+use tiny_keccak::{Hasher, Keccak};
+
+use crate::FieldExt;
+
+use crate::transcript::Transcript;
+
+pub fn rlc_rows<F: FieldExt>(x: Vec<Vec<F>>, r: &[F]) -> Vec<F> {
+    debug_assert_eq!(x.len(), r.len());
+    let num_cols = x[0].len();
+    let mut result = vec![F::ZERO; num_cols];
+    for (row, r_i) in x.iter().zip(r.iter()) {
+        for j in 0..num_cols {
+            result[j] += row[j] * r_i
+        }
+    }
+    result
+}
+
+pub fn dot_prod<F: FieldExt>(x: &[F], y: &[F]) -> F {
+    assert_eq!(x.len(), y.len());
+    let mut result = F::ZERO;
+    for i in 0..x.len() {
+        result += x[i] * y[i];
+    }
+    result
+}
+
+pub fn hash_two(values: &[[u8; 32]; 2]) -> [u8; 32] {
+    let mut hasher = Keccak::v256();
+    hasher.update(&values[0]);
+    hasher.update(&values[1]);
+    let mut hash = [0u8; 32];
+    hasher.finalize(&mut hash);
+    hash
+}
+
+pub fn hash_all(values: &[[u8; 32]]) -> [u8; 32] {
+    let mut hasher = Keccak::v256();
+    for value in values {
+        hasher.update(value);
+    }
+    let mut hash = [0u8; 32];
+    hasher.finalize(&mut hash);
+    hash
+}
+
+fn sample_index(random_bytes: [u8; 64], size: usize) -> usize {
+    let mut acc: u64 = 0;
+    for b in random_bytes {
+        acc = acc << 8 ^ (b as u64);
+    }
+
+    (acc % (size as u64)) as usize
+}
+
+pub fn sample_indices<F: FieldExt>(
+    num_indices: usize,
+    max_index: usize,
+    transcript: &mut Transcript<F>,
+) -> Vec<usize> {
+    assert!(
+        num_indices <= max_index,
+        "max_index {:?} num_indices {:?}",
+        max_index,
+        num_indices
+    );
+
+    let mut indices = Vec::with_capacity(num_indices);
+    let mut counter: u32 = 0;
+
+    // TODO: Don't sample at n and n + N
+    while indices.len() < num_indices {
+        let mut random_bytes = [0u8; 64];
+
+        transcript.append_bytes(&counter.to_le_bytes());
+        transcript.challenge_bytes(&mut random_bytes);
+
+        let index = sample_index(random_bytes, max_index);
+        if !indices.contains(&index)
+        // || !indices.contains(&(index + (max_index / 2)))
+        // || !indices.contains(&(index - (max_index / 2)))
+        {
+            indices.push(index);
+        }
+        counter += 1;
+    }
+
+    indices
+}