feat(proof-compression): implement proof systems

Author: saitima

crates/proof-compression/src/artifacts.rs

@@ -3,31 +3,44 @@ use std::io::Read;
 use boojum::cs::implementations::fast_serialization::MemcopySerializable;
 use boojum::cs::implementations::setup::FinalizationHintsForProver;
 use boojum::cs::traits::gate::FinalizationHintSerialized;
-use circuit_definitions::circuit_definitions::aux_layer::compression_modes::CompressionTreeHasherForWrapper;
-use circuit_definitions::circuit_definitions::aux_layer::{
-    CompressionProofsTreeHasher, ZkSyncCompressionVerificationKey,
-};
+use circuit_definitions::circuit_definitions::aux_layer::ZkSyncCompressionVerificationKey;
 use circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionVerificationKey;
-use shivini::cs::GpuSetup as BoojumDeviceSetup;
 
 use super::*;
 
 pub trait BlobStorage {
-    fn save<T>(artifact: T)
+    fn save<W>(data: W)
     where
-        T: MemcopySerializable;
+        W: std::io::Write;
 }
 
 pub struct AsyncHandler<T> {
-    precomputation: std::sync::mpsc::Receiver<T>, // This is indeed a receiver
+    receiver: std::sync::mpsc::Receiver<T>,
 }
 
-impl<T> AsyncHandler<T> {
-    pub fn into_inner(self) -> T {
-        todo!()
+impl<T> AsyncHandler<T>
+where
+    T: Send + Sync + 'static,
+{
+    pub fn spawn<F>(f: F) -> Self
+    where
+        F: FnOnce() -> std::sync::mpsc::Receiver<T> + Send + Sync + 'static,
+    {
+        let receiver = std::thread::spawn(f);
+
+        Self {
+            receiver: receiver.join().unwrap(),
+        }
+    }
+
+    pub fn wait(self) -> T {
+        self.receiver.recv().unwrap()
     }
 }
-pub trait ArtifactLoader: Sized {
+
+use shivini::cs::GpuSetup;
+
+pub trait ArtifactLoader: Sized + Send + Sync {
     fn init<BS>(bs: BS) -> Self
     where
         BS: BlobStorage;
@@ -54,18 +67,16 @@ pub trait ArtifactLoader: Sized {
     fn get_compression_layer_precomputation(
         &self,
         circuit_id: u8,
-    ) -> AsyncHandler<TreeHasherCompatibleGpuSetup<CompressionProofsTreeHasher>>;
+    ) -> Box<dyn Read + Send + Sync + 'static>;
     fn get_compression_wrapper_precomputation(
         &self,
         circuit_id: u8,
-    ) -> AsyncHandler<TreeHasherCompatibleGpuSetup<CompressionTreeHasherForWrapper>>;
+    ) -> Box<dyn Read + Send + Sync + 'static>;
     fn get_plonk_snark_wrapper_precomputation(
         &self,
         circuit_id: u8,
-    ) -> AsyncHandler<PlonkSnarkVerifierCircuitDeviceSetupWrapper>;
-    fn get_fflonk_snark_wrapper_precomputation(
-        &self,
-    ) -> AsyncHandler<FflonkSnarkVerifierCircuitDeviceSetupWrapper>;
+    ) -> Box<dyn Read + Send + Sync + 'static>;
+    fn get_fflonk_snark_wrapper_precomputation(&self) -> Box<dyn Read + Send + Sync + 'static>;
 }
 
 pub struct SimpleArtifactLoader;
@@ -139,27 +150,25 @@ impl ArtifactLoader for SimpleArtifactLoader {
     fn get_compression_layer_precomputation(
         &self,
         circuit_id: u8,
-    ) -> AsyncHandler<TreeHasherCompatibleGpuSetup<CompressionProofsTreeHasher>> {
+    ) -> Box<dyn Read + Send + Sync + 'static> {
         todo!()
     }
 
     fn get_compression_wrapper_precomputation(
         &self,
         circuit_id: u8,
-    ) -> AsyncHandler<TreeHasherCompatibleGpuSetup<CompressionTreeHasherForWrapper>> {
+    ) -> Box<dyn Read + Send + Sync + 'static> {
         todo!()
     }
 
     fn get_plonk_snark_wrapper_precomputation(
         &self,
         circuit_id: u8,
-    ) -> AsyncHandler<PlonkSnarkVerifierCircuitDeviceSetupWrapper> {
+    ) -> Box<dyn Read + Send + Sync + 'static> {
         todo!()
     }
 
-    fn get_fflonk_snark_wrapper_precomputation(
-        &self,
-    ) -> AsyncHandler<FflonkSnarkVerifierCircuitDeviceSetupWrapper> {
+    fn get_fflonk_snark_wrapper_precomputation(&self) -> Box<dyn Read + Send + Sync + 'static> {
         todo!()
     }
 }

crates/proof-compression/src/chain.rs

@@ -1,200 +1,88 @@
-use circuit_definitions::circuit_definitions::aux_layer::{
-    compression::ProofCompressionFunction,
-    compression_modes::{
-        CompressionMode1, CompressionMode1ForWrapper, CompressionMode2, CompressionMode3,
-        CompressionMode4, CompressionMode5ForWrapper,
-    },
-};
-
 use super::*;
 
-pub struct CompressionChain<C: CompressionStep, U>(std::marker::PhantomData<(C, U)>);
-
-impl<C, U> CompressionChain<C, U>
-where
-    C: CompressionStep,
-{
-    pub fn new(step: C) -> CompressionChain<C, U>
-    where
-        U: CompressionStep,
-    {
-        todo!()
-    }
-
-    pub fn next(self) -> U {
-        todo!()
-    }
+pub enum SnarkWrapper {
+    Plonk,
+    FFfonk,
 }
-
-impl<C, U> StepDefinition for CompressionChain<C, U>
-where
-    C: CompressionStep,
-{
-    type PreviousProofSystem = C::PreviousProofSystem;
-    type ThisProofSystem = C::ThisProofSystem;
+pub enum SnarkWrapperProof {
+    Plonk(PlonkSnarkVerifierCircuitProof),
+    FFfonk(FflonkSnarkVerifierCircuitProof),
 }
-
-impl<C, U> ProofCompressionFunction for CompressionChain<C, U>
-where
-    C: CompressionStep,
-{
-    type PreviousLayerPoW = C::PreviousLayerPoW;
-
-    type ThisLayerPoW = C::ThisLayerPoW;
-
-    type ThisLayerHasher = C::ThisLayerHasher;
-
-    type ThisLayerTranscript = C::ThisLayerTranscript;
-
-    fn this_layer_transcript_parameters(
-    ) -> <Self::ThisLayerTranscript as boojum::cs::implementations::transcript::Transcript<
-        GoldilocksField,
-    >>::TransciptParameters {
-        C::this_layer_transcript_parameters()
-    }
-
-    fn description_for_compression_step() -> String {
-        C::description_for_compression_step()
-    }
-
-    fn size_hint_for_compression_step() -> (usize, usize) {
-        C::size_hint_for_compression_step()
-    }
-
-    fn geometry_for_compression_step() -> boojum::cs::CSGeometry {
-        C::geometry_for_compression_step()
-    }
-
-    fn lookup_parameters_for_compression_step() -> boojum::cs::LookupParameters {
-        C::lookup_parameters_for_compression_step()
-    }
-
-    fn configure_builder_for_compression_step<
-        T: boojum::cs::cs_builder::CsBuilderImpl<GoldilocksField, T>,
-        GC: boojum::cs::GateConfigurationHolder<GoldilocksField>,
-        TB: boojum::cs::StaticToolboxHolder,
-    >(
-        builder: boojum::cs::cs_builder::CsBuilder<T, GoldilocksField, GC, TB>,
-    ) -> boojum::cs::cs_builder::CsBuilder<
-        T,
-        GoldilocksField,
-        impl boojum::cs::GateConfigurationHolder<GoldilocksField>,
-        impl boojum::cs::StaticToolboxHolder,
-    > {
-        C::configure_builder_for_compression_step(builder)
-    }
-
-    fn previous_step_builder_for_compression<
-        CS: boojum::cs::traits::cs::ConstraintSystem<GoldilocksField> + 'static,
-    >() -> Box<
-        dyn boojum::cs::traits::circuit::ErasedBuilderForRecursiveVerifier<
-            GoldilocksField,
-            GoldilocksExt2,
-            CS,
-        >,
-    > {
-        C::previous_step_builder_for_compression()
-    }
-
-    fn proof_config_for_compression_step() -> boojum::cs::implementations::prover::ProofConfig {
-        C::proof_config_for_compression_step()
-    }
-}
-
-impl<C, T> CompressionStep for CompressionChain<C, T>
-where
-    C: CompressionStep,
-{
-    const MODE: u8 = C::MODE;
-    const IS_WRAPPER: bool = C::IS_WRAPPER;
-}
-
-pub struct CompressionChainBuilder<C: CompressionStep>(std::marker::PhantomData<C>);
-
-impl<C> CompressionChainBuilder<C>
-where
-    C: CompressionStep,
-{
-    pub fn new(step: C) -> CompressionChainBuilder<C> {
-        todo!()
-    }
-}
-
-impl<C> CompressionChainBuilder<C>
-where
-    C: CompressionStep,
-{
-    pub fn link_compression_step<CC>(
-        self,
-        step: CC,
-    ) -> CompressionChainBuilder<CompressionChain<C, CC>>
-    where
-        CC: CompressionStep,
-    {
-        todo!()
-    }
-
-    pub fn build(self) -> C {
-        todo!()
-    }
-}
-
-pub fn build_full_chain_with_fflonk() -> CompressionChain<
-    CompressionChain<
-        CompressionChain<CompressionChain<CompressionMode1, CompressionMode2>, CompressionMode3>,
-        CompressionMode4,
-    >,
-    CompressionMode5ForWrapper,
-> {
-    let chain = CompressionChainBuilder::new(CompressionMode1)
-        .link_compression_step(CompressionMode2)
-        .link_compression_step(CompressionMode3)
-        .link_compression_step(CompressionMode4)
-        .link_compression_step(CompressionMode5ForWrapper)
-        .build();
-
-    chain
-}
-
-pub fn build_full_chain_with_plonk() -> CompressionChain<CompressionMode1ForWrapper, ()> {
-    todo!()
-}
-
-pub fn run_step_chain_with_fflonk<BS>(
+pub fn wrap_proof<BS>(
     input_proof: SchedulerProof,
+    snark_wrapper: SnarkWrapper,
     blob_storage: BS,
-) -> FflonkSnarkVerifierCircuitProof
+) -> SnarkWrapperProof
 where
     BS: BlobStorage,
 {
-    let chain = build_full_chain_with_fflonk();
-    run_step_chain::<_, _, FflonkSnarkWrapper, _>(chain, input_proof, blob_storage)
+    match snark_wrapper {
+        SnarkWrapper::Plonk => {
+            let proof = run_step_chain_with_plonk(input_proof, blob_storage);
+            SnarkWrapperProof::Plonk(proof)
+        }
+        SnarkWrapper::FFfonk => {
+            let proof = run_step_chain_with_fflonk(input_proof, blob_storage);
+            SnarkWrapperProof::FFfonk(proof)
+        }
+    }
 }
 
-pub fn run_step_chain_with_plonk<BS>(
+pub(crate) fn run_step_chain_with_fflonk<BS>(
     input_proof: SchedulerProof,
     blob_storage: BS,
-) -> PlonkSnarkVerifierCircuitProof
+) -> FflonkSnarkVerifierCircuitProof
 where
     BS: BlobStorage,
 {
-    let chain = build_full_chain_with_plonk();
-    todo!()
+    let artifact_loader = SimpleArtifactLoader::init(blob_storage);
+    let context_initializor = SimpelContextInitializor::new();
+
+    let next_proof = CompressionMode1::prove_compression_step::<_, SimpelContextInitializor>(
+        input_proof,
+        &artifact_loader,
+    );
+    let next_proof = CompressionMode2::prove_compression_step::<_, SimpelContextInitializor>(
+        next_proof,
+        &artifact_loader,
+    );
+    let next_proof = CompressionMode3::prove_compression_step::<_, SimpelContextInitializor>(
+        next_proof,
+        &artifact_loader,
+    );
+    let next_proof = CompressionMode4::prove_compression_step::<_, SimpelContextInitializor>(
+        next_proof,
+        &artifact_loader,
+    );
+    let next_proof = CompressionMode5ForWrapper::prove_compression_step::<
+        _,
+        SimpelContextInitializor,
+    >(next_proof, &artifact_loader);
+    let final_proof = FflonkSnarkWrapper::prove_snark_wrapper_step::<_, SimpelContextInitializor>(
+        next_proof,
+        &artifact_loader,
+    );
+
+    final_proof
 }
 
-pub fn run_step_chain<C1, C2, SW, BS>(
-    chain: CompressionChain<C1, C2>,
+pub(crate) fn run_step_chain_with_plonk<BS>(
     input_proof: SchedulerProof,
     blob_storage: BS,
-) -> <SW::ThisProofSystem as ProofSystemDefinition>::Proof
+) -> PlonkSnarkVerifierCircuitProof
 where
-    C1: CompressionStep,
-    C2: CompressionStep,
-    SW: SnarkWrapperStep,
     BS: BlobStorage,
 {
     let artifact_loader = SimpleArtifactLoader::init(blob_storage);
+    let next_proof = CompressionMode1ForWrapper::prove_compression_step::<
+        _,
+        SimpelContextInitializor,
+    >(input_proof, &artifact_loader);
+
+    let final_proof = PlonkSnarkWrapper::prove_snark_wrapper_step::<_, SimpelContextInitializor>(
+        next_proof,
+        &artifact_loader,
+    );
 
-    let next_step = chain.next();
-    todo!()
+    final_proof
 }

crates/proof-compression/src/common.rs

@@ -1,7 +1,6 @@
+use super::*;
 use circuit_definitions::circuit_definitions::{
     aux_layer::{
-        compression::{CompressionLayerCircuit, ProofCompressionFunction},
-        compression_modes::CompressionTranscriptForWrapper,
         CompressionProofsTreeHasher, CompressionProofsTreeHasherForWrapper,
         ZkSyncCompressionForWrapperCircuit, ZkSyncCompressionLayerCircuit, ZkSyncCompressionProof,
         ZkSyncCompressionProofForWrapper, ZkSyncCompressionVerificationKey,