refactor(mithril-stm): refactor and stabilize the SNARK setups#3349
refactor(mithril-stm): refactor and stabilize the SNARK setups#3349hjeljeli32 wants to merge 14 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR refactors and stabilizes the STM SNARK setup code paths (non-recursive certificate circuit + recursive IVC circuit) by removing test-only unsafe setup plumbing, adding typed on-disk key caching with unified key (de)serialization, and aligning setup/type naming between circuits.
Changes:
- Introduces a
CircuitKeySerializationtrait and circuit-specific implementations to uniformly serialize/deserialize proving/verifying keys. - Refactors the on-disk
CircuitKeyCacheinto a typed cache with staleness detection and typed VK/PK accessors. - Reworks recursive/non-recursive setup plumbing: removes temporary unsafe providers, adds cache-backed key providers, aligns naming (
SnarkSetup→SnarkProverSetup), and standardizes the recursive circuit degree constant (RECURSIVE_CIRCUIT_DEGREE).
Reviewed changes
Copilot reviewed 33 out of 33 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| mithril-stm/src/protocol/aggregate_signature/signature.rs | Updates tests to use SnarkProverSetup naming. |
| mithril-stm/src/proof_system/mod.rs | Updates re-exports to align with the new SNARK setup types. |
| mithril-stm/src/proof_system/ivc_halo2_snark/verifier_setup.rs | Renames/aligns recursive verifier key types (PlonkVerifyingKey). |
| mithril-stm/src/proof_system/ivc_halo2_snark/unsafe_setup_helpers.rs | Removes temporary unsafe setup helper module. |
| mithril-stm/src/proof_system/ivc_halo2_snark/prover_setup.rs | Refactors IvcProverSetup::load to orchestrate key providers around a single downsized SRS; updates tests to use caches. |
| mithril-stm/src/proof_system/ivc_halo2_snark/prover_input.rs | Updates slow tests to use cache-backed key providers and RECURSIVE_CIRCUIT_DEGREE. |
| mithril-stm/src/proof_system/ivc_halo2_snark/proof.rs | Updates IVC proving/verification code and tests to use new key/provider types. |
| mithril-stm/src/proof_system/ivc_halo2_snark/mod.rs | Removes unsafe helpers module and re-exports recursive PLONK key aliases from circuit layer. |
| mithril-stm/src/proof_system/ivc_halo2_snark/key_providers.rs | Adds cache-backed key providers for certificate and recursive circuits. |
| mithril-stm/src/proof_system/halo2_snark/unsafe_helpers.rs | Removes old SNARK setup implementation file. |
| mithril-stm/src/proof_system/halo2_snark/setup.rs | Adds the new SnarkProverSetup / SnarkVerifierSetup implementation using the typed key cache. |
| mithril-stm/src/proof_system/halo2_snark/proof.rs | Updates proof/prover code to use SnarkProverSetup and isolates test caches. |
| mithril-stm/src/proof_system/halo2_snark/mod.rs | Wires in the new setup module and updates exports. |
| mithril-stm/src/proof_system/halo2_snark/circuit_verification_key.rs | Switches VK byte conversion to the shared CircuitKeySerialization trait. |
| mithril-stm/src/circuits/mod.rs | Adds the new key_serialization module. |
| mithril-stm/src/circuits/key_serialization.rs | Introduces the CircuitKeySerialization trait. |
| mithril-stm/src/circuits/key_cache.rs | Refactors cache state, adds typed VK/PK accessors, and adds store_key_pair. |
| mithril-stm/src/circuits/halo2/mod.rs | Adds the certificate-circuit key serialization module. |
| mithril-stm/src/circuits/halo2/key_serialization.rs | Implements key (de)serialization for MidnightVK / MidnightPK. |
| mithril-stm/src/circuits/halo2_ivc/tests/verification_key_computation.rs | Updates to RECURSIVE_CIRCUIT_DEGREE. |
| mithril-stm/src/circuits/halo2_ivc/tests/off_circuit/mod.rs | Updates documentation references to RECURSIVE_CIRCUIT_DEGREE. |
| mithril-stm/src/circuits/halo2_ivc/tests/off_circuit/circuit_validation.rs | Updates to RECURSIVE_CIRCUIT_DEGREE in assertions/errors. |
| mithril-stm/src/circuits/halo2_ivc/tests/golden/mod.rs | Removes redundant local RECURSIVE_CIRCUIT_DEGREE constant. |
| mithril-stm/src/circuits/halo2_ivc/tests/common/mod.rs | Removes redundant local RECURSIVE_CIRCUIT_DEGREE constant. |
| mithril-stm/src/circuits/halo2_ivc/tests/common/helpers.rs | Updates MockProver invocations to use RECURSIVE_CIRCUIT_DEGREE. |
| mithril-stm/src/circuits/halo2_ivc/tests/common/generators/verification_key.rs | Updates imports to use RECURSIVE_CIRCUIT_DEGREE from module root. |
| mithril-stm/src/circuits/halo2_ivc/tests/common/generators/setup.rs | Updates to use module-level RECURSIVE_CIRCUIT_DEGREE. |
| mithril-stm/src/circuits/halo2_ivc/mod.rs | Introduces shared raw PLONK key aliases and renames K to RECURSIVE_CIRCUIT_DEGREE. |
| mithril-stm/src/circuits/halo2_ivc/key_serialization.rs | Implements key (de)serialization for raw PLONK verifying/proving keys. |
| mithril-stm/src/circuits/halo2_ivc/gadget.rs | Updates gadget configuration to use RECURSIVE_CIRCUIT_DEGREE. |
| mithril-stm/src/circuits/halo2_ivc/errors.rs | Updates docs to refer to RECURSIVE_CIRCUIT_DEGREE. |
| mithril-stm/src/circuits/halo2_ivc/circuit.rs | Updates validation and domain sizing to use RECURSIVE_CIRCUIT_DEGREE. |
| mithril-stm/src/circuits/halo2_ivc/certificate_proof.rs | Switches verifier key type to PlonkVerifyingKey for accumulator preparation. |
Test Results 5 files ±0 206 suites ±0 3h 10m 54s ⏱️ + 2m 43s Results for commit d5e79b1. ± Comparison against base commit 5d25797. This pull request removes 19 and adds 25 tests. Note that renamed tests count towards both.♻️ This comment has been updated with latest results. |
9b62741 to
23d3d73
Compare
damrobi
left a comment
There was a problem hiding this comment.
LGTM overall 👍
I think @jpraynaud comments on the PR #3348 about downsizing the srs will apply to the key provider of the IVC
A more general remark, are we planning to store/download the srs of degree 19 for the IVC circuit instead of downsizing? The downsizing operation appears to be pretty slow sometimes.
curiecrypt
left a comment
There was a problem hiding this comment.
Thanks @hjeljeli32, this PR will resolve a lot of issues we have been struggling.
I have a suggestion that would also avoid some ambiguotions we might have and @jpraynaud 's comment about using MidnightVK.
The certificate and recursive circuits share the same raw key type (VerifyingKey<F, KZGCommitmentScheme<E>>), and MidnightVK::vk() hands back that same type. Since PlonkVerifyingKey is only an alias for it, the CircuitKeySerialization impl is effectively defined on the shared raw type, and deserialize_key always rebuilds with IvcCircuitData. It works today because only recursive keys are ever deserialized raw, but nothing stops a certificate key (same type) from being decoded as the IVC circuit, and that would only surface at proving or verification time.
Proposal: give each circuit its own clearly named key types, so the names answer the "why not MidnightVK" question and the serialization impls become unambiguous.
// circuits/halo2 (non-recursive)
pub struct NonRecursiveCircuitVerifyingKey(MidnightVK);
pub struct NonRecursiveCircuitProvingKey(MidnightPK<StmCertificateCircuit>);
// circuits/halo2_ivc (recursive)
pub struct RecursiveCircuitVerifyingKey(VerifyingKey<F, KZGCommitmentScheme<E>>);
pub struct RecursiveCircuitProvingKey(ProvingKey<F, KZGCommitmentScheme<E>>);A few notes so this lands cleanly:
- Make the recursive pair real newtypes, not type aliases. A renamed alias is still the shared raw type, so it reads better but does not stop a certificate key from matching the recursive impl. If you want to keep them as aliases, add a degree check (
== RECURSIVE_CIRCUIT_DEGREE) inside the recursivedeserialize_keyso a wrong-circuit decode fails loudly instead of silently. circuit_verification_key.rsis not dead code.CircuitVerificationKeyis the serde-embedded VK field inSnarkProofand exposesget_midnight_vk()andfrom_bytes(), both of which are used. So this is a relocate and rename into the non-recursive keys file, carryingnew,from_bytes,get_midnight_vk, the serde derive, and themidnight_vk_serdemodule. It is not a delete.- Once the impl lives on
NonRecursiveCircuitVerifyingKey,midnight_vk_serdeandCircuitKeySerialization for MidnightVKencode the same "MidnightVK to RawBytes" logic. Point the#[serde(with = ...)]atserialize_keyanddeserialize_keyso there is a single source of truth. CircuitVerificationKey::to_bytesis only used by a#[cfg(test)]golden test, so gate it with#[cfg(test)]rather than#[allow(dead_code)].MidnightPKis generic, so the non-recursive proving key wrapsMidnightPK<StmCertificateCircuit>.
This also sets up the later move to generic TryToBytes / TryFromBytes traits in codec.rs: distinct newtypes avoid conflicting impls on the shared foreign key type.
We can have a discussion about my proposal and decide on the best choice.
4633f9b to
8382b4b
Compare
089bdcc to
5a0d3a8
Compare
@jpraynaud @damrobi IMO, the newtype direction is the right fix: it makes the cross-circuit decode impossible by construction, gives each key its own codec. A couple of your points are already in this PR (the generic The one challenge is scope: the aliases are used in ~37 places across 9 files, and it touches the same IVC area @damrobi is changing in #3348, so folding it into this PR would risk conflicts. My suggestion: do the newtypes as a follow-up ticket, coordinated with @damrobi once #3348 is merged. WDYT? |
5a0d3a8 to
8f6b387
Compare
ff7e18a to
fbe74b2
Compare
…erifyingKey/PlonkProvingKey
…ircuits, delegating CircuitVerificationKey
355a3f2 to
4ee3c5c
Compare
damrobi
left a comment
There was a problem hiding this comment.
Left a few non blocking comments I think. I haven't had time to look at the tests in details so I will do so on Monday and approve
| return Err(error).with_context(|| "Failed to store the proving key in the cache"); | ||
| } | ||
| if let Err(error) = fs::rename(&verification_key_temp, &self.verification_key_path) { | ||
| let _ = fs::remove_file(&verification_key_temp); |
There was a problem hiding this comment.
We don't want to remove the proving key file in that case ?
| merkle_tree_depth: u32, | ||
| /// Circuit verification key | ||
| circuit_verification_key: CircuitVerificationKey, | ||
| circuit_verification_key: NonRecursiveCircuitVerifyingKey, |
There was a problem hiding this comment.
Should we continue to store the verification key with the proof or try to get it from the cache when trying to verify the proof?
| pub(crate) verification_key: NonRecursiveCircuitVerifyingKey, | ||
| /// Proving key for the SNARK proof. |
There was a problem hiding this comment.
This might not be needed anymore if we remove the verification key from the proof but if we leave it this can stay.
| ) -> StmResult<Self> { | ||
| let circuit = provider.circuit().clone(); | ||
| zk::downsize_srs_for_relation(&mut srs, &circuit); | ||
| let (verification_key, proving_key) = provider.key_pair(&srs)?; |
There was a problem hiding this comment.
This key_pair() function can clone the srs to downsize it again if the keys are not available. Maybe it's worth checking if we can remove the second downsizing as it demands cloning the full SRS.
| tau_g2: G2Affine, | ||
| /// Verifying key of the IVC circuit. | ||
| ivc_verifying_key: CircuitVerifyingKey, | ||
| ivc_verifying_key: PlonkVerifyingKey, |
There was a problem hiding this comment.
I think this should be the new type RecursiveCircuitVerifyingKey.
| ) -> StmResult<(MidnightVK, MidnightPK<StmCertificateCircuit>)> { | ||
| // The proving key is only read once the verifying key is present, so an orphan proving key left | ||
| // by a failed store (which removes the verifying key) is recomputed rather than surfaced as a | ||
| // deserialization error. | ||
| if let Some(verification_key) = disk_cache.get_verification_key::<MidnightVK>()? | ||
| && let Some(proving_key) = | ||
| disk_cache.get_proving_key::<MidnightPK<StmCertificateCircuit>>()? | ||
| { | ||
| return Ok(pair); | ||
| return Ok((verification_key, proving_key)); | ||
| } |
There was a problem hiding this comment.
This should be a guarantee of the CircuitKeyCache (and I believe it is once the validate function has been called).
| } | ||
|
|
||
| /// Borrows the wrapped Midnight verifying key. | ||
| pub(crate) fn midnight_vk(&self) -> &MidnightVK { |
There was a problem hiding this comment.
This could be implemented with Deref trait.
|
|
||
| impl NonRecursiveCircuitProvingKey { | ||
| /// Borrows the wrapped Midnight proving key, for proof generation. | ||
| pub(crate) fn midnight_pk(&self) -> &MidnightPK<StmCertificateCircuit> { |
There was a problem hiding this comment.
This also could be implemented with Deref trait.
| } | ||
| } | ||
|
|
||
| impl CircuitVerificationKeyProvider<StmCertificateCircuit> { |
There was a problem hiding this comment.
Maybe this should live in the CircuitVerificationKeyProvider module?
| } | ||
|
|
||
| #[test] | ||
| fn provider_returns_cached_verifying_key_on_hit() { |
There was a problem hiding this comment.
This test looks like you are trying to test underlying implementation details and can probably be removed.
| // once they ship. | ||
| /// `recursive_provider_factory` builds the recursive provider once the certificate verifying key | ||
| /// is known; production passes [`CircuitVerificationKeyProvider::for_recursive_circuit`]. | ||
| pub(crate) fn load( |
There was a problem hiding this comment.
We should try to align as much as possible the function names between recursive and non-recursive.
| recursive_provider_factory: impl FnOnce( | ||
| &NonRecursiveCircuitVerifyingKey, | ||
| ) -> StmResult< | ||
| CircuitVerificationKeyProvider<IvcCircuitData>, | ||
| >, |
There was a problem hiding this comment.
Why not inject the non-recursive CircuitKeyProvider in the recursive one and provide on the recursive one to this function instead? This would avoid this complex constructor.
| use crate::circuits::{ | ||
| halo2::NON_RECURSIVE_CIRCUIT_VERIFICATION_KEY_FOR_PRODUCTION, | ||
| halo2_ivc::RECURSIVE_CIRCUIT_VERIFICATION_KEY_FOR_PRODUCTION, | ||
| test_utils::key_cache::with_shared_key_cache, | ||
| trusted_setup::{UNSAFE_SRS_SEED, build_provider_with_unsafe_srs}, | ||
| }; |
There was a problem hiding this comment.
Can you avoid imports in function body?
| /// bare-key shape (a `MidnightVK` and a `PlonkVerifyingKey`, each via a raw-bytes `#[serde(with)]` | ||
| /// module), serializes it through the same codec, then asserts the current struct decodes those | ||
| /// bytes and re-encodes them byte-for-byte. | ||
| mod wire_format_compatibility { |
There was a problem hiding this comment.
I don't really understand what this tests do.
There was a problem hiding this comment.
This module should stay as temporary or better be removed as it is only here to avoid re-computing SRS and circuit keys constantly.
Content
This PR refactors and stabilizes both SNARK setups (the non-recursive certificate circuit and the recursive IVC circuit). It replaces the test-only "unsafe" setup plumbing and the duplicated key-computation paths with a single compute-on-miss key provider backed by a typed on-disk cache, introduces per-circuit key newtypes with their own (de)serialization, aligns the prover/verifier setup types and names across both circuits, and makes the test setup deterministic, race-free, and key-reusing.
Changes
CircuitVerificationKeyProvider<C>owns the on-disk key cache and computes the key pair on a miss through the circuit'sCircuitKeyGenerator; the standalone key cache was folded into the provider. The cached verifying key is checked for staleness against the embedded production verifying-key bytes.NonRecursiveCircuit{Verifying,Proving}KeyandRecursiveCircuit{Verifying,Proving}Keywrap the Midnight and raw PLONK key types, each with its own key (de)serialization. The rawPlonkVerifyingKey/PlonkProvingKeyaliases remain only at the Halo2/verifier boundary.SnarkSetup→SnarkProverSetup.IvcProverSetup→IvcSnarkProverSetup; it loads a single SRS, derives the certificate verifying key through the provider, downsizes the SRS to the recursive circuit degree, builds the recursive provider from that certificate key, and derives the IVC key pair — aligning the setup types and names across both circuits.temp_ivc_prover_cache→ivc_prover_setup_cache(it memoizes the assembled setup on the clerk path, not key computation).RECURSIVE_CIRCUIT_DEGREE.Tests
SnarkProof::try_new_with_unsafe_srsis documented as being only for the deterministic unsafe SRS.Pre-submit checklist
Issue(s)
Closes #3300