From 5f615d7fb30751c0e5a246c1db5357242123fd19 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Garillot?=
 <4142+huitseeker@users.noreply.github.com>
Date: Thu, 21 Mar 2024 20:49:52 -0400
Subject: [PATCH] refactor: Refactor Supernova test suite and streamline
 functions (#372)

- Refactored test functions within the supernova snark module to reduce redundancy.
- Consolidated `test_compression_with` and `test_compression_with_circuit_size_difference_with` functions into a single function.
- Introduced a 'circuit factory' function to generate test circuits.
- Removed `test_nivc_trivial_with_compression_with` method, unified under `test_compression_with`.
- Adjusted individual test cases to conform to new structured changes.
---
 src/supernova/snark.rs | 101 +++++++++++++----------------------------
 1 file changed, 31 insertions(+), 70 deletions(-)

diff --git a/src/supernova/snark.rs b/src/supernova/snark.rs
index 8645ee0ec..ab7b55c60 100644
--- a/src/supernova/snark.rs
+++ b/src/supernova/snark.rs
@@ -458,65 +458,6 @@ mod test {
     }
   }
 
-  fn test_nivc_trivial_with_compression_with<E1, S1, S2>()
-  where
-    E1: CurveCycleEquipped,
-    S1: BatchedRelaxedR1CSSNARKTrait<E1>,
-    S2: RelaxedR1CSSNARKTrait<Dual<E1>>,
-    <E1::Scalar as PrimeField>::Repr: Abomonation,
-    <<Dual<E1> as Engine>::Scalar as PrimeField>::Repr: Abomonation,
-  {
-    const NUM_STEPS: usize = 6;
-
-    let secondary_circuit = TrivialSecondaryCircuit::default();
-    let test_circuits = TestCircuit::new(NUM_STEPS);
-
-    let pp = PublicParams::setup(&test_circuits[0], &*S1::ck_floor(), &*S2::ck_floor());
-
-    let z0_primary = vec![E1::Scalar::from(17u64)];
-    let z0_secondary = vec![<Dual<E1> as Engine>::Scalar::ZERO];
-
-    let mut recursive_snark = RecursiveSNARK::new(
-      &pp,
-      &test_circuits[0],
-      &test_circuits[0],
-      &secondary_circuit,
-      &z0_primary,
-      &z0_secondary,
-    )
-    .unwrap();
-
-    for circuit in test_circuits.iter().take(NUM_STEPS) {
-      recursive_snark
-        .prove_step(&pp, circuit, &secondary_circuit)
-        .unwrap();
-
-      recursive_snark
-        .verify(&pp, &z0_primary, &z0_secondary)
-        .unwrap();
-    }
-
-    let (prover_key, verifier_key) = CompressedSNARK::<_, S1, S2>::setup(&pp).unwrap();
-
-    let compressed_snark = CompressedSNARK::prove(&pp, &prover_key, &recursive_snark).unwrap();
-
-    compressed_snark
-      .verify(&pp, &verifier_key, &z0_primary, &z0_secondary)
-      .unwrap();
-  }
-
-  #[test]
-  fn test_nivc_trivial_with_compression() {
-    // ppSNARK
-    test_nivc_trivial_with_compression_with::<PallasEngine, S1PP<_>, S2<_>>();
-    test_nivc_trivial_with_compression_with::<Bn256EngineIPA, S1PP<_>, S2<_>>();
-    test_nivc_trivial_with_compression_with::<Secp256k1Engine, S1PP<_>, S2<_>>();
-    // classic SNARK
-    test_nivc_trivial_with_compression_with::<PallasEngine, S1<_>, S2<_>>();
-    test_nivc_trivial_with_compression_with::<Bn256EngineIPA, S1<_>, S2<_>>();
-    test_nivc_trivial_with_compression_with::<Secp256k1Engine, S1<_>, S2<_>>();
-  }
-
   #[derive(Clone)]
   struct BigPowerCircuit<E> {
     _p: PhantomData<E>,
@@ -637,18 +578,19 @@ mod test {
     }
   }
 
-  fn test_compression_with_circuit_size_difference_with<E1, S1, S2>()
+  fn test_compression_with<E1, S1, S2, F, C>(num_steps: usize, circuits_factory: F)
   where
     E1: CurveCycleEquipped,
     S1: BatchedRelaxedR1CSSNARKTrait<E1>,
     S2: RelaxedR1CSSNARKTrait<Dual<E1>>,
     <E1::Scalar as PrimeField>::Repr: Abomonation,
     <<Dual<E1> as Engine>::Scalar as PrimeField>::Repr: Abomonation,
+    C: NonUniformCircuit<E1, C1 = C, C2 = TrivialSecondaryCircuit<<Dual<E1> as Engine>::Scalar>>
+      + StepCircuit<E1::Scalar>,
+    F: Fn(usize) -> Vec<C>,
   {
-    const NUM_STEPS: usize = 4;
-
     let secondary_circuit = TrivialSecondaryCircuit::default();
-    let test_circuits = BigTestCircuit::new(NUM_STEPS);
+    let test_circuits = circuits_factory(num_steps);
 
     let pp = PublicParams::setup(&test_circuits[0], &*S1::ck_floor(), &*S2::ck_floor());
 
@@ -665,7 +607,7 @@ mod test {
     )
     .unwrap();
 
-    for circuit in test_circuits.iter().take(NUM_STEPS) {
+    for circuit in test_circuits.iter().take(num_steps) {
       recursive_snark
         .prove_step(&pp, circuit, &secondary_circuit)
         .unwrap();
@@ -684,15 +626,34 @@ mod test {
       .unwrap();
   }
 
+  #[test]
+  fn test_nivc_trivial_with_compression() {
+    const NUM_STEPS: usize = 6;
+
+    // ppSNARK
+    test_compression_with::<PallasEngine, S1PP<_>, S2<_>, _, _>(NUM_STEPS, TestCircuit::new);
+
+    test_compression_with::<Bn256EngineIPA, S1PP<_>, S2<_>, _, _>(NUM_STEPS, TestCircuit::new);
+    test_compression_with::<Secp256k1Engine, S1PP<_>, S2<_>, _, _>(NUM_STEPS, TestCircuit::new);
+
+    // classic SNARK
+    test_compression_with::<PallasEngine, S1<_>, S2<_>, _, _>(NUM_STEPS, TestCircuit::new);
+    test_compression_with::<Bn256EngineIPA, S1<_>, S2<_>, _, _>(NUM_STEPS, TestCircuit::new);
+    test_compression_with::<Secp256k1Engine, S1<_>, S2<_>, _, _>(NUM_STEPS, TestCircuit::new);
+  }
+
   #[test]
   fn test_compression_with_circuit_size_difference() {
+    const NUM_STEPS: usize = 4;
+
     // ppSNARK
-    test_compression_with_circuit_size_difference_with::<PallasEngine, S1PP<_>, S2<_>>();
-    test_compression_with_circuit_size_difference_with::<Bn256EngineIPA, S1PP<_>, S2<_>>();
-    test_compression_with_circuit_size_difference_with::<Secp256k1Engine, S1PP<_>, S2<_>>();
+    test_compression_with::<PallasEngine, S1PP<_>, S2<_>, _, _>(NUM_STEPS, BigTestCircuit::new);
+    test_compression_with::<Bn256EngineIPA, S1PP<_>, S2<_>, _, _>(NUM_STEPS, BigTestCircuit::new);
+    test_compression_with::<Secp256k1Engine, S1PP<_>, S2<_>, _, _>(NUM_STEPS, BigTestCircuit::new);
+
     // classic SNARK
-    test_compression_with_circuit_size_difference_with::<PallasEngine, S1<_>, S2<_>>();
-    test_compression_with_circuit_size_difference_with::<Bn256EngineIPA, S1<_>, S2<_>>();
-    test_compression_with_circuit_size_difference_with::<Secp256k1Engine, S1<_>, S2<_>>();
+    test_compression_with::<PallasEngine, S1<_>, S2<_>, _, _>(NUM_STEPS, BigTestCircuit::new);
+    test_compression_with::<Bn256EngineIPA, S1<_>, S2<_>, _, _>(NUM_STEPS, BigTestCircuit::new);
+    test_compression_with::<Secp256k1Engine, S1<_>, S2<_>, _, _>(NUM_STEPS, BigTestCircuit::new);
   }
 }