refactor: remove shared preprocessing (#1069)

Author: andrewmilson

examples/recursion/src/main.rs

@@ -477,17 +477,12 @@ fn run_recursion_proof(
         jolt_sdk::JoltVerifierPreprocessing::from(&recursion_prover_preprocessing);
 
     // update program_size in memory_config now that we know it
-    recursion.memory_config.program_size = Some(
-        recursion_verifier_preprocessing
-            .shared
-            .memory_layout
-            .program_size,
-    );
+    recursion.memory_config.program_size =
+        Some(recursion_verifier_preprocessing.memory_layout.program_size);
 
     let mut output_bytes = vec![
         0;
         recursion_verifier_preprocessing
-            .shared
             .memory_layout
             .max_output_size as usize
     ];

jolt-core/src/guest/verifier.rs

@@ -45,13 +45,13 @@ where
 {
     use common::jolt_device::JoltDevice;
     let memory_config = MemoryConfig {
-        max_untrusted_advice_size: preprocessing.shared.memory_layout.max_untrusted_advice_size,
-        max_trusted_advice_size: preprocessing.shared.memory_layout.max_trusted_advice_size,
-        max_input_size: preprocessing.shared.memory_layout.max_input_size,
-        max_output_size: preprocessing.shared.memory_layout.max_output_size,
-        stack_size: preprocessing.shared.memory_layout.stack_size,
-        memory_size: preprocessing.shared.memory_layout.memory_size,
-        program_size: Some(preprocessing.shared.memory_layout.program_size),
+        max_untrusted_advice_size: preprocessing.memory_layout.max_untrusted_advice_size,
+        max_trusted_advice_size: preprocessing.memory_layout.max_trusted_advice_size,
+        max_input_size: preprocessing.memory_layout.max_input_size,
+        max_output_size: preprocessing.memory_layout.max_output_size,
+        stack_size: preprocessing.memory_layout.stack_size,
+        memory_size: preprocessing.memory_layout.memory_size,
+        program_size: Some(preprocessing.memory_layout.program_size),
     };
     let mut io_device = JoltDevice::new(&memory_config);
 

jolt-core/src/zkvm/bytecode/mod.rs

@@ -171,7 +171,7 @@ fn gen_ra_one_hot_provers<F: JoltField>(
     transcript: &mut impl Transcript,
 ) -> (HammingWeightSumcheckProver<F>, BooleanitySumcheckProver<F>) {
     let (preprocessing, _, trace, _, _) = state_manager.get_prover_data();
-    let bytecode_preprocessing = &preprocessing.shared.bytecode;
+    let bytecode_preprocessing = &preprocessing.bytecode;
 
     let r_cycle: Vec<F::Challenge> = opening_accumulator
         .get_virtual_polynomial_opening(VirtualPolynomial::UnexpandedPC, SumcheckId::SpartanOuter)

jolt-core/src/zkvm/bytecode/read_raf_checking.rs

@@ -131,7 +131,7 @@ impl<F: JoltField> ReadRafSumcheckProver<F> {
         let (preprocessing, _, trace, _, _) = state_manager.get_prover_data();
 
         let params = ReadRafSumcheckParams::gen(
-            &preprocessing.shared.bytecode,
+            &preprocessing.bytecode,
             trace.len().log_2(),
             opening_accumulator,
             transcript,
@@ -169,7 +169,7 @@ impl<F: JoltField> ReadRafSumcheckProver<F> {
                     array::from_fn(|_| unsafe_allocate_zero_vec::<F>(params.K));
 
                 for (j, cycle) in trace_chunk.iter().enumerate() {
-                    let pc = preprocessing.shared.bytecode.get_pc(cycle);
+                    let pc = preprocessing.bytecode.get_pc(cycle);
                     for stage in 0..N_STAGES {
                         res_per_stage[stage][pc] += eq_evals[stage][j]
                     }
@@ -221,7 +221,7 @@ impl<F: JoltField> ReadRafSumcheckProver<F> {
 
         let pc = trace
             .par_iter()
-            .map(|cycle| preprocessing.shared.bytecode.get_pc(cycle))
+            .map(|cycle| preprocessing.bytecode.get_pc(cycle))
             .collect();
 
         Self {

jolt-core/src/zkvm/dag/jolt_dag.rs

@@ -88,9 +88,9 @@ pub fn prove_jolt_dag<
     let trace_length = trace.len();
     let padded_trace_length = trace_length.next_power_of_two();
 
-    tracing::info!("bytecode size: {}", preprocessing.shared.bytecode.code_size);
+    tracing::info!("bytecode size: {}", preprocessing.bytecode.code_size);
 
-    let bytecode_d = preprocessing.shared.bytecode.d;
+    let bytecode_d = preprocessing.bytecode.d;
 
     // Commit to untrusted_advice
     let _untrusted_advice_opening_proof_hints =
@@ -524,7 +524,7 @@ pub fn prove_jolt_dag<
         reduced_opening_proof,
         trace_length: prover_state.trace.len(),
         ram_K: state_manager.ram_K,
-        bytecode_d: prover_state.preprocessing.shared.bytecode.d,
+        bytecode_d: prover_state.preprocessing.bytecode.d,
         twist_sumcheck_switch_index: state_manager.twist_sumcheck_switch_index,
     };
 
@@ -554,7 +554,7 @@ impl<'a, 'b, 'c, F: JoltField, ProofTranscript: Transcript, PCS: CommitmentSchem
         self.state_manager.fiat_shamir_preamble(self.transcript);
 
         let ram_K = self.proof.ram_K;
-        let bytecode_d = self.state_manager.get_verifier_data().0.shared.bytecode.d;
+        let bytecode_d = self.preprocessing.bytecode.d;
         let _guard = AllCommittedPolynomials::initialize(ram_K, bytecode_d);
 
         // Append commitments to transcript
@@ -712,7 +712,7 @@ impl<'a, 'b, 'c, F: JoltField, ProofTranscript: Transcript, PCS: CommitmentSchem
         );
         let initial_ram_state = ram::gen_ram_initial_memory_state::<F>(
             self.proof.ram_K,
-            &self.preprocessing.shared.ram,
+            &self.preprocessing.ram,
             &self.state_manager.program_io,
         );
         let ram_val_evaluation = RamValEvaluationSumcheckVerifier::new(
@@ -777,13 +777,13 @@ impl<'a, 'b, 'c, F: JoltField, ProofTranscript: Transcript, PCS: CommitmentSchem
     fn verify_stage6(&mut self) -> Result<(), anyhow::Error> {
         let n_cycle_vars = self.proof.trace_length.log_2();
         let bytecode_read_raf = BytecodeReadRafSumcheckVerifier::gen(
-            &self.preprocessing.shared.bytecode,
+            &self.preprocessing.bytecode,
             n_cycle_vars,
             &self.opening_accumulator,
             self.transcript,
         );
         let (bytecode_hamming_weight, bytecode_booleanity) = bytecode::new_ra_one_hot_verifiers(
-            &self.preprocessing.shared.bytecode,
+            &self.preprocessing.bytecode,
             n_cycle_vars,
             self.transcript,
         );

jolt-core/src/zkvm/dag/state_manager.rs

@@ -6,10 +6,10 @@ use crate::poly::multilinear_polynomial::MultilinearPolynomial;
 use crate::transcripts::Transcript;
 use crate::utils::math::Math;
 use crate::zkvm::witness::compute_d_parameter;
-use crate::zkvm::{JoltProverPreprocessing, JoltSharedPreprocessing, JoltVerifierPreprocessing};
+use crate::zkvm::{JoltProverPreprocessing, JoltVerifierPreprocessing};
 use rayon::prelude::*;
 use tracer::emulator::memory::Memory;
-use tracer::instruction::{Cycle, Instruction};
+use tracer::instruction::Cycle;
 use tracer::{JoltDevice, LazyTraceIterator};
 
 pub struct ProverState<'a, F: JoltField, PCS>
@@ -62,18 +62,18 @@ where
             .filter_map(|cycle| {
                 crate::zkvm::ram::remap_address(
                     cycle.ram_access().address() as u64,
-                    &preprocessing.shared.memory_layout,
+                    &preprocessing.memory_layout,
                 )
             })
             .max()
             .unwrap_or(0)
             .max(
                 crate::zkvm::ram::remap_address(
-                    preprocessing.shared.ram.min_bytecode_address,
-                    &preprocessing.shared.memory_layout,
+                    preprocessing.ram.min_bytecode_address,
+                    &preprocessing.memory_layout,
                 )
                 .unwrap_or(0)
-                    + preprocessing.shared.ram.bytecode_words.len() as u64
+                    + preprocessing.ram.bytecode_words.len() as u64
                     + 1,
             )
             .next_power_of_two() as usize;
@@ -103,33 +103,6 @@ where
         }
     }
 
-    /// Only used in tests; in practice, the verifier state manager is
-    /// constructed using `JoltProof::to_verifier_state_manager`
-    #[cfg(test)]
-    pub fn new_verifier(
-        preprocessing: &'a JoltVerifierPreprocessing<F, PCS>,
-        program_io: JoltDevice,
-        trace_length: usize,
-        ram_K: usize,
-        twist_sumcheck_switch_index: usize,
-    ) -> Self {
-        let ram_d = compute_d_parameter(ram_K);
-
-        StateManager {
-            untrusted_advice_commitment: None,
-            trusted_advice_commitment: None,
-            program_io,
-            ram_K,
-            ram_d,
-            twist_sumcheck_switch_index,
-            prover_state: None,
-            verifier_state: Some(VerifierState {
-                preprocessing,
-                trace_length,
-            }),
-        }
-    }
-
     pub fn get_prover_data(
         &self,
     ) -> (
@@ -160,48 +133,6 @@ where
         }
     }
 
-    pub fn get_verifier_data(&self) -> (&'a JoltVerifierPreprocessing<F, PCS>, &JoltDevice, usize) {
-        if let Some(ref verifier_state) = self.verifier_state {
-            (
-                verifier_state.preprocessing,
-                &self.program_io,
-                verifier_state.trace_length,
-            )
-        } else {
-            panic!("Verifier state not initialized");
-        }
-    }
-
-    pub fn get_trace_len(&self) -> usize {
-        if let Some(ref verifier_state) = self.verifier_state {
-            verifier_state.trace_length
-        } else if let Some(ref prover_state) = self.prover_state {
-            prover_state.trace.len()
-        } else {
-            panic!("Neither prover nor verifier state initialized");
-        }
-    }
-
-    pub fn get_bytecode(&self) -> &[Instruction] {
-        if let Some(ref verifier_state) = self.verifier_state {
-            &verifier_state.preprocessing.shared.bytecode.bytecode
-        } else if let Some(ref prover_state) = self.prover_state {
-            &prover_state.preprocessing.shared.bytecode.bytecode
-        } else {
-            panic!("Neither prover nor verifier state initialized");
-        }
-    }
-
-    pub fn get_shared_preprocessing(&self) -> &JoltSharedPreprocessing {
-        if let Some(ref verifier_state) = self.verifier_state {
-            &verifier_state.preprocessing.shared
-        } else if let Some(ref prover_state) = self.prover_state {
-            &prover_state.preprocessing.shared
-        } else {
-            panic!("Neither prover nor verifier state initialized");
-        }
-    }
-
     pub fn fiat_shamir_preamble(&self, transcript: &mut impl Transcript) {
         transcript.append_u64(self.program_io.memory_layout.max_input_size);
         transcript.append_u64(self.program_io.memory_layout.max_output_size);

jolt-core/src/zkvm/instruction_lookups/read_raf_checking.rs

@@ -1204,10 +1204,7 @@ mod tests {
     use crate::transcripts::Blake2bTranscript;
     use crate::{
         poly::commitment::mock::MockCommitScheme,
-        zkvm::{
-            bytecode::BytecodePreprocessing, ram::RAMPreprocessing, JoltProverPreprocessing,
-            JoltSharedPreprocessing,
-        },
+        zkvm::{bytecode::BytecodePreprocessing, ram::RAMPreprocessing, JoltProverPreprocessing},
     };
     use ark_bn254::Fr;
     use ark_std::Zero;
@@ -1302,15 +1299,12 @@ mod tests {
         let bytecode = vec![];
         let bytecode_preprocessing = BytecodePreprocessing::preprocess(bytecode);
         let memory_layout = MemoryLayout::default();
-        let shared_preprocessing = JoltSharedPreprocessing {
-            bytecode: bytecode_preprocessing,
-            ram: RAMPreprocessing::preprocess(vec![]),
-            memory_layout: memory_layout.clone(),
-        };
         let prover_preprocessing: JoltProverPreprocessing<Fr, MockCommitScheme<Fr>> =
             JoltProverPreprocessing {
                 generators: (),
-                shared: shared_preprocessing.clone(),
+                bytecode: bytecode_preprocessing,
+                ram: RAMPreprocessing::preprocess(vec![]),
+                memory_layout: memory_layout.clone(),
             };
 
         let program_io = JoltDevice {

jolt-core/src/zkvm/mod.rs

@@ -124,7 +124,9 @@ where
     PCS: CommitmentScheme<Field = F>,
 {
     pub generators: PCS::VerifierSetup,
-    pub shared: JoltSharedPreprocessing,
+    pub bytecode: BytecodePreprocessing,
+    pub ram: RAMPreprocessing,
+    pub memory_layout: MemoryLayout,
 }
 
 impl<F, PCS> Serializable for JoltVerifierPreprocessing<F, PCS>
@@ -164,7 +166,9 @@ where
     PCS: CommitmentScheme<Field = F>,
 {
     pub generators: PCS::ProverSetup,
-    pub shared: JoltSharedPreprocessing,
+    pub bytecode: BytecodePreprocessing,
+    pub ram: RAMPreprocessing,
+    pub memory_layout: MemoryLayout,
 }
 
 impl<F, PCS> Serializable for JoltProverPreprocessing<F, PCS>
@@ -206,7 +210,9 @@ where
         let generators = PCS::setup_verifier(&preprocessing.generators);
         JoltVerifierPreprocessing {
             generators,
-            shared: preprocessing.shared.clone(),
+            bytecode: preprocessing.bytecode.clone(),
+            ram: preprocessing.ram.clone(),
+            memory_layout: preprocessing.memory_layout.clone(),
         }
     }
 }
@@ -224,35 +230,25 @@ where
 }
 
 pub trait Jolt<F: JoltField, PCS: StreamingCommitmentScheme<Field = F>, FS: Transcript> {
-    fn shared_preprocess(
-        bytecode: Vec<Instruction>,
-        memory_layout: MemoryLayout,
-        memory_init: Vec<(u64, u8)>,
-    ) -> JoltSharedPreprocessing {
-        let bytecode_preprocessing = BytecodePreprocessing::preprocess(bytecode);
-        let ram_preprocessing = RAMPreprocessing::preprocess(memory_init);
-
-        JoltSharedPreprocessing {
-            memory_layout,
-            bytecode: bytecode_preprocessing,
-            ram: ram_preprocessing,
-        }
-    }
-
     #[tracing::instrument(skip_all, name = "Jolt::prover_preprocess")]
     fn prover_preprocess(
         bytecode: Vec<Instruction>,
         memory_layout: MemoryLayout,
         memory_init: Vec<(u64, u8)>,
         max_trace_length: usize,
     ) -> JoltProverPreprocessing<F, PCS> {
-        let shared = Self::shared_preprocess(bytecode, memory_layout, memory_init);
+        let bytecode = BytecodePreprocessing::preprocess(bytecode);
+        let ram = RAMPreprocessing::preprocess(memory_init);
 
         let max_T: usize = max_trace_length.next_power_of_two();
-
         let generators = PCS::setup_prover(DTH_ROOT_OF_K.log_2() + max_T.log_2());
 
-        JoltProverPreprocessing { generators, shared }
+        JoltProverPreprocessing {
+            generators,
+            bytecode,
+            ram,
+            memory_layout,
+        }
     }
 
     #[allow(clippy::type_complexity)]
@@ -276,13 +272,13 @@ pub trait Jolt<F: JoltField, PCS: StreamingCommitmentScheme<Field = F>, FS: Tran
         use tracer::instruction::Cycle;
 
         let memory_config = MemoryConfig {
-            max_untrusted_advice_size: preprocessing.shared.memory_layout.max_untrusted_advice_size,
-            max_trusted_advice_size: preprocessing.shared.memory_layout.max_trusted_advice_size,
-            max_input_size: preprocessing.shared.memory_layout.max_input_size,
-            max_output_size: preprocessing.shared.memory_layout.max_output_size,
-            stack_size: preprocessing.shared.memory_layout.stack_size,
-            memory_size: preprocessing.shared.memory_layout.memory_size,
-            program_size: Some(preprocessing.shared.memory_layout.program_size),
+            max_untrusted_advice_size: preprocessing.memory_layout.max_untrusted_advice_size,
+            max_trusted_advice_size: preprocessing.memory_layout.max_trusted_advice_size,
+            max_input_size: preprocessing.memory_layout.max_input_size,
+            max_output_size: preprocessing.memory_layout.max_output_size,
+            stack_size: preprocessing.memory_layout.stack_size,
+            memory_size: preprocessing.memory_layout.memory_size,
+            program_size: Some(preprocessing.memory_layout.program_size),
         };
 
         let (lazy_trace, mut trace, final_memory_state, mut program_io) = {
@@ -382,13 +378,13 @@ pub trait Jolt<F: JoltField, PCS: StreamingCommitmentScheme<Field = F>, FS: Tran
         let _guard = DoryGlobals::initialize(DTH_ROOT_OF_K, T);
 
         // Memory layout checks
-        if program_io.memory_layout != preprocessing.shared.memory_layout {
+        if program_io.memory_layout != preprocessing.memory_layout {
             return Err(ProofVerifyError::MemoryLayoutMismatch);
         }
-        if program_io.inputs.len() > preprocessing.shared.memory_layout.max_input_size as usize {
+        if program_io.inputs.len() > preprocessing.memory_layout.max_input_size as usize {
             return Err(ProofVerifyError::InputTooLarge);
         }
-        if program_io.outputs.len() > preprocessing.shared.memory_layout.max_output_size as usize {
+        if program_io.outputs.len() > preprocessing.memory_layout.max_output_size as usize {
             return Err(ProofVerifyError::OutputTooLarge);
         }