axiom_query/keygen/agg/
subquery_agg.rs

1use std::{collections::BTreeMap, path::Path, sync::Arc};
2
3use axiom_eth::{
4    halo2_base::gates::circuit::CircuitBuilderStage,
5    halo2_proofs::{plonk::ProvingKey, poly::kzg::commitment::ParamsKZG},
6    halo2curves::bn256::{Bn256, Fr, G1Affine},
7    snark_verifier::verifier::plonk::PlonkProtocol,
8    snark_verifier_sdk::{
9        halo2::{
10            aggregation::AggregationConfigParams,
11            utils::{
12                AggregationDependencyIntent, AggregationDependencyIntentOwned,
13                KeygenAggregationCircuitIntent,
14            },
15        },
16        Snark,
17    },
18    utils::{
19        build_utils::{
20            aggregation::get_dummy_aggregation_params,
21            keygen::{compile_agg_dep_to_protocol, read_srs_from_dir},
22            pinning::aggregation::{AggTreeId, GenericAggPinning},
23        },
24        snark_verifier::EnhancedSnark,
25    },
26};
27use itertools::Itertools;
28use serde::{Deserialize, Serialize};
29use serde_with::serde_as;
30
31use crate::{
32    keygen::{ProvingKeySerializer, SupportedPinning},
33    subquery_aggregation::types::{
34        InputSubqueryAggregation, SUBQUERY_AGGREGATION_AGG_VKEY_HASH_IDX,
35    },
36};
37
38use super::{
39    common::{parse_agg_intent, ForceBasicConfigParams},
40    impl_keygen_intent_for_aggregation, impl_pkey_serializer_for_aggregation,
41    single_type::SupportedIntentTreeSingleType,
42};
43
44/// ** !! IMPORTANT !! **
45/// Do not change the order of this enum, which determines how inputs are parsed.
46// Determines order of circuit IDs in `to_agg`
47#[derive(Clone, Copy, Serialize, Deserialize, Debug, PartialEq, Eq, PartialOrd, Ord)]
48pub enum SubqueryAggInputSnark {
49    Header,
50    Account,
51    Storage,
52    Tx,
53    Receipt,
54    SolidityMapping,
55    ResultsRoot,
56}
57
58#[serde_as]
59#[derive(Clone, Serialize, Deserialize, Debug)]
60pub struct SubqueryAggParams {
61    /// The compiled verification keys of the dependency circuits to aggregate.
62    /// Since Subquery Aggregation is universal aggregation, we remove the `domain` and `preprocessed` from `PlonkProtocol` since those
63    /// are loaded as witness es.
64    #[serde_as(as = "BTreeMap<_, axiom_eth::utils::snark_verifier::Base64Bytes>")]
65    pub to_agg: BTreeMap<SubqueryAggInputSnark, PlonkProtocol<G1Affine>>,
66    pub agg_params: AggregationConfigParams,
67}
68
69pub type SubqueryAggPinning = GenericAggPinning<SubqueryAggParams>;
70
71/// Only implements [ProvingKeySerializer] and not [KeygenCircuitIntent].
72#[derive(Serialize, Deserialize)]
73pub struct RecursiveSubqueryAggIntent {
74    pub deps: Vec<SupportedIntentTreeSingleType>,
75    pub k: u32,
76    /// For different versions of this circuit to be aggregated by the same universal aggregation circuit,
77    /// we may wish to force configure the circuit to have a certain number of columns without auto-configuration.
78    #[serde(skip_serializing_if = "Option::is_none")]
79    pub force_params: Option<ForceBasicConfigParams>,
80}
81
82/// Non-recursive intent. Currently only used internally as an intermediary for recursive intent.
83/// This will implement [KeygenCircuitIntent] where the pinning is not "wrapped" into an enum.
84/// The pinning type is `GenericAggPinning<SubqueryAggParams>`.
85#[derive(Clone, Debug)]
86struct SubqueryAggIntent {
87    // This is from bad UX; only svk = kzg_params.get_g()[0] is used
88    pub kzg_params: Arc<ParamsKZG<Bn256>>,
89    /// For passing to SubqueryAggParams via macro
90    pub to_agg: BTreeMap<SubqueryAggInputSnark, PlonkProtocol<G1Affine>>,
91    /// Circuit intents for subset of circuit types that are enabled
92    pub deps: BTreeMap<SubqueryAggInputSnark, (AggTreeId, AggregationDependencyIntentOwned)>,
93    /// The log_2 domain size of the current aggregation circuit
94    pub k: u32,
95    /// For different versions of this circuit to be aggregated by the same universal aggregation circuit,
96    /// we may wish to force configure the circuit to have a certain number of columns without auto-configuration.
97    pub force_params: Option<ForceBasicConfigParams>,
98}
99
100impl SubqueryAggIntent {
101    pub fn children(&self) -> Vec<AggTreeId> {
102        self.deps.values().map(|(tree, _)| tree.clone()).collect()
103    }
104}
105
106impl KeygenAggregationCircuitIntent for SubqueryAggIntent {
107    fn intent_of_dependencies(&self) -> Vec<AggregationDependencyIntent> {
108        self.deps.values().map(|(_, d)| d.into()).collect()
109    }
110    fn build_keygen_circuit_from_snarks(self, snarks: Vec<Snark>) -> Self::AggregationCircuit {
111        let mut deps = BTreeMap::from_iter(self.deps.keys().cloned().zip_eq(snarks));
112        let mut remove_and_wrap =
113            |k: &SubqueryAggInputSnark| deps.remove(k).map(|s| EnhancedSnark::new(s, None));
114        // TODO: don't do manual conversion
115        let snark_header = remove_and_wrap(&SubqueryAggInputSnark::Header);
116        let snark_results_root = remove_and_wrap(&SubqueryAggInputSnark::ResultsRoot);
117        let snark_account = remove_and_wrap(&SubqueryAggInputSnark::Account);
118        let snark_storage = remove_and_wrap(&SubqueryAggInputSnark::Storage);
119        let snark_tx = remove_and_wrap(&SubqueryAggInputSnark::Tx);
120        let snark_receipt = remove_and_wrap(&SubqueryAggInputSnark::Receipt);
121        let snark_solidity_mapping = remove_and_wrap(&SubqueryAggInputSnark::SolidityMapping);
122        let promise_commit_keccak = Fr::zero(); // just a dummy
123
124        let input = InputSubqueryAggregation {
125            snark_header: snark_header.unwrap(),
126            snark_account,
127            snark_storage,
128            snark_solidity_mapping,
129            snark_tx,
130            snark_receipt,
131            promise_commit_keccak,
132            snark_results_root: snark_results_root.unwrap(),
133        };
134        let mut force = false;
135        let agg_params = if let Some(force_params) = self.force_params {
136            force = true;
137            force_params.into_agg_params(self.k)
138        } else {
139            get_dummy_aggregation_params(self.k as usize)
140        };
141        let mut circuit =
142            input.build(CircuitBuilderStage::Keygen, agg_params, &self.kzg_params).unwrap();
143        if !force {
144            circuit.calculate_params(Some(20));
145        }
146        circuit
147    }
148}
149
150impl_keygen_intent_for_aggregation!(
151    SubqueryAggIntent,
152    SubqueryAggParams,
153    Some(SUBQUERY_AGGREGATION_AGG_VKEY_HASH_IDX)
154);
155impl_pkey_serializer_for_aggregation!(SubqueryAggIntent, SubqueryAggParams, SubqueryAggregation);
156
157fn get_key(supported: &SupportedIntentTreeSingleType) -> SubqueryAggInputSnark {
158    type S = SupportedIntentTreeSingleType;
159    type I = SubqueryAggInputSnark;
160    match supported {
161        S::Header(_) => I::Header,
162        S::Account(_) => I::Account,
163        S::Storage(_) => I::Storage,
164        S::Tx(_) => I::Tx,
165        S::Receipt(_) => I::Receipt,
166        S::SolidityMapping(_) => I::SolidityMapping,
167        S::ResultsRoot(_) => I::ResultsRoot,
168    }
169}
170
171impl ProvingKeySerializer for RecursiveSubqueryAggIntent {
172    fn create_and_serialize_proving_key(
173        self,
174        params_dir: &Path,
175        data_dir: &Path,
176    ) -> anyhow::Result<(AggTreeId, ProvingKey<G1Affine>, SupportedPinning)> {
177        let mut deps = BTreeMap::new();
178        for intent in self.deps {
179            let key = get_key(&intent);
180            let (child_tree_id, pk, pinning) =
181                intent.create_and_serialize_proving_key(params_dir, data_dir)?;
182            let intent = parse_agg_intent(pk.get_vk(), pinning);
183            assert!(deps.insert(key, (child_tree_id, intent)).is_none());
184        }
185        let kzg_params = Arc::new(read_srs_from_dir(params_dir, self.k)?);
186        let to_agg = deps
187            .iter()
188            .map(|(&k, (_, dep))| (k, compile_agg_dep_to_protocol(&kzg_params, dep, true)))
189            .collect();
190        let subquery_agg_intent = SubqueryAggIntent {
191            to_agg,
192            deps,
193            k: self.k,
194            kzg_params,
195            force_params: self.force_params,
196        };
197        subquery_agg_intent.create_and_serialize_proving_key(params_dir, data_dir)
198    }
199}
axiom_query/keygen/agg/subquery_agg.rs

axiom_query/keygen/agg/
subquery_agg.rs
