Add tests that use non scalar ntt witnesses (#130)

We currently have all tests based from scalar ntt witness, in which all
basefield elements are the same
Change this so we can test the folding scheme more thoroughly

latticefold/src/arith/r1cs.rs

@@ -1,3 +1,4 @@
+use cyclotomic_rings::rings::SuitableRing;
 use lattirust_linear_algebra::sparse_matrix::dense_matrix_u64_to_sparse;
 use lattirust_linear_algebra::SparseMatrix;
 use lattirust_ring::Ring;
@@ -158,7 +159,7 @@ pub fn create_dummy_identity_sparse_matrix<R: Ring>(
 }
 
 pub fn get_test_z<R: Ring>(input: usize) -> Vec<R> {
-    // z = (1, io, w)
+    // z = (io, 1, w)
     to_F_vec(vec![
         input, // io
         1,
@@ -169,20 +170,40 @@ pub fn get_test_z<R: Ring>(input: usize) -> Vec<R> {
     ])
 }
 
-pub fn get_test_z_split<R: Ring>(input: usize) -> (R, Vec<R>, Vec<R>) {
-    // z = (1, io, w)
-    (
-        R::one(),
-        to_F_vec(vec![
+pub fn get_test_z_ntt<R: SuitableRing>() -> Vec<R> {
+    let mut res = Vec::new();
+    for input in 0..R::dimension() {
+        // z = (io, 1, w)
+        res.push(to_F_vec::<R::BaseRing>(vec![
             input, // io
-        ]),
-        to_F_vec(vec![
+            1,
             input * input * input + input + 5, // x^3 + x + 5
             input * input,                     // x^2
             input * input * input,             // x^2 * x
             input * input * input + input,     // x^3 + x
-        ]),
-    )
+        ]))
+    }
+
+    let mut ret: Vec<R> = Vec::new();
+    for j in 0..res[0].len() {
+        let mut vec = Vec::new();
+        for witness in &res {
+            vec.push(witness[j]);
+        }
+        ret.push(R::from(vec));
+    }
+
+    ret
+}
+
+pub fn get_test_z_split<R: Ring>(input: usize) -> (R, Vec<R>, Vec<R>) {
+    let z = get_test_z(input);
+    (z[1], vec![z[0]], z[2..].to_vec())
+}
+
+pub fn get_test_z_ntt_split<R: SuitableRing>() -> (R, Vec<R>, Vec<R>) {
+    let z = get_test_z_ntt();
+    (z[1], vec![z[0]], z[2..].to_vec())
 }
 
 pub fn get_test_dummy_z_split<R: Ring, const X_LEN: usize, const WIT_LEN: usize>(

latticefold/src/decomposition_parameters.rs

@@ -97,4 +97,13 @@ pub mod test_params {
         const B_SMALL: usize = 2;
         const K: usize = 8;
     }
+    #[derive(Clone)]
+    pub struct FrogDP;
+
+    impl DecompositionParams for FrogDP {
+        const B: u128 = 1 << 8;
+        const L: usize = 8;
+        const B_SMALL: usize = 2;
+        const K: usize = 10;
+    }
 }

latticefold/src/nifs/folding/tests/mod.rs

@@ -1,6 +1,8 @@
 use crate::arith::{CCS, LCCCS};
 use crate::ark_base::Vec;
-use crate::decomposition_parameters::test_params::{BabyBearDP, StarkFoldingDP, DP};
+use crate::decomposition_parameters::test_params::{
+    BabyBearDP, FrogDP, GoldilocksDP, StarkFoldingDP,
+};
 use crate::nifs::folding::utils::SqueezeAlphaBetaZetaMu;
 use crate::nifs::folding::{
     prepare_public_output,
@@ -13,7 +15,7 @@ use crate::transcript::{Transcript, TranscriptWithShortChallenges};
 use crate::utils::sumcheck::prover::ProverState;
 use crate::utils::sumcheck::{virtual_polynomial::VPAuxInfo, MLSumcheck};
 use crate::{
-    arith::{r1cs::get_test_z_split, tests::get_test_ccs, Witness, CCCS},
+    arith::{r1cs::get_test_z_ntt_split, tests::get_test_ccs, Witness, CCCS},
     commitment::AjtaiCommitmentScheme,
     decomposition_parameters::DecompositionParams,
     nifs::{
@@ -45,7 +47,6 @@ use lattirust_ring::cyclotomic_ring::models::{
 use lattirust_ring::cyclotomic_ring::{CRT, ICRT};
 use lattirust_ring::Ring;
 use num_traits::{One, Zero};
-use rand::Rng;
 
 const C: usize = 4;
 const WIT_LEN: usize = 4;
@@ -66,8 +67,10 @@ where
     DP: DecompositionParams,
 {
     let ccs = get_test_ccs::<RqNTT>(W, DP::L);
+
     let mut rng = test_rng();
-    let (_, x_ccs, w_ccs) = get_test_z_split::<RqNTT>(rng.gen_range(0..64));
+    let (_, x_ccs, w_ccs) = get_test_z_ntt_split::<RqNTT>();
+
     let scheme = AjtaiCommitmentScheme::rand(&mut rng);
 
     let wit = Witness::from_w_ccs::<DP>(w_ccs);
@@ -129,7 +132,7 @@ where
     };
 
     let folding_proof = if generate_proof {
-        Some(generate_folding_proof(
+        Some(generate_folding_proof::<_, _, C, DP>(
             &ccs,
             &mut prover_transcript,
             &lcccs,
@@ -150,7 +153,7 @@ where
     )
 }
 
-fn generate_folding_proof<RqNTT, CS, const C: usize>(
+fn generate_folding_proof<RqNTT, CS, const C: usize, DP>(
     ccs: &CCS<RqNTT>,
     prover_transcript: &mut PoseidonTranscript<RqNTT, CS>,
     lcccs: &[LCCCS<C, RqNTT>],
@@ -160,6 +163,7 @@ fn generate_folding_proof<RqNTT, CS, const C: usize>(
 where
     RqNTT: SuitableRing,
     CS: LatticefoldChallengeSet<RqNTT>,
+    DP: DecompositionParams,
 {
     let (_, _, folding_proof) = LFFoldingProver::<RqNTT, PoseidonTranscript<RqNTT, CS>>::prove::<
         C,
@@ -655,6 +659,7 @@ fn test_full_prove() {
 fn test_proof_serialization() {
     type RqNTT = GoldilocksRqNTT;
     type CS = GoldilocksChallengeSet;
+    type DP = GoldilocksDP;
     const W: usize = WIT_LEN * DP::L;
 
     let (_, _, _, _, proof, _) = setup_test_environment::<RqNTT, CS, DP, C, W>(true);
@@ -680,6 +685,7 @@ fn test_proof_serialization() {
 fn test_verify_evaluation() {
     type RqNTT = GoldilocksRqNTT;
     type CS = GoldilocksChallengeSet;
+    type DP = GoldilocksDP;
     const W: usize = WIT_LEN * DP::L;
 
     let (lccs_vec, _, mut transcript, ccs, proof, _) =
@@ -723,6 +729,7 @@ fn test_verify_evaluation() {
 fn test_calculate_claims() {
     type RqNTT = GoldilocksRqNTT;
     type CS = GoldilocksChallengeSet;
+    type DP = GoldilocksDP;
     const W: usize = WIT_LEN * DP::L;
 
     let (lcccs_vec, _, _, _, _, _) = setup_test_environment::<RqNTT, CS, DP, C, W>(false);
@@ -756,6 +763,8 @@ fn test_calculate_claims() {
 fn test_verify_sumcheck_proof() {
     type RqNTT = FrogRqNTT;
     type CS = FrogChallengeSet;
+    type DP = FrogDP;
+
     const W: usize = WIT_LEN * DP::L;
 
     let (lcccs_vec, _, mut transcript, ccs, proof, _) =
@@ -791,7 +800,7 @@ fn test_verify_sumcheck_proof() {
 fn test_full_verify() {
     type RqNTT = GoldilocksRqNTT;
     type CS = GoldilocksChallengeSet;
-
+    type DP = GoldilocksDP;
     const W: usize = WIT_LEN * DP::L;
 
     let (lccs_vec, _, mut transcript, ccs, proof, _) =