From f07f2f6fa9d82da02fecbdf8b456b6bea79bc0fb Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Wed, 27 Dec 2023 22:10:23 -0500
Subject: [PATCH] feat: dummy snark should have correct `agg_vkey_hash`

if it is for a universal aggregation circuit
---
 Cargo.lock                            |  6 ++--
 snark-verifier-sdk/src/halo2/utils.rs | 52 +++++++++++++++++++--------
 2 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 251a3944..9f7a723d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1136,7 +1136,7 @@ checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7"
 
 [[package]]
 name = "halo2-axiom"
-version = "0.4.1"
+version = "0.4.2"
 dependencies = [
  "blake2b_simd",
  "crossbeam",
@@ -1230,9 +1230,9 @@ dependencies = [
 
 [[package]]
 name = "halo2curves-axiom"
-version = "0.4.4"
+version = "0.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d1f0b8f4b7dc27d1e438dae6573e8d95c9f6ee568d1d21146a4c34c17db2ad38"
+checksum = "4a11fd188635e09153d007ac33300d0af76972342b2f5ca3f116cba8c759f57d"
 dependencies = [
  "blake2b_simd",
  "ff",
diff --git a/snark-verifier-sdk/src/halo2/utils.rs b/snark-verifier-sdk/src/halo2/utils.rs
index d5719cd1..d7588860 100644
--- a/snark-verifier-sdk/src/halo2/utils.rs
+++ b/snark-verifier-sdk/src/halo2/utils.rs
@@ -14,6 +14,8 @@ pub struct AggregationDependencyIntent<'a> {
     pub vk: &'a VerifyingKey<G1Affine>,
     pub num_instance: &'a [usize],
     pub accumulator_indices: Option<&'a [(usize, usize)]>,
+    /// See [`AggregationDependencyIntentOwned::agg_vk_hash_data`].
+    pub agg_vk_hash_data: Option<((usize, usize), Fr)>,
 }
 
 #[derive(Clone, Debug)]
@@ -21,6 +23,8 @@ pub struct AggregationDependencyIntentOwned {
     pub vk: VerifyingKey<G1Affine>,
     pub num_instance: Vec<usize>,
     pub accumulator_indices: Option<Vec<(usize, usize)>>,
+    /// If this dependency is itself from a universal aggregation circuit, this should contain (index, agg_vkey_hash), where `index = (i,j)` is the pair recording that the agg_vkey_hash is located at `instances[i][j]`.
+    pub agg_vk_hash_data: Option<((usize, usize), Fr)>,
 }
 
 /// This trait should be implemented on the minimal circuit configuration data necessary to
@@ -56,22 +60,39 @@ pub trait KeygenAggregationCircuitIntent {
         Self: Sized,
     {
         let mut rng = StdRng::seed_from_u64(0u64);
-        let snarks = self
-            .intent_of_dependencies()
-            .into_iter()
-            .map(|AggregationDependencyIntent { vk, num_instance, accumulator_indices }| {
-                let k = vk.get_domain().k();
-                // In KZG `gen_dummy_snark_from_vk` calls `compile`, which only uses `params` for `params.k()` so we can just use a random untrusted setup.
-                // Moreover since this is a dummy snark, the trusted setup shouldn't matter.
-                let params = ParamsKZG::setup(k, &mut rng);
-                gen_dummy_snark_from_vk::<SHPLONK>(
-                    &params,
-                    vk,
-                    num_instance.to_vec(),
-                    accumulator_indices.map(|v| v.to_vec()),
+        let snarks =
+            self.intent_of_dependencies()
+                .into_iter()
+                .map(
+                    |AggregationDependencyIntent {
+                         vk,
+                         num_instance,
+                         accumulator_indices,
+                         agg_vk_hash_data,
+                     }| {
+                        let k = vk.get_domain().k();
+                        // In KZG `gen_dummy_snark_from_vk` calls `compile`, which only uses `params` for `params.k()` so we can just use a random untrusted setup.
+                        // Moreover since this is a dummy snark, the trusted setup shouldn't matter.
+                        let params = ParamsKZG::setup(k, &mut rng);
+                        let mut snark = gen_dummy_snark_from_vk::<SHPLONK>(
+                            &params,
+                            vk,
+                            num_instance.to_vec(),
+                            accumulator_indices.map(|v| v.to_vec()),
+                        );
+                        // We set the current agg_vk_hash in the dummy snark so that the final agg_vk_hash is correct at the end of keygen.
+                        if let Some(((i, j), agg_vk_hash)) = agg_vk_hash_data {
+                            assert!(
+                            i < snark.instances.len(),
+                            "Invalid agg_vk_hash index: ({i},{j}), num_instance: {num_instance:?}");
+                            assert!(j < snark.instances[i].len(),
+                            "Invalid agg_vk_hash index: ({i},{j}), num_instance: {num_instance:?}");
+                            snark.instances[i][j] = agg_vk_hash;
+                        }
+                        snark
+                    },
                 )
-            })
-            .collect();
+                .collect();
         self.build_keygen_circuit_from_snarks(snarks)
     }
 }
@@ -82,6 +103,7 @@ impl<'a> From<&'a AggregationDependencyIntentOwned> for AggregationDependencyInt
             vk: &intent.vk,
             num_instance: &intent.num_instance,
             accumulator_indices: intent.accumulator_indices.as_deref(),
+            agg_vk_hash_data: intent.agg_vk_hash_data,
         }
     }
 }