From 872e85c4f6f65255673fec0aa77f8881adc5016b Mon Sep 17 00:00:00 2001
From: Daniyar Itegulov <ditegulov@gmail.com>
Date: Thu, 10 Oct 2024 16:34:48 +1100
Subject: [PATCH] avoid calling gpu methods

---
 prover/Cargo.lock                             |  1 +
 .../Cargo.toml                                |  1 +
 .../src/main.rs                               | 59 +++++++++++++++----
 3 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/prover/Cargo.lock b/prover/Cargo.lock
index 53c069a54b25..e1d7756c6ee5 100644
--- a/prover/Cargo.lock
+++ b/prover/Cargo.lock
@@ -8339,6 +8339,7 @@ dependencies = [
  "proof-compression-gpu",
  "proptest",
  "serde_json",
+ "shivini",
  "toml_edit 0.14.4",
  "tracing",
  "tracing-subscriber",
diff --git a/prover/crates/bin/vk_setup_data_generator_server_fri/Cargo.toml b/prover/crates/bin/vk_setup_data_generator_server_fri/Cargo.toml
index 0168f480a023..a8bd66b10f6e 100644
--- a/prover/crates/bin/vk_setup_data_generator_server_fri/Cargo.toml
+++ b/prover/crates/bin/vk_setup_data_generator_server_fri/Cargo.toml
@@ -24,6 +24,7 @@ zksync_utils.workspace = true
 zkevm_test_harness.workspace = true
 circuit_definitions = { workspace = true, features = ["log_tracing"] }
 proof-compression-gpu.workspace = true
+shivini.workspace = true
 
 anyhow.workspace = true
 clap = { workspace = true, features = ["derive"] }
diff --git a/prover/crates/bin/vk_setup_data_generator_server_fri/src/main.rs b/prover/crates/bin/vk_setup_data_generator_server_fri/src/main.rs
index 514ef3bae7b5..1ddd2c837a2a 100644
--- a/prover/crates/bin/vk_setup_data_generator_server_fri/src/main.rs
+++ b/prover/crates/bin/vk_setup_data_generator_server_fri/src/main.rs
@@ -1,9 +1,14 @@
+#![feature(allocator_api)]
+
 //! Tool to generate different types of keys used by the proving system.
 //!
 //! It can generate verification keys, setup keys, and also commitments.
-use std::{collections::HashMap, path::PathBuf};
-
 use anyhow::Context as _;
+use circuit_definitions::boojum::config::SetupCSConfig;
+use circuit_definitions::circuit_definitions::aux_layer::compression_modes::{
+    CompressionTranscriptForWrapper, CompressionTreeHasherForWrapper,
+};
+use circuit_definitions::circuit_definitions::aux_layer::ZkSyncCompressionForWrapperCircuit;
 use circuit_definitions::{
     circuit_definitions::aux_layer::{ZkSyncCompressionProof, ZkSyncCompressionVerificationKey},
     snark_wrapper::franklin_crypto,
@@ -12,6 +17,9 @@ use clap::{Parser, Subcommand};
 use commitment_generator::read_and_update_contract_toml;
 use indicatif::{ProgressBar, ProgressStyle};
 use proof_compression_gpu::{CompressionInput, CompressionMode};
+use shivini::cs::{GpuProverSetupData, GpuSetup};
+use std::alloc::Global;
+use std::{collections::HashMap, path::PathBuf};
 use tracing::level_filters::LevelFilter;
 use zkevm_test_harness::{
     compute_setups::{
@@ -26,7 +34,7 @@ use zkevm_test_harness::{
 };
 use zksync_prover_fri_types::{
     circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerStorageType,
-    FriProofWrapper, ProverServiceDataKey,
+    ProverServiceDataKey,
 };
 use zksync_prover_keystore::{
     keystore::Keystore,
@@ -248,23 +256,48 @@ fn generate_setup_keys(
 }
 
 fn generate_fflonk_setup_keys(keystore: &Keystore) -> anyhow::Result<()> {
-    let mut setup_data = None;
     let worker = franklin_crypto::boojum::worker::Worker::new();
     let proof_bytes = std::fs::read("data/compression_proof.bin").unwrap();
     let proof: ZkSyncCompressionProof = bincode::deserialize(&proof_bytes).unwrap();
     let vk_bytes = std::fs::read("data/compression_vk.json").unwrap();
     let vk: ZkSyncCompressionVerificationKey = serde_json::from_slice(&vk_bytes).unwrap();
     let input = CompressionInput::CompressionWrapper(Some(proof), vk, CompressionMode::Five);
-    proof_compression_gpu::prove_compression_wrapper_circuit(
-        input.into_compression_wrapper_circuit(),
-        &mut setup_data,
-        &worker,
-    );
-    if let Some(setup_data) = setup_data {
-        keystore.save_setup_data_for_fflonk(setup_data)?;
-    } else {
-        unreachable!("Proof compression was supposed to produce setup data but did not")
+
+    let setup_data = {
+        let circuit = input.into_compression_wrapper_circuit();
+        let proof_cfg = circuit.proof_config_for_compression_step();
+        match circuit {
+            ZkSyncCompressionForWrapperCircuit::CompressionMode5Circuit(inner) => {
+                let (setup_cs, finalization_hint) =
+                    shivini::synthesis_utils::synthesize_compression_circuit::<
+                        _,
+                        SetupCSConfig,
+                        Global,
+                    >(inner, true, None);
+                let (base_setup, _setup, setup_tree, variables_hint, witnesses_hint, vk) =
+                    setup_cs.prepare_base_setup_with_precomputations_and_vk::<CompressionTranscriptForWrapper, CompressionTreeHasherForWrapper>(
+                        proof_cfg.clone(),
+                        &worker
+                    );
+                let gpu_setup = GpuSetup::from_setup_and_hints(
+                    base_setup,
+                    setup_tree,
+                    variables_hint,
+                    witnesses_hint,
+                    &worker,
+                )
+                .unwrap();
+                GpuProverSetupData {
+                    setup: gpu_setup,
+                    vk,
+                    finalization_hint: finalization_hint.unwrap(),
+                }
+            }
+            _ => unreachable!(),
+        }
     };
+
+    keystore.save_setup_data_for_fflonk(setup_data)?;
     Ok(())
 }