axiom_core/keygen/
mod.rs

1use std::{collections::BTreeMap, path::Path, sync::Arc};
2
3use axiom_eth::{
4    block_header::GENESIS_BLOCK_RLP,
5    halo2_base::{
6        gates::circuit::CircuitBuilderStage,
7        utils::halo2::{KeygenCircuitIntent, ProvingKeyGenerator},
8    },
9    halo2_proofs::{
10        plonk::{Circuit, ProvingKey},
11        poly::{commitment::ParamsProver, kzg::commitment::ParamsKZG},
12    },
13    halo2curves::bn256::{Bn256, Fr, G1Affine},
14    snark_verifier_sdk::{
15        halo2::{
16            aggregation::AggregationCircuit,
17            utils::{
18                AggregationDependencyIntent, AggregationDependencyIntentOwned,
19                KeygenAggregationCircuitIntent,
20            },
21        },
22        CircuitExt, Snark,
23    },
24    utils::{
25        build_utils::{
26            aggregation::get_dummy_aggregation_params,
27            keygen::{
28                compile_agg_dep_to_protocol, get_dummy_rlc_keccak_params, read_srs_from_dir,
29                write_pk_and_pinning,
30            },
31            pinning::aggregation::{AggTreeId, GenericAggParams, GenericAggPinning},
32        },
33        merkle_aggregation::keygen::AggIntentMerkle,
34        DEFAULT_RLC_CACHE_BITS,
35    },
36};
37use serde::{Deserialize, Serialize};
38
39use crate::{
40    aggregation::{
41        final_merkle::{
42            EthBlockHeaderChainRootAggregationCircuit, EthBlockHeaderChainRootAggregationInput,
43        },
44        intermediate::EthBlockHeaderChainIntermediateAggregationInput,
45    },
46    header_chain::{EthBlockHeaderChainCircuit, EthBlockHeaderChainInput},
47    types::{
48        CoreNodeParams, CoreNodeType, CorePinningIntermediate, CorePinningLeaf, CorePinningRoot,
49    },
50};
51
52/// Recursive intent for a node in the aggregation tree that can construct proving keys for this node and all its children.
53#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
54pub struct RecursiveCoreIntent {
55    /// `k_at_depth[i]` is the log2 domain size at depth `i`. So it starts from the current node and goes down the
56    /// layers of the tree. Our aggregation structure is such that each layer of the tree only has a single circuit type so only one `k` is needed.
57    pub k_at_depth: Vec<u32>,
58    /// Different chains (e.g., Goerli) can have different maximum number of bytes in the extra data field of the block header.
59    /// We configure the circuits differently based on this.
60    pub max_extra_data_bytes: usize,
61    pub params: CoreNodeParams,
62}
63
64impl RecursiveCoreIntent {
65    pub fn new(k_at_depth: Vec<u32>, max_extra_data_bytes: usize, params: CoreNodeParams) -> Self {
66        Self { k_at_depth, max_extra_data_bytes, params }
67    }
68    /// Each layer of tree has a unique circuit type, so this is the child circuit type.
69    pub fn child(&self) -> Option<Self> {
70        assert!(!self.k_at_depth.is_empty());
71        self.params.child(Some(self.max_extra_data_bytes)).map(|params| Self {
72            k_at_depth: self.k_at_depth[1..].to_vec(),
73            max_extra_data_bytes: self.max_extra_data_bytes,
74            params,
75        })
76    }
77}
78
79#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)]
80pub struct CoreIntentLeaf {
81    pub k: u32,
82    /// Different chains (e.g., Goerli) can have different maximum number of bytes in the extra data field of the block header.
83    /// We configure the circuits differently based on this.
84    pub max_extra_data_bytes: usize,
85    /// The leaf layer of the aggregation starts with max number of block headers equal to 2<sup>depth</sup>.
86    pub depth: usize,
87}
88
89#[derive(Clone, Debug)]
90pub(crate) struct CoreIntentIntermediate {
91    pub k: u32,
92    // This is from bad UX; only svk = kzg_params.get_g()[0] is used
93    pub kzg_params: Arc<ParamsKZG<Bn256>>,
94    // There will be exponential duplication since both children have the same circuit type, but it seems better for clarity
95    pub to_agg: Vec<AggTreeId>,
96    /// There are always two children of the same type, so we only specify the intent for one of them
97    pub child_intent: AggregationDependencyIntentOwned,
98    /// The maximum number of block headers in the chain at this level of the tree is 2<sup>depth</sup>.
99    pub depth: usize,
100    /// The leaf layer of the aggregation starts with max number of block headers equal to 2<sup>initial_depth</sup>.
101    pub initial_depth: usize,
102}
103
104#[derive(Clone, Debug)]
105pub(crate) struct CoreIntentRoot {
106    pub k: u32,
107    // This is from bad UX; only svk = kzg_params.get_g()[0] is used
108    pub(crate) kzg_params: Arc<ParamsKZG<Bn256>>,
109    // There will be exponential duplication since both children have the same circuit type, but it seems better for clarity
110    pub to_agg: Vec<AggTreeId>,
111    /// There are always two children of the same type, so we only specify the intent for one of them
112    pub child_intent: AggregationDependencyIntentOwned,
113    /// The maximum number of block headers in the chain at this level of the tree is 2<sup>depth</sup>.
114    pub depth: usize,
115    /// The leaf layer of the aggregation starts with max number of block headers equal to 2<sup>initial_depth</sup>.
116    pub initial_depth: usize,
117}
118
119/// Passthrough wrapper aggregation.
120/// Internal version doesn't need any additional context.
121#[derive(Clone, Debug)]
122pub(crate) struct CoreIntentEvm {
123    pub k: u32,
124    // This is from bad UX; only svk = kzg_params.get_g()[0] is used
125    pub kzg_params: Arc<ParamsKZG<Bn256>>,
126    /// Tree of the single child
127    pub to_agg: AggTreeId,
128    /// Wrap single child
129    pub child_intent: AggregationDependencyIntentOwned,
130}
131
132impl KeygenCircuitIntent<Fr> for CoreIntentLeaf {
133    type ConcreteCircuit = EthBlockHeaderChainCircuit<Fr>;
134    type Pinning = CorePinningLeaf;
135    fn get_k(&self) -> u32 {
136        self.k
137    }
138    fn build_keygen_circuit(self) -> Self::ConcreteCircuit {
139        let dummy_input = EthBlockHeaderChainInput::<Fr>::new(
140            vec![GENESIS_BLOCK_RLP.to_vec(); 1 << self.depth], // the resizing of each header RLP is handled by constructor
141            1,
142            self.depth,
143            self.max_extra_data_bytes,
144        );
145        let circuit_params = get_dummy_rlc_keccak_params(self.k as usize, 8);
146        let mut circuit = EthBlockHeaderChainCircuit::new_impl(
147            CircuitBuilderStage::Keygen,
148            dummy_input,
149            circuit_params,
150            DEFAULT_RLC_CACHE_BITS,
151        );
152        circuit.calculate_params();
153        circuit
154    }
155    fn get_pinning_after_keygen(
156        self,
157        kzg_params: &ParamsKZG<Bn256>,
158        circuit: &Self::ConcreteCircuit,
159    ) -> Self::Pinning {
160        let params = circuit.params();
161        let break_points = circuit.break_points();
162        let num_instance = circuit.num_instance();
163        let dk = (kzg_params.get_g()[0], kzg_params.g2(), kzg_params.s_g2());
164        CorePinningLeaf { params, break_points, num_instance, dk: dk.into() }
165    }
166}
167
168impl KeygenAggregationCircuitIntent for CoreIntentIntermediate {
169    fn intent_of_dependencies(&self) -> Vec<AggregationDependencyIntent> {
170        vec![(&self.child_intent).into(); 2]
171    }
172    fn build_keygen_circuit_from_snarks(self, snarks: Vec<Snark>) -> Self::AggregationCircuit {
173        assert_eq!(snarks.len(), 2);
174
175        let input = EthBlockHeaderChainIntermediateAggregationInput::new(
176            snarks,
177            1,
178            self.depth,
179            self.initial_depth,
180        );
181        let circuit_params = get_dummy_aggregation_params(self.k as usize);
182        let mut circuit =
183            input.build(CircuitBuilderStage::Keygen, circuit_params, &self.kzg_params).unwrap();
184        circuit.0.calculate_params(Some(20));
185        circuit.0
186    }
187}
188
189impl KeygenCircuitIntent<Fr> for CoreIntentIntermediate {
190    type ConcreteCircuit = AggregationCircuit;
191    type Pinning = CorePinningIntermediate;
192    fn get_k(&self) -> u32 {
193        self.k
194    }
195    fn build_keygen_circuit(self) -> Self::ConcreteCircuit {
196        self.build_keygen_circuit_shplonk()
197    }
198    fn get_pinning_after_keygen(
199        self,
200        kzg_params: &ParamsKZG<Bn256>,
201        circuit: &Self::ConcreteCircuit,
202    ) -> Self::Pinning {
203        let to_agg = compile_agg_dep_to_protocol(kzg_params, &self.child_intent, false);
204        let dk = (kzg_params.get_g()[0], kzg_params.g2(), kzg_params.s_g2());
205        CorePinningIntermediate {
206            params: circuit.params(),
207            to_agg: vec![to_agg; self.to_agg.len()],
208            break_points: circuit.break_points(),
209            num_instance: circuit.num_instance(),
210            dk: dk.into(),
211        }
212    }
213}
214
215impl KeygenAggregationCircuitIntent for CoreIntentRoot {
216    type AggregationCircuit = EthBlockHeaderChainRootAggregationCircuit;
217    fn intent_of_dependencies(&self) -> Vec<AggregationDependencyIntent> {
218        vec![(&self.child_intent).into(); 2]
219    }
220    fn build_keygen_circuit_from_snarks(self, snarks: Vec<Snark>) -> Self::AggregationCircuit {
221        assert_eq!(snarks.len(), 2);
222
223        let input = EthBlockHeaderChainRootAggregationInput::new(
224            snarks,
225            1,
226            self.depth,
227            self.initial_depth,
228            &self.kzg_params,
229        )
230        .unwrap();
231        // This is aggregation circuit, so set lookup bits to max
232        let circuit_params = get_dummy_rlc_keccak_params(self.k as usize, self.k as usize - 1);
233        // This is from bad UX; only svk = kzg_params.get_g()[0] is used
234        let mut circuit = EthBlockHeaderChainRootAggregationCircuit::new_impl(
235            CircuitBuilderStage::Keygen,
236            input,
237            circuit_params,
238            0, // note: rlc is not used
239        );
240        circuit.calculate_params();
241        circuit
242    }
243}
244
245impl KeygenCircuitIntent<Fr> for CoreIntentRoot {
246    type ConcreteCircuit = EthBlockHeaderChainRootAggregationCircuit;
247    type Pinning = CorePinningRoot;
248    fn get_k(&self) -> u32 {
249        self.k
250    }
251    fn build_keygen_circuit(self) -> Self::ConcreteCircuit {
252        self.build_keygen_circuit_shplonk()
253    }
254    fn get_pinning_after_keygen(
255        self,
256        kzg_params: &ParamsKZG<Bn256>,
257        circuit: &Self::ConcreteCircuit,
258    ) -> Self::Pinning {
259        let params = circuit.params();
260        let break_points = circuit.break_points();
261        let to_agg = compile_agg_dep_to_protocol(kzg_params, &self.child_intent, false);
262        let dk = (kzg_params.get_g()[0], kzg_params.g2(), kzg_params.s_g2());
263        CorePinningRoot {
264            params,
265            to_agg: vec![to_agg; self.to_agg.len()],
266            num_instance: circuit.num_instance(),
267            break_points,
268            dk: dk.into(),
269        }
270    }
271}
272
273impl From<CoreIntentEvm> for AggIntentMerkle {
274    fn from(value: CoreIntentEvm) -> Self {
275        AggIntentMerkle {
276            kzg_params: value.kzg_params,
277            to_agg: vec![value.to_agg],
278            deps: vec![value.child_intent],
279            k: value.k,
280        }
281    }
282}
283
284impl KeygenCircuitIntent<Fr> for CoreIntentEvm {
285    type ConcreteCircuit = AggregationCircuit;
286    type Pinning = GenericAggPinning<GenericAggParams>;
287    fn get_k(&self) -> u32 {
288        self.k
289    }
290    fn build_keygen_circuit(self) -> Self::ConcreteCircuit {
291        AggIntentMerkle::from(self).build_keygen_circuit()
292    }
293    fn get_pinning_after_keygen(
294        self,
295        kzg_params: &ParamsKZG<Bn256>,
296        circuit: &Self::ConcreteCircuit,
297    ) -> Self::Pinning {
298        AggIntentMerkle::from(self).get_pinning_after_keygen(kzg_params, circuit)
299    }
300}
301
302impl RecursiveCoreIntent {
303    /// Recursively creates and serializes proving keys and pinnings.
304    ///
305    /// Computes `circuit_id` as the blake3 hash of the halo2 VerifyingKey written to bytes. Writes proving key to `circuit_id.pk`, verifying key to `circuit_id.vk` and pinning to `circuit_id.json` in the `data_dir` directory.
306    ///
307    /// Returns the `circuit_id, proving_key, pinning`.
308    /// * `cid_repo` stores a mapping from the [CoreNodeParams] to the corresponding circuit ID. In an Axiom Core aggregation tree, the [CoreNodeParams] determines the node.
309    /// * `PARAMS_DIR` **must** be set because the aggregation circuit creation requires reading trusted setup files (this can be removed later).
310    pub fn create_and_serialize_proving_key(
311        self,
312        srs_dir: &Path,
313        data_dir: &Path,
314        cid_repo: &mut BTreeMap<CoreNodeParams, String>,
315    ) -> anyhow::Result<(AggTreeId, ProvingKey<G1Affine>, serde_json::Value)> {
316        // If there is child, do it first
317        let child = if let Some(child_intent) = self.child() {
318            let is_aggregation = !matches!(child_intent.params.node_type, CoreNodeType::Leaf(_));
319            let (child_id, child_pk, child_pinning) =
320                child_intent.create_and_serialize_proving_key(srs_dir, data_dir, cid_repo)?;
321            let num_instance: Vec<usize> =
322                serde_json::from_value(child_pinning["num_instance"].clone())?;
323            // !! ** ASSERTION: all aggregation circuits in Axiom Core have accumulator in indices 0..12 ** !!
324            // No aggregation circuits in Axiom Core are universal.
325            let agg_intent = AggregationDependencyIntentOwned {
326                vk: child_pk.get_vk().clone(),
327                num_instance,
328                accumulator_indices: is_aggregation
329                    .then(|| AggregationCircuit::accumulator_indices().unwrap()),
330                agg_vk_hash_data: None,
331            };
332            Some((child_id, agg_intent))
333        } else {
334            None
335        };
336        assert!(!self.k_at_depth.is_empty());
337        let k = self.k_at_depth[0];
338        let kzg_params = Arc::new(read_srs_from_dir(srs_dir, k)?);
339        let ((pk, pinning), children) = match self.params.node_type {
340            CoreNodeType::Leaf(max_extra_data_bytes) => {
341                let intent =
342                    CoreIntentLeaf { k, max_extra_data_bytes, depth: self.params.initial_depth };
343                (intent.create_pk_and_pinning(&kzg_params), vec![])
344            }
345            CoreNodeType::Intermediate => {
346                let (child_id, child_intent) = child.unwrap();
347                let to_agg = vec![child_id; 2];
348                let intent = CoreIntentIntermediate {
349                    k,
350                    kzg_params: kzg_params.clone(),
351                    to_agg: to_agg.clone(),
352                    child_intent,
353                    depth: self.params.depth,
354                    initial_depth: self.params.initial_depth,
355                };
356                (intent.create_pk_and_pinning(&kzg_params), to_agg)
357            }
358            CoreNodeType::Root => {
359                let (child_id, child_intent) = child.unwrap();
360                let to_agg = vec![child_id; 2];
361                let intent = CoreIntentRoot {
362                    k,
363                    kzg_params: kzg_params.clone(),
364                    to_agg: to_agg.clone(),
365                    child_intent,
366                    depth: self.params.depth,
367                    initial_depth: self.params.initial_depth,
368                };
369                (intent.create_pk_and_pinning(&kzg_params), to_agg)
370            }
371            CoreNodeType::Evm(_) => {
372                let (child_id, child_intent) = child.unwrap();
373                let to_agg = vec![child_id.clone()];
374                let intent = CoreIntentEvm {
375                    k,
376                    to_agg: child_id,
377                    child_intent,
378                    kzg_params: kzg_params.clone(),
379                };
380                (intent.create_pk_and_pinning(&kzg_params), to_agg)
381            }
382        };
383        let circuit_id = write_pk_and_pinning(data_dir, &pk, &pinning)?;
384        if let Some(old_cid) = cid_repo.insert(self.params, circuit_id.clone()) {
385            if old_cid != circuit_id {
386                anyhow::bail!("Different circuit ID for the same node params")
387            }
388        }
389        let tree_id = AggTreeId { circuit_id, children, aggregate_vk_hash: None };
390        Ok((tree_id, pk, pinning))
391    }
392}
axiom_core/keygen/mod.rs

axiom_core/keygen/
mod.rs
