fix: Move JoltDevice to preprocessing

a16z · Oct 31, 2024 · dcbd290 · dcbd290
1 parent e61e51e
commit dcbd290
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 28 deletions.
diff --git a/jolt-core/src/benches/bench.rs b/jolt-core/src/benches/bench.rs
@@ -131,11 +131,10 @@ where
         let (io_device, trace) = program.trace();
 
         let preprocessing: crate::jolt::vm::JoltPreprocessing<C, F, PCS, ProofTranscript> =
-            RV32IJoltVM::preprocess(bytecode.clone(), memory_init, 1 << 20, 1 << 20, 1 << 22);
+            RV32IJoltVM::preprocess(bytecode.clone(), io_device, memory_init, 1 << 20, 1 << 20, 1 << 22);
 
         let (jolt_proof, jolt_commitments, _) =
             <RV32IJoltVM as Jolt<_, PCS, C, M, ProofTranscript>>::prove(
-                io_device,
                 trace,
                 preprocessing.clone(),
             );
@@ -187,11 +186,10 @@ where
         let (io_device, trace) = program.trace();
 
         let preprocessing: crate::jolt::vm::JoltPreprocessing<C, F, PCS, ProofTranscript> =
-            RV32IJoltVM::preprocess(bytecode.clone(), memory_init, 1 << 20, 1 << 20, 1 << 22);
+            RV32IJoltVM::preprocess(bytecode.clone(), io_device, memory_init, 1 << 20, 1 << 20, 1 << 22);
 
         let (jolt_proof, jolt_commitments, _) =
             <RV32IJoltVM as Jolt<_, PCS, C, M, ProofTranscript>>::prove(
-                io_device,
                 trace,
                 preprocessing.clone(),
             );

diff --git a/jolt-core/src/jolt/vm/mod.rs b/jolt-core/src/jolt/vm/mod.rs
@@ -60,6 +60,7 @@ where
     pub instruction_lookups: InstructionLookupsPreprocessing<C, F>,
     pub bytecode: BytecodePreprocessing<F>,
     pub read_write_memory: ReadWriteMemoryPreprocessing,
+    pub program_io: JoltDevice,
 }
 
 #[derive(Clone, Serialize, Deserialize, Debug)]
@@ -123,7 +124,6 @@ pub struct JoltProof<
     ProofTranscript: Transcript,
 {
     pub trace_length: usize,
-    pub program_io: JoltDevice,
     pub bytecode: BytecodeProof<F, PCS, ProofTranscript>,
     pub read_write_memory: ReadWriteMemoryProof<F, PCS, ProofTranscript>,
     pub instruction_lookups:
@@ -276,6 +276,7 @@ where
     #[tracing::instrument(skip_all, name = "Jolt::preprocess")]
     fn preprocess(
         bytecode: Vec<ELFInstruction>,
+        program_io: JoltDevice,
         memory_init: Vec<(u64, u8)>,
         max_bytecode_size: usize,
         max_memory_address: usize,
@@ -339,12 +340,12 @@ where
             instruction_lookups: instruction_lookups_preprocessing,
             bytecode: bytecode_preprocessing,
             read_write_memory: read_write_memory_preprocessing,
+            program_io,
         }
     }
 
     #[tracing::instrument(skip_all, name = "Jolt::prove")]
     fn prove(
-        program_io: JoltDevice,
         mut trace: Vec<JoltTraceStep<Self::InstructionSet>>,
         preprocessing: JoltPreprocessing<C, F, PCS, ProofTranscript>,
     ) -> (
@@ -368,7 +369,7 @@ where
         JoltTraceStep::pad(&mut trace);
 
         let mut transcript = ProofTranscript::new(b"Jolt transcript");
-        Self::fiat_shamir_preamble(&mut transcript, &program_io, trace_length);
+        Self::fiat_shamir_preamble(&mut transcript, &preprocessing.program_io, trace_length);
 
         let instruction_polynomials =
             InstructionLookupsProof::<
@@ -383,7 +384,7 @@ where
 
         let load_store_flags = &instruction_polynomials.instruction_flags[5..10];
         let (memory_polynomials, read_timestamps) = ReadWriteMemoryPolynomials::generate_witness(
-            &program_io,
+            &preprocessing.program_io,
             load_store_flags,
             &preprocessing.read_write_memory,
             &trace,
@@ -405,7 +406,7 @@ where
 
         let r1cs_builder = Self::Constraints::construct_constraints(
             padded_trace_length,
-            RAM_START_ADDRESS - program_io.memory_layout.ram_witness_offset,
+            RAM_START_ADDRESS - preprocessing.program_io.memory_layout.ram_witness_offset,
         );
         let spartan_key = spartan::UniformSpartanProof::<
             C,
@@ -468,7 +469,7 @@ where
             &preprocessing.generators,
             &preprocessing.read_write_memory,
             &jolt_polynomials,
-            &program_io,
+            &preprocessing.program_io,
             &mut opening_accumulator,
             &mut transcript,
         );
@@ -495,7 +496,6 @@ where
 
         let jolt_proof = JoltProof {
             trace_length,
-            program_io,
             bytecode: bytecode_proof,
             read_write_memory: memory_proof,
             instruction_lookups: instruction_proof,
@@ -539,11 +539,11 @@ where
             opening_accumulator
                 .compare_to(debug_info.opening_accumulator, &preprocessing.generators);
         }
-        Self::fiat_shamir_preamble(&mut transcript, &proof.program_io, proof.trace_length);
+        Self::fiat_shamir_preamble(&mut transcript, &preprocessing.program_io, proof.trace_length);
 
         // Regenerate the uniform Spartan key
         let padded_trace_length = proof.trace_length.next_power_of_two();
-        let memory_start = RAM_START_ADDRESS - proof.program_io.memory_layout.ram_witness_offset;
+        let memory_start = RAM_START_ADDRESS - preprocessing.program_io.memory_layout.ram_witness_offset;
         let r1cs_builder =
             Self::Constraints::construct_constraints(padded_trace_length, memory_start);
         let spartan_key = spartan::UniformSpartanProof::<C, _, F, ProofTranscript>::setup(
@@ -588,7 +588,7 @@ where
             &preprocessing.generators,
             proof.read_write_memory,
             &commitments,
-            proof.program_io,
+            preprocessing.program_io,
             &mut opening_accumulator,
             &mut transcript,
         )?;

diff --git a/jolt-core/src/jolt/vm/rv32i_vm.rs b/jolt-core/src/jolt/vm/rv32i_vm.rs
@@ -313,10 +313,9 @@ mod tests {
         drop(artifact_guard);
 
         let preprocessing =
-            RV32IJoltVM::preprocess(bytecode.clone(), memory_init, 1 << 20, 1 << 20, 1 << 20);
+            RV32IJoltVM::preprocess(bytecode.clone(), io_device, memory_init, 1 << 20, 1 << 20, 1 << 20);
         let (proof, commitments, debug_info) =
             <RV32IJoltVM as Jolt<F, PCS, C, M, ProofTranscript>>::prove(
-                io_device,
                 trace,
                 preprocessing.clone(),
             );
@@ -372,15 +371,15 @@ mod tests {
         let (io_device, trace) = program.trace();
 
         let preprocessing =
-            RV32IJoltVM::preprocess(bytecode.clone(), memory_init, 1 << 20, 1 << 20, 1 << 20);
+            RV32IJoltVM::preprocess(bytecode.clone(), io_device, memory_init, 1 << 20, 1 << 20, 1 << 20);
         let (jolt_proof, jolt_commitments, debug_info) =
             <RV32IJoltVM as Jolt<
                 _,
                 HyraxScheme<G1Projective, KeccakTranscript>,
                 C,
                 M,
                 KeccakTranscript,
-            >>::prove(io_device, trace, preprocessing.clone());
+            >>::prove(trace, preprocessing.clone());
         let verification_result =
             RV32IJoltVM::verify(preprocessing, jolt_proof, jolt_commitments, debug_info);
         assert!(
@@ -402,15 +401,15 @@ mod tests {
         drop(guard);
 
         let preprocessing =
-            RV32IJoltVM::preprocess(bytecode.clone(), memory_init, 1 << 20, 1 << 20, 1 << 20);
+            RV32IJoltVM::preprocess(bytecode.clone(), io_device, memory_init, 1 << 20, 1 << 20, 1 << 20);
         let (jolt_proof, jolt_commitments, debug_info) =
             <RV32IJoltVM as Jolt<
                 _,
                 HyraxScheme<G1Projective, KeccakTranscript>,
                 C,
                 M,
                 KeccakTranscript,
-            >>::prove(io_device, trace, preprocessing.clone());
+            >>::prove(trace, preprocessing.clone());
 
         let verification_result =
             RV32IJoltVM::verify(preprocessing, jolt_proof, jolt_commitments, debug_info);
@@ -432,15 +431,15 @@ mod tests {
         drop(guard);
 
         let preprocessing =
-            RV32IJoltVM::preprocess(bytecode.clone(), memory_init, 1 << 20, 1 << 20, 1 << 20);
+            RV32IJoltVM::preprocess(bytecode.clone(), io_device, memory_init, 1 << 20, 1 << 20, 1 << 20);
         let (jolt_proof, jolt_commitments, debug_info) = <RV32IJoltVM as Jolt<
             _,
             Zeromorph<Bn254, KeccakTranscript>,
             C,
             M,
             KeccakTranscript,
         >>::prove(
-            io_device, trace, preprocessing.clone()
+            trace, preprocessing.clone()
         );
 
         let verification_result =
@@ -463,15 +462,15 @@ mod tests {
         drop(guard);
 
         let preprocessing =
-            RV32IJoltVM::preprocess(bytecode.clone(), memory_init, 1 << 20, 1 << 20, 1 << 20);
+            RV32IJoltVM::preprocess(bytecode.clone(), io_device, memory_init, 1 << 20, 1 << 20, 1 << 20);
         let (jolt_proof, jolt_commitments, debug_info) = <RV32IJoltVM as Jolt<
             _,
             HyperKZG<Bn254, KeccakTranscript>,
             C,
             M,
             KeccakTranscript,
         >>::prove(
-            io_device, trace, preprocessing.clone()
+            trace, preprocessing.clone()
         );
 
         let verification_result =
@@ -496,15 +495,15 @@ mod tests {
         drop(artifact_guard);
 
         let preprocessing =
-            RV32IJoltVM::preprocess(bytecode.clone(), memory_init, 1 << 20, 1 << 20, 1 << 20);
+            RV32IJoltVM::preprocess(bytecode.clone(), io_device, memory_init, 1 << 20, 1 << 20, 1 << 20);
         let (proof, commitments, debug_info) = <RV32IJoltVM as Jolt<
             Fr,
             HyperKZG<Bn254, KeccakTranscript>,
             C,
             M,
             KeccakTranscript,
         >>::prove(
-            io_device, trace, preprocessing.clone()
+            trace, preprocessing.clone()
         );
         let verification_result =
             RV32IJoltVM::verify(preprocessing, proof, commitments, debug_info);

diff --git a/tracer/src/emulator/mmu.rs b/tracer/src/emulator/mmu.rs
@@ -23,7 +23,7 @@ use super::terminal::Terminal;
 
 /// Emulates Memory Management Unit. It holds the Main memory and peripheral
 /// devices, maps address to them, and accesses them depending on address.
-/// It also manages virtual-physical address translation and memoty protection.
+/// It also manages virtual-physical address translation and memory protection.
 /// It may also be said Bus.
 /// @TODO: Memory protection is not implemented yet. We should support.
 pub struct Mmu {
@@ -494,7 +494,7 @@ impl Mmu {
             false => match effective_address {
                 // I don't know why but dtb data seems to be stored from 0x1020 on Linux.
                 // It might be from self.x[0xb] initialization?
-                // And DTB size is arbitray.
+                // And DTB size is arbitrary.
                 0x00001020..=0x00001fff => self.dtb[effective_address as usize - 0x1020],
                 0x02000000..=0x0200ffff => self.clint.load(effective_address),
                 0x0C000000..=0x0fffffff => self.plic.load(effective_address),