From e695d84bf46572e32a37d091d23e688c5135d8ef Mon Sep 17 00:00:00 2001
From: Ethan Cemer <tylercemer@gmail.com>
Date: Thu, 1 Jun 2023 08:41:37 -0500
Subject: [PATCH] fix: add public/public solidity tests (#278)

---
 src/eth.rs                 |  8 ++++++--
 tests/integration_tests.rs | 16 ++++++++++++----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/eth.rs b/src/eth.rs
index b51b85a00..2d3719a30 100644
--- a/src/eth.rs
+++ b/src/eth.rs
@@ -87,7 +87,9 @@ pub async fn verify_proof_via_solidity(
     let contract = Verifier::new(addr, client.clone());
 
     let mut public_inputs = vec![];
-    for val in &proof.instances[0] {
+    let flattened_instances = proof.instances.into_iter().flatten();
+    
+    for val in flattened_instances {
         let bytes = val.to_repr();
         let u = U256::from_little_endian(bytes.as_slice());
         public_inputs.push(u);
@@ -306,7 +308,9 @@ pub async fn send_proof<M: 'static + Middleware>(
         let contract = Verifier::new(addr, client.clone());
 
         let mut public_inputs = vec![];
-        for val in &snark.instances[0] {
+        let flattened_instances = snark.instances.into_iter().flatten();
+
+        for val in flattened_instances {
             let bytes = val.to_repr();
             let u = U256::from_little_endian(bytes.as_slice());
             public_inputs.push(u);
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index a972f7452..6bdf718f7 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -382,7 +382,9 @@ mod native_tests {
                 fn kzg_evm_prove_and_verify_(test: &str) {
                     crate::native_tests::init_binary();
                     crate::native_tests::init_params_17();
-                    kzg_evm_prove_and_verify(test.to_string(), TESTS_SOLIDITY.contains(&test));
+                    kzg_evm_prove_and_verify(test.to_string(), TESTS_SOLIDITY.contains(&test), false, true);
+                    kzg_evm_prove_and_verify(test.to_string(), TESTS_SOLIDITY.contains(&test), true, true);
+                    kzg_evm_prove_and_verify(test.to_string(), TESTS_SOLIDITY.contains(&test), true, false);
                 }
 
                 #(#[test_case(TESTS_EVM[N])])*
@@ -397,7 +399,9 @@ mod native_tests {
                 fn kzg_evm_aggr_prove_and_verify_(test: &str) {
                     crate::native_tests::init_binary();
                     crate::native_tests::init_params_23();
-                    kzg_evm_aggr_prove_and_verify(test.to_string(), TESTS_SOLIDITY.contains(&test));
+                    kzg_evm_aggr_prove_and_verify(test.to_string(), TESTS_SOLIDITY.contains(&test), false, true);
+                    kzg_evm_aggr_prove_and_verify(test.to_string(), TESTS_SOLIDITY.contains(&test), true, true);
+                    kzg_evm_aggr_prove_and_verify(test.to_string(), TESTS_SOLIDITY.contains(&test), true, false);
                 }
 
             });
@@ -791,7 +795,7 @@ mod native_tests {
     }
 
     // prove-serialize-verify, the usual full path
-    fn kzg_evm_aggr_prove_and_verify(example_name: String, with_solidity: bool) {
+    fn kzg_evm_aggr_prove_and_verify(example_name: String, with_solidity: bool, public_inputs: bool, public_outputs: bool) {
         let status = Command::new(format!("{}/release/ezkl", *CARGO_TARGET_DIR))
             .args([
                 "setup",
@@ -822,6 +826,8 @@ mod native_tests {
                 ),
                 "--bits=16",
                 "-K=17",
+                &format!("--public-inputs={}", public_inputs),
+                &format!("--public-outputs={}", public_outputs)
             ])
             .status()
             .expect("failed to execute process");
@@ -1079,7 +1085,7 @@ mod native_tests {
     }
 
     // prove-serialize-verify, the usual full path
-    fn kzg_evm_prove_and_verify(example_name: String, with_solidity: bool) {
+    fn kzg_evm_prove_and_verify(example_name: String, with_solidity: bool, public_inputs: bool, public_outputs: bool) {
         let status = Command::new(format!("{}/release/ezkl", *CARGO_TARGET_DIR))
             .args([
                 "setup",
@@ -1102,6 +1108,8 @@ mod native_tests {
                 ),
                 "--bits=16",
                 "-K=17",
+                &format!("--public-inputs={}", public_inputs),
+                &format!("--public-outputs={}", public_outputs)
             ])
             .status()
             .expect("failed to execute process");