Update Jolt to v0.3.0-alpha (#181)

.github/workflows/test-zkvm-jolt.yml

@@ -14,4 +14,4 @@ jobs:
       packages: write
     with:
       zkvm: jolt
-      toolchain: 1.86.0
+      toolchain: 1.88.0

.github/workflows/test-zkvm-zisk.yml

@@ -15,5 +15,4 @@ jobs:
     with:
       zkvm: zisk
       toolchain: 1.86.0
-      free_up_disk_space: true
       skip_prove_test: true

.github/workflows/test-zkvm.yml

@@ -12,11 +12,6 @@ on:
         required: false
         type: string
         default: 1.86.0
-      free_up_disk_space:
-        description: 'Whether free up disk space or not'
-        required: false
-        type: boolean
-        default: false
       # Remove when we use larger runners, currently only needed to skip for zisk
       skip_prove_test:
         description: 'Whether to skip prove test and ere-dockerized test or not'
@@ -47,7 +42,6 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Free up disk space
-        if: ${{ inputs.free_up_disk_space }}
         run: bash .github/scripts/free-up-disk-space.sh
 
       - name: Set up Docker Buildx
@@ -151,7 +145,6 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Free up disk space
-        if: ${{ inputs.free_up_disk_space }}
         run: bash .github/scripts/free-up-disk-space.sh
 
       - name: Cache dependencies
@@ -201,7 +194,6 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Free up disk space
-        if: ${{ inputs.free_up_disk_space }}
         run: bash .github/scripts/free-up-disk-space.sh
 
       - name: Cache dependencies
@@ -249,7 +241,6 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Free up disk space
-        if: ${{ inputs.free_up_disk_space }}
         run: bash .github/scripts/free-up-disk-space.sh
 
       - name: Install protoc

.gitignore

@@ -1,2 +1,3 @@
 target
-notes.md
+notes.md
+dory_srs_*_variables.srs

Cargo.toml

@@ -68,9 +68,10 @@ twirp-build = "0.9.0"
 airbender_execution_utils = { git = "https://github.com/matter-labs/zksync-airbender", package = "execution_utils", tag = "v0.5.0" }
 
 # Jolt dependencies
-ark-serialize = "0.5.0"
-common = { git = "https://github.com/a16z/jolt.git", rev = "55b9830a3944dde55d33a55c42522b81dd49f87a" }
-jolt = { git = "https://github.com/a16z/jolt.git", rev = "55b9830a3944dde55d33a55c42522b81dd49f87a" }
+jolt-ark-serialize = { git = "https://github.com/a16z/arkworks-algebra", package = "ark-serialize", branch = "feat/fewer-reductions" }
+jolt-common = { git = "https://github.com/a16z/jolt.git", package = "common", tag = "v0.3.0-alpha" }
+jolt-core = { git = "https://github.com/a16z/jolt.git", tag = "v0.3.0-alpha" }
+jolt-sdk = { git = "https://github.com/a16z/jolt.git", tag = "v0.3.0-alpha" }
 
 # Miden dependencies
 miden-assembly = { git = "https://github.com/0xPolygonMiden/miden-vm.git", tag = "v0.17.1" }
@@ -137,6 +138,8 @@ opt-level = 3
 
 [patch.crates-io]
 # These patches are only needed by Jolt
-ark-ff = { git = "https://github.com/a16z/arkworks-algebra", branch = "v0.5.0-optimize-mul-u64" }
-ark-ec = { git = "https://github.com/a16z/arkworks-algebra", branch = "v0.5.0-optimize-mul-u64" }
-ark-serialize = { git = "https://github.com/a16z/arkworks-algebra", branch = "v0.5.0-optimize-mul-u64" }
+ark-ff = { git = "https://github.com/a16z/arkworks-algebra", branch = "feat/fewer-reductions" }
+ark-ec = { git = "https://github.com/a16z/arkworks-algebra", branch = "feat/fewer-reductions" }
+jolt-optimizations = { git = "https://github.com/a16z/arkworks-algebra", branch = "feat/fewer-reductions" }
+ark-serialize = { git = "https://github.com/a16z/arkworks-algebra", branch = "feat/fewer-reductions" }
+allocative = { git = "https://github.com/facebookexperimental/allocative", rev = "85b773d85d526d068ce94724ff7a7b81203fc95e" }

crates/dockerized/compiler/src/main.rs

@@ -56,9 +56,9 @@ fn compile(guest_path: PathBuf) -> Result<impl Serialize, Error> {
 
     #[cfg(feature = "jolt")]
     let result = if use_stock_rust() {
-        ere_jolt::compiler::RustRv32ima.compile(&guest_path)
+        ere_jolt::compiler::RustRv64imac.compile(&guest_path)
     } else {
-        ere_jolt::compiler::RustRv32imaCustomized.compile(&guest_path)
+        ere_jolt::compiler::RustRv64imacCustomized.compile(&guest_path)
     };
 
     #[cfg(feature = "miden")]

crates/dockerized/dockerized/src/lib.rs

@@ -654,6 +654,16 @@ mod test {
 
     mod jolt {
         test_compile!(Jolt, "basic");
+        test_execute!(
+            Jolt,
+            BasicProgramInput::valid(),
+            [Vec::new(), BasicProgramInput::invalid().serialized_input()]
+        );
+        test_prove!(
+            Jolt,
+            BasicProgramInput::valid(),
+            [Vec::new(), BasicProgramInput::invalid().serialized_input()]
+        );
     }
 
     mod miden {

crates/zkvm/airbender/src/lib.rs

@@ -146,7 +146,6 @@ mod tests {
         let zkvm = EreAirbender::new(program, ProverResourceType::Cpu).unwrap();
 
         let test_case = BasicProgramInput::valid().into_output_sha256();
-        run_zkvm_execute(&zkvm, &test_case);
         run_zkvm_prove(&zkvm, &test_case);
     }
 

crates/zkvm/jolt/Cargo.toml

@@ -11,9 +11,10 @@ tempfile.workspace = true
 thiserror.workspace = true
 
 # Jolt dependencies
-ark-serialize = { workspace = true, features = ["derive"] }
-common.workspace = true
-jolt = { workspace = true, features = ["host"] }
+jolt-ark-serialize = { workspace = true, features = ["derive"] }
+jolt-common.workspace = true
+jolt-core.workspace = true
+jolt-sdk = { workspace = true, features = ["host"] }
 
 # Local dependencies
 ere-compile-utils.workspace = true

crates/zkvm/jolt/src/client.rs

@@ -0,0 +1,123 @@
+use crate::error::JoltError;
+use ere_zkvm_interface::{CommonError, PublicValues};
+use jolt_ark_serialize::{self as ark_serialize, CanonicalDeserialize, CanonicalSerialize};
+use jolt_common::constants::{
+    DEFAULT_MAX_INPUT_SIZE, DEFAULT_MAX_OUTPUT_SIZE, DEFAULT_MAX_TRACE_LENGTH, DEFAULT_MEMORY_SIZE,
+    DEFAULT_STACK_SIZE,
+};
+use jolt_core::{
+    poly::commitment::commitment_scheme::CommitmentScheme, transcripts::Blake2bTranscript as FS,
+    utils::math::Math, zkvm::witness::DTH_ROOT_OF_K,
+};
+use jolt_sdk::{
+    F, Jolt, JoltDevice, JoltProverPreprocessing, JoltRV64IMAC, JoltVerifierPreprocessing,
+    MemoryConfig, MemoryLayout, PCS,
+    guest::program::{decode, trace},
+    postcard,
+};
+
+#[derive(CanonicalSerialize, CanonicalDeserialize)]
+pub struct JoltProof {
+    proof: jolt_sdk::JoltProof<F, PCS, FS>,
+    // FIXME: Remove `inputs` when Jolt supports proving with private input.
+    //        Issue for tracking: https://github.com/eth-act/ere/issues/4.
+    inputs: Vec<u8>,
+    outputs: Vec<u8>,
+}
+
+pub struct JoltSdk {
+    elf: Vec<u8>,
+    memory_config: MemoryConfig,
+    pk: JoltProverPreprocessing<F, PCS>,
+    vk: JoltVerifierPreprocessing<F, PCS>,
+}
+
+impl JoltSdk {
+    pub fn new(elf: &[u8]) -> Self {
+        let (bytecode, memory_init, program_size) = decode(elf);
+        let memory_config = MemoryConfig {
+            max_input_size: DEFAULT_MAX_INPUT_SIZE,
+            max_output_size: DEFAULT_MAX_OUTPUT_SIZE,
+            stack_size: DEFAULT_STACK_SIZE,
+            memory_size: DEFAULT_MEMORY_SIZE,
+            program_size: Some(program_size),
+        };
+        let memory_layout = MemoryLayout::new(&memory_config);
+        let max_trace_length = DEFAULT_MAX_TRACE_LENGTH as usize;
+        let pk = {
+            // FIXME: Use public trusted setup or switch to other transparent PCS.
+            let max_trace_length = max_trace_length.next_power_of_two();
+            let generators = PCS::setup_prover(DTH_ROOT_OF_K.log_2() + max_trace_length.log_2());
+
+            let shared = JoltRV64IMAC::shared_preprocess(bytecode, memory_layout, memory_init);
+
+            JoltProverPreprocessing { generators, shared }
+        };
+        let vk = JoltVerifierPreprocessing::from(&pk);
+        Self {
+            elf: elf.to_vec(),
+            memory_config,
+            pk,
+            vk,
+        }
+    }
+
+    pub fn execute(&self, input: &[u8]) -> Result<(PublicValues, u64), JoltError> {
+        let (cycles, _, io) = trace(
+            &self.elf,
+            None,
+            &serialize_input(input)?,
+            &self.memory_config,
+        );
+        if io.panic {
+            return Err(JoltError::ExecutionPanic);
+        }
+        let public_values = deserialize_output(&io.outputs)?;
+        Ok((public_values, cycles.len() as _))
+    }
+
+    pub fn prove(&self, input: &[u8]) -> Result<(PublicValues, JoltProof), JoltError> {
+        let (proof, io, _) = JoltRV64IMAC::prove(&self.pk, &self.elf, &serialize_input(input)?);
+        if io.panic {
+            return Err(JoltError::ExecutionPanic);
+        }
+        let public_values = deserialize_output(&io.outputs)?;
+        let proof = JoltProof {
+            proof,
+            inputs: io.inputs,
+            outputs: io.outputs,
+        };
+        Ok((public_values, proof))
+    }
+
+    pub fn verify(&self, proof: JoltProof) -> Result<PublicValues, JoltError> {
+        JoltRV64IMAC::verify(
+            &self.vk,
+            proof.proof,
+            JoltDevice {
+                inputs: proof.inputs.clone(),
+                outputs: proof.outputs.clone(),
+                panic: false,
+                memory_layout: MemoryLayout::new(&self.memory_config),
+            },
+            None,
+        )?;
+        let public_values = deserialize_output(&proof.outputs)?;
+        Ok(public_values)
+    }
+}
+
+fn serialize_input(bytes: &[u8]) -> Result<Vec<u8>, JoltError> {
+    Ok(postcard::to_stdvec(bytes)
+        .map_err(|err| CommonError::serialize("input", "postcard", err))?)
+}
+
+fn deserialize_output(output: &[u8]) -> Result<Vec<u8>, JoltError> {
+    Ok(if output.is_empty() {
+        Vec::new()
+    } else {
+        postcard::take_from_bytes(output)
+            .map_err(|err| CommonError::deserialize("output", "postcard", err))?
+            .0
+    })
+}

crates/zkvm/jolt/src/compiler.rs

@@ -1,7 +1,7 @@
-mod rust_rv32ima;
-mod rust_rv32ima_customized;
+mod rust_rv64imac;
+mod rust_rv64imac_customized;
 
-pub use rust_rv32ima::RustRv32ima;
-pub use rust_rv32ima_customized::RustRv32imaCustomized;
+pub use rust_rv64imac::RustRv64imac;
+pub use rust_rv64imac_customized::RustRv64imacCustomized;
 
 pub type JoltProgram = Vec<u8>;

crates/zkvm/jolt/src/compiler/rust_rv64imac.rs

@@ -1,10 +1,13 @@
 use crate::{compiler::JoltProgram, error::CompileError};
 use ere_compile_utils::CargoBuildCmd;
 use ere_zkvm_interface::Compiler;
+use jolt_common::constants::{
+    DEFAULT_MEMORY_SIZE, DEFAULT_STACK_SIZE, EMULATOR_MEMORY_CAPACITY, STACK_CANARY_SIZE,
+};
 use std::{env, path::Path};
 
-const TARGET_TRIPLE: &str = "riscv32ima-unknown-none-elf";
-// According to https://github.com/a16z/jolt/blob/55b9830a3944dde55d33a55c42522b81dd49f87a/jolt-core/src/host/mod.rs#L95
+const TARGET_TRIPLE: &str = "riscv64imac-unknown-none-elf";
+// According to https://github.com/a16z/jolt/blob/v0.3.0-alpha/jolt-core/src/host/program.rs#L82
 const RUSTFLAGS: &[&str] = &[
     "-C",
     "passes=lower-atomic",
@@ -14,6 +17,8 @@ const RUSTFLAGS: &[&str] = &[
     "strip=symbols",
     "-C",
     "opt-level=z",
+    "--cfg",
+    "getrandom_backend=\"custom\"",
 ];
 const CARGO_BUILD_OPTIONS: &[&str] = &[
     "--features",
@@ -22,20 +27,20 @@ const CARGO_BUILD_OPTIONS: &[&str] = &[
     "-Zbuild-std=core,alloc",
 ];
 
-const DEFAULT_MEMORY_SIZE: u64 = 10 * 1024 * 1024;
-const DEFAULT_STACK_SIZE: u64 = 4096;
-const LINKER_SCRIPT_TEMPLATE: &str = include_str!("rust_rv32ima/template.ld");
+const LINKER_SCRIPT_TEMPLATE: &str = include_str!("rust_rv64imac/template.ld");
 
 fn make_linker_script() -> String {
     LINKER_SCRIPT_TEMPLATE
+        .replace("{EMULATOR_MEMORY}", &EMULATOR_MEMORY_CAPACITY.to_string())
+        .replace("{STACK_CANARY}", &STACK_CANARY_SIZE.to_string())
         .replace("{MEMORY_SIZE}", &DEFAULT_MEMORY_SIZE.to_string())
         .replace("{STACK_SIZE}", &DEFAULT_STACK_SIZE.to_string())
 }
 
-/// Compiler for Rust guest program to RV32IMA architecture.
-pub struct RustRv32ima;
+/// Compiler for Rust guest program to RV64IMAC architecture.
+pub struct RustRv64imac;
 
-impl Compiler for RustRv32ima {
+impl Compiler for RustRv64imac {
     type Error = CompileError;
 
     type Program = JoltProgram;
@@ -54,21 +59,21 @@ impl Compiler for RustRv32ima {
 
 #[cfg(test)]
 mod tests {
-    use crate::{EreJolt, compiler::RustRv32ima};
+    use crate::{EreJolt, compiler::RustRv64imac};
     use ere_test_utils::host::testing_guest_directory;
     use ere_zkvm_interface::{Compiler, ProverResourceType, zkVM};
 
     #[test]
     fn test_compile() {
         let guest_directory = testing_guest_directory("jolt", "stock_nightly_no_std");
-        let elf = RustRv32ima.compile(&guest_directory).unwrap();
+        let elf = RustRv64imac.compile(&guest_directory).unwrap();
         assert!(!elf.is_empty(), "ELF bytes should not be empty.");
     }
 
     #[test]
     fn test_execute() {
         let guest_directory = testing_guest_directory("jolt", "stock_nightly_no_std");
-        let program = RustRv32ima.compile(&guest_directory).unwrap();
+        let program = RustRv64imac.compile(&guest_directory).unwrap();
         let zkvm = EreJolt::new(program, ProverResourceType::Cpu).unwrap();
 
         zkvm.execute(&[]).unwrap();

crates/zkvm/jolt/src/compiler/rust_rv64imac/template.ld

@@ -1,5 +1,7 @@
+/* Copied from https://github.com/a16z/jolt/blob/v0.3.0-alpha/jolt-core/src/host/mod.rs#L28 */
+
 MEMORY {
-  program (rwx) : ORIGIN = 0x80000000, LENGTH = {MEMORY_SIZE}
+  program (rwx) : ORIGIN = 0x80000000, LENGTH = {EMULATOR_MEMORY}
 }
 
 SECTIONS {
@@ -15,13 +17,16 @@ SECTIONS {
     *(.data)
   } > program
 
-  .bss : {
+  .bss (NOLOAD) : {
     *(.bss)
   } > program
 
   . = ALIGN(8);
+  _STACK_END = .;
+  . = . + {STACK_CANARY};
   . = . + {STACK_SIZE};
   _STACK_PTR = .;
+
   . = ALIGN(8);
   _HEAP_PTR = .;
 }

crates/zkvm/jolt/src/compiler/rust_rv64imac_customized.rs

@@ -1,14 +1,15 @@
 use crate::{compiler::JoltProgram, error::CompileError};
 use ere_compile_utils::{CommonError, cargo_metadata};
 use ere_zkvm_interface::Compiler;
-use jolt::host::DEFAULT_TARGET_DIR;
+use jolt_sdk::host::Program;
 use std::{env::set_current_dir, fs, path::Path};
+use tempfile::tempdir;
 
-/// Compiler for Rust guest program to RV32IMA architecture, using customized
+/// Compiler for Rust guest program to RV64IMAC architecture, using customized
 /// Rust toolchain of Jolt.
-pub struct RustRv32imaCustomized;
+pub struct RustRv64imacCustomized;
 
-impl Compiler for RustRv32imaCustomized {
+impl Compiler for RustRv64imacCustomized {
     type Error = CompileError;
 
     type Program = JoltProgram;
@@ -23,11 +24,13 @@ impl Compiler for RustRv32imaCustomized {
         let metadata = cargo_metadata(guest_directory)?;
         let package_name = &metadata.root_package().unwrap().name;
 
+        let tempdir = tempdir().map_err(CommonError::tempdir)?;
+
         // Note that if this fails, it will panic, hence we need to catch it.
         let elf_path = std::panic::catch_unwind(|| {
-            let mut program = jolt::host::Program::new(package_name);
+            let mut program = Program::new(package_name);
             program.set_std(true);
-            program.build(DEFAULT_TARGET_DIR);
+            program.build(&tempdir.path().to_string_lossy());
             program.elf.unwrap()
         })
         .map_err(|_| CompileError::BuildFailed)?;
@@ -41,23 +44,14 @@ impl Compiler for RustRv32imaCustomized {
 
 #[cfg(test)]
 mod tests {
-    use crate::{EreJolt, compiler::RustRv32imaCustomized};
+    use crate::compiler::RustRv64imacCustomized;
     use ere_test_utils::host::testing_guest_directory;
-    use ere_zkvm_interface::{Compiler, ProverResourceType, zkVM};
+    use ere_zkvm_interface::Compiler;
 
     #[test]
     fn test_compile() {
         let guest_directory = testing_guest_directory("jolt", "basic");
-        let elf = RustRv32imaCustomized.compile(&guest_directory).unwrap();
+        let elf = RustRv64imacCustomized.compile(&guest_directory).unwrap();
         assert!(!elf.is_empty(), "ELF bytes should not be empty.");
     }
-
-    #[test]
-    fn test_execute() {
-        let guest_directory = testing_guest_directory("jolt", "basic");
-        let program = RustRv32imaCustomized.compile(&guest_directory).unwrap();
-        let zkvm = EreJolt::new(program, ProverResourceType::Cpu).unwrap();
-
-        zkvm.execute(&[]).unwrap();
-    }
 }

crates/zkvm/jolt/src/error.rs

@@ -1,5 +1,5 @@
 use ere_compile_utils::CommonError;
-use jolt::jolt_core::utils::errors::ProofVerifyError;
+use jolt_core::utils::errors::ProofVerifyError;
 use std::{io, path::PathBuf};
 use thiserror::Error;
 
@@ -24,6 +24,10 @@ pub enum JoltError {
     #[error(transparent)]
     CommonError(#[from] ere_zkvm_interface::CommonError),
 
+    // Execute
+    #[error("Execution panics")]
+    ExecutionPanic,
+
     // Verify
     #[error("Failed to verify proof: {0}")]
     VerifyProofFailed(#[from] ProofVerifyError),

crates/zkvm/jolt/src/jolt_methods.rs

@@ -1,88 +0,0 @@
-use crate::{EreJoltProof, error::JoltError};
-use common::constants::{DEFAULT_MAX_BYTECODE_SIZE, DEFAULT_MAX_TRACE_LENGTH, DEFAULT_MEMORY_SIZE};
-use jolt::{
-    Jolt, JoltHyperKZGProof, JoltProverPreprocessing, JoltVerifierPreprocessing, MemoryConfig,
-    MemoryLayout, RV32IJoltVM, tracer::JoltDevice,
-};
-
-pub fn preprocess_prover(
-    program: &jolt::host::Program,
-) -> JoltProverPreprocessing<4, jolt::F, jolt::PCS, jolt::ProofTranscript> {
-    let (bytecode, memory_init) = program.decode();
-    let memory_layout = MemoryLayout::new(&MemoryConfig::default());
-    let preprocessing: JoltProverPreprocessing<4, jolt::F, jolt::PCS, jolt::ProofTranscript> =
-        RV32IJoltVM::prover_preprocess(
-            bytecode,
-            memory_layout,
-            memory_init,
-            DEFAULT_MAX_BYTECODE_SIZE as usize,
-            DEFAULT_MEMORY_SIZE as usize,
-            DEFAULT_MAX_TRACE_LENGTH as usize,
-        );
-    preprocessing
-}
-
-pub fn preprocess_verifier(
-    program: &jolt::host::Program,
-) -> JoltVerifierPreprocessing<4, jolt::F, jolt::PCS, jolt::ProofTranscript> {
-    let (bytecode, memory_init) = program.decode();
-    let memory_layout = MemoryLayout::new(&MemoryConfig::default());
-    let preprocessing: JoltVerifierPreprocessing<4, jolt::F, jolt::PCS, jolt::ProofTranscript> =
-        RV32IJoltVM::verifier_preprocess(
-            bytecode,
-            memory_layout,
-            memory_init,
-            DEFAULT_MAX_BYTECODE_SIZE as usize,
-            DEFAULT_MEMORY_SIZE as usize,
-            DEFAULT_MAX_TRACE_LENGTH as usize,
-        );
-    preprocessing
-}
-
-pub fn prove_generic(
-    program: &jolt::host::Program,
-    preprocessing: JoltProverPreprocessing<4, jolt::F, jolt::PCS, jolt::ProofTranscript>,
-    _input: &[u8],
-) -> EreJoltProof {
-    let mut program = program.clone();
-
-    // TODO: Check how to pass private input to jolt, issue for tracking:
-    //       https://github.com/a16z/jolt/issues/371.
-    let input_bytes = Vec::new();
-
-    let (io_device, trace) = program.trace(&input_bytes);
-
-    let (jolt_proof, jolt_commitments, io_device, _) =
-        RV32IJoltVM::prove(io_device, trace, preprocessing);
-
-    EreJoltProof {
-        proof: JoltHyperKZGProof {
-            proof: jolt_proof,
-            commitments: jolt_commitments,
-        },
-        public_outputs: io_device.outputs,
-    }
-}
-
-pub fn verify_generic(
-    proof: EreJoltProof,
-    preprocessing: JoltVerifierPreprocessing<4, jolt::F, jolt::PCS, jolt::ProofTranscript>,
-) -> Result<(), JoltError> {
-    let mut io_device = JoltDevice::new(&MemoryConfig {
-        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,
-    });
-    io_device.outputs = proof.public_outputs;
-
-    RV32IJoltVM::verify(
-        preprocessing,
-        proof.proof.proof,
-        proof.proof.commitments,
-        io_device,
-        None,
-    )?;
-
-    Ok(())
-}

crates/zkvm/jolt/src/lib.rs

@@ -1,36 +1,26 @@
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
 
 use crate::{
+    client::{JoltProof, JoltSdk},
     compiler::JoltProgram,
     error::JoltError,
-    jolt_methods::{preprocess_prover, preprocess_verifier, prove_generic, verify_generic},
 };
 use anyhow::bail;
-use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use ere_zkvm_interface::{
     CommonError, ProgramExecutionReport, ProgramProvingReport, Proof, ProofKind,
     ProverResourceType, PublicValues, zkVM,
 };
-use jolt::{JoltHyperKZGProof, JoltProverPreprocessing, JoltVerifierPreprocessing};
-use std::{env, fs, io::Cursor};
-use tempfile::TempDir;
+use jolt_ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
+use std::{env, io::Cursor, time::Instant};
 
 include!(concat!(env!("OUT_DIR"), "/name_and_sdk_version.rs"));
 
+mod client;
 pub mod compiler;
 pub mod error;
-mod jolt_methods;
-
-#[derive(CanonicalSerialize, CanonicalDeserialize)]
-pub struct EreJoltProof {
-    proof: JoltHyperKZGProof,
-    public_outputs: Vec<u8>,
-}
 
 pub struct EreJolt {
-    elf: JoltProgram,
-    prover_preprocessing: JoltProverPreprocessing<4, jolt::F, jolt::PCS, jolt::ProofTranscript>,
-    verifier_preprocessing: JoltVerifierPreprocessing<4, jolt::F, jolt::PCS, jolt::ProofTranscript>,
+    sdk: JoltSdk,
     _resource: ProverResourceType,
 }
 
@@ -39,31 +29,28 @@ impl EreJolt {
         if !matches!(resource, ProverResourceType::Cpu) {
             panic!("Network or GPU proving not yet implemented for Miden. Use CPU resource type.");
         }
-        let (_tempdir, program) = program(&elf)?;
-        let prover_preprocessing = preprocess_prover(&program);
-        let verifier_preprocessing = preprocess_verifier(&program);
+        let sdk = JoltSdk::new(&elf);
         Ok(EreJolt {
-            elf,
-            prover_preprocessing,
-            verifier_preprocessing,
+            sdk,
             _resource: resource,
         })
     }
 }
 
 impl zkVM for EreJolt {
-    fn execute(&self, _input: &[u8]) -> anyhow::Result<(PublicValues, ProgramExecutionReport)> {
-        let (_tempdir, program) = program(&self.elf)?;
+    fn execute(&self, input: &[u8]) -> anyhow::Result<(PublicValues, ProgramExecutionReport)> {
+        let start = Instant::now();
+        let (public_values, total_num_cycles) = self.sdk.execute(input)?;
+        let execution_duration = start.elapsed();
 
-        // TODO: Check how to pass private input to jolt, issue for tracking:
-        //       https://github.com/a16z/jolt/issues/371.
-        let summary = program.clone().trace_analyze::<jolt::F>(&[]);
-        let trace_len = summary.trace_len();
-
-        // TODO: Public values
-        let public_values = Vec::new();
-
-        Ok((public_values, ProgramExecutionReport::new(trace_len as u64)))
+        Ok((
+            public_values,
+            ProgramExecutionReport {
+                total_num_cycles,
+                execution_duration,
+                ..Default::default()
+            },
+        ))
     }
 
     fn prove(
@@ -77,25 +64,19 @@ impl zkVM for EreJolt {
                 [ProofKind::Compressed]
             ))
         }
-
-        let (_tempdir, program) = program(&self.elf)?;
-
-        let now = std::time::Instant::now();
-        let proof = prove_generic(&program, self.prover_preprocessing.clone(), input);
-        let elapsed = now.elapsed();
+        let start = Instant::now();
+        let (public_values, proof) = self.sdk.prove(input)?;
+        let proving_time = start.elapsed();
 
         let mut proof_bytes = Vec::new();
         proof
             .serialize_compressed(&mut proof_bytes)
             .map_err(|err| CommonError::serialize("proof", "jolt", err))?;
 
-        // TODO: Public values
-        let public_values = Vec::new();
-
         Ok((
             public_values,
             Proof::Compressed(proof_bytes),
-            ProgramProvingReport::new(elapsed),
+            ProgramProvingReport::new(proving_time),
         ))
     }
 
@@ -107,13 +88,10 @@ impl zkVM for EreJolt {
             ))
         };
 
-        let proof = EreJoltProof::deserialize_compressed(&mut Cursor::new(proof))
+        let proof = JoltProof::deserialize_compressed(&mut Cursor::new(proof))
             .map_err(|err| CommonError::deserialize("proof", "jolt", err))?;
 
-        verify_generic(proof, self.verifier_preprocessing.clone())?;
-
-        // TODO: Public values
-        let public_values = Vec::new();
+        let public_values = self.sdk.verify(proof)?;
 
         Ok(public_values)
     }
@@ -127,15 +105,74 @@ impl zkVM for EreJolt {
     }
 }
 
-/// Create `jolt::host::Program` by storing the compiled `elf` to a temporary
-/// file, and set the elf path for `program`, so methods like `decode`, `trace`
-/// and `trace_analyze` that depend on elf path will work.
-pub fn program(elf: &[u8]) -> Result<(TempDir, jolt::host::Program), JoltError> {
-    let tempdir = TempDir::new().map_err(CommonError::tempdir)?;
-    let elf_path = tempdir.path().join("guest.elf");
-    fs::write(&elf_path, elf).map_err(|err| CommonError::write_file("elf", &elf_path, err))?;
-    // Set a dummy package name because we don't need to compile anymore.
-    let mut program = jolt::host::Program::new("");
-    program.elf = Some(elf_path);
-    Ok((tempdir, program))
+#[cfg(test)]
+mod tests {
+    use crate::{
+        EreJolt,
+        compiler::{JoltProgram, RustRv64imacCustomized},
+    };
+    use ere_test_utils::{
+        host::{TestCase, run_zkvm_execute, run_zkvm_prove, testing_guest_directory},
+        program::basic::BasicProgramInput,
+    };
+    use ere_zkvm_interface::{Compiler, ProofKind, ProverResourceType, zkVM};
+    use std::sync::{Mutex, OnceLock};
+
+    /// While proving, Jolt uses global static variables to store some
+    /// parameters, that might cause panics if we prove concurrently, so we put
+    /// a lock here for the test to work without the need to set test threads.
+    static PROVE_LOCK: Mutex<()> = Mutex::new(());
+
+    fn basic_program() -> JoltProgram {
+        static PROGRAM: OnceLock<JoltProgram> = OnceLock::new();
+        PROGRAM
+            .get_or_init(|| {
+                RustRv64imacCustomized
+                    .compile(&testing_guest_directory("jolt", "basic"))
+                    .unwrap()
+            })
+            .clone()
+    }
+
+    #[test]
+    fn test_execute() {
+        let program = basic_program();
+        let zkvm = EreJolt::new(program, ProverResourceType::Cpu).unwrap();
+
+        let test_case = BasicProgramInput::valid();
+        run_zkvm_execute(&zkvm, &test_case);
+    }
+
+    #[test]
+    fn test_execute_invalid_input() {
+        let program = basic_program();
+        let zkvm = EreJolt::new(program, ProverResourceType::Cpu).unwrap();
+
+        for input in [Vec::new(), BasicProgramInput::invalid().serialized_input()] {
+            zkvm.execute(&input).unwrap_err();
+        }
+    }
+
+    #[test]
+    fn test_prove() {
+        let program = basic_program();
+        let zkvm = EreJolt::new(program, ProverResourceType::Cpu).unwrap();
+
+        let _guard = PROVE_LOCK.lock().unwrap();
+
+        let test_case = BasicProgramInput::valid();
+        run_zkvm_prove(&zkvm, &test_case);
+    }
+
+    #[test]
+    fn test_prove_invalid_input() {
+        let program = basic_program();
+        let zkvm = EreJolt::new(program, ProverResourceType::Cpu).unwrap();
+
+        let _guard = PROVE_LOCK.lock().unwrap();
+
+        for input in [Vec::new(), BasicProgramInput::invalid().serialized_input()] {
+            zkvm.prove(&input, ProofKind::default()).unwrap_err();
+        }
+    }
 }

docker/jolt/Dockerfile.base

@@ -5,6 +5,9 @@ FROM $BASE_IMAGE
 # The ere-base image provides Rust, Cargo (with a default nightly), and common tools.
 # We operate as root for SDK installation.
 
+# Set default toolchain to 1.88.0
+RUN rustup default 1.88.0
+
 # Copy the Jolt SDK (CLI) installer script from the workspace context
 COPY --chmod=755 scripts/sdk_installers/install_jolt_sdk.sh /tmp/install_jolt_sdk.sh
 

scripts/sdk_installers/install_jolt_sdk.sh

@@ -28,13 +28,13 @@ ensure_tool_installed "rustup" "to manage Rust toolchains (though Jolt uses defa
 ensure_tool_installed "git" "to install Jolt from a git repository"
 ensure_tool_installed "cargo" "to build and install Rust packages"
 
-JOLT_REVISION="55b9830a3944dde55d33a55c42522b81dd49f87a"
+JOLT_VERSION_TAG="v0.3.0-alpha"
 
 # Install Jolt CLI using cargo install with +nightly
 # This installs the 'jolt' binary directly to $HOME/.cargo/bin
 # The ere-base image should have a compatible default nightly toolchain.
 echo "Installing Jolt CLI from GitHub repository (a16z/jolt)..."
-cargo +nightly install --git https://github.com/a16z/jolt --force --bins jolt --rev "$JOLT_REVISION"
+cargo +nightly install --git https://github.com/a16z/jolt --force --bins jolt --tag "$JOLT_VERSION_TAG"
 
 # Install Jolt's toolchain
 jolt install-toolchain

tests/jolt/basic/Cargo.toml

@@ -1,12 +1,13 @@
 [package]
-name = "guest"
+name = "ere-jolt-guest"
 version = "0.1.0"
 edition = "2021"
 
+[dependencies]
+jolt = { git = "https://github.com/a16z/jolt", package = "jolt-sdk", tag = "v0.3.0-alpha" }
+ere-test-utils = { path = "../../../crates/test-utils" }
+
 [features]
 guest = []
 
-[dependencies]
-jolt = { package = "jolt-sdk", git = "https://github.com/a16z/jolt", rev = "55b9830a3944dde55d33a55c42522b81dd49f87a" }
-
-[workspace]
+[workspace]

tests/jolt/basic/src/lib.rs

@@ -1,15 +0,0 @@
-#![cfg_attr(feature = "guest", no_std)]
-
-#[jolt::provable]
-fn fib(n: u32) -> u128 {
-    let mut a: u128 = 0;
-    let mut b: u128 = 1;
-    let mut sum: u128;
-    for _ in 1..n {
-        sum = a + b;
-        a = b;
-        b = sum;
-    }
-
-    b
-}

tests/jolt/basic/src/main.rs

@@ -1,5 +1,33 @@
-#![cfg_attr(feature = "guest", no_std)]
+#![no_std]
 #![no_main]
 
-#[allow(unused_imports)]
-use guest::*;
+extern crate alloc;
+
+use alloc::vec::Vec;
+use core::mem;
+use ere_test_utils::{
+    guest::Platform,
+    program::{basic::BasicProgram, Program},
+};
+
+static mut INPUT: Vec<u8> = Vec::new();
+static mut OUTPUT: Vec<u8> = Vec::new();
+
+struct JoltPlatform;
+
+impl Platform for JoltPlatform {
+    fn read_input() -> Vec<u8> {
+        unsafe { mem::take(&mut INPUT) }
+    }
+
+    fn write_output(output: &[u8]) {
+        unsafe { mem::replace(&mut OUTPUT, output.to_vec()) };
+    }
+}
+
+#[jolt::provable(guest_only)]
+fn main(input: Vec<u8>) -> Vec<u8> {
+    unsafe { mem::replace(&mut INPUT, input) };
+    BasicProgram::run::<JoltPlatform>();
+    unsafe { mem::take(&mut OUTPUT) }
+}

tests/jolt/stock_nightly_no_std/Cargo.toml

@@ -2,10 +2,10 @@
 name = "addition_no_std"
 edition = "2021"
 
+[dependencies]
+jolt = { git = "https://github.com/a16z/jolt", package = "jolt-sdk", tag = "v0.3.0-alpha" }
+
 [features]
 guest = []
 
-[dependencies]
-jolt-sdk = { git = "https://github.com/a16z/jolt", rev = "55b9830a3944dde55d33a55c42522b81dd49f87a" }
-
 [workspace]

tests/jolt/stock_nightly_no_std/src/main.rs

@@ -1,14 +1,13 @@
-#![cfg_attr(feature = "guest", no_std)]
+#![no_std]
 #![no_main]
 extern crate alloc;
 
 use alloc::vec::Vec;
-use core::sync::atomic::Ordering;
 use core::sync::atomic::AtomicU16;
-use jolt_sdk as jolt;
+use core::sync::atomic::Ordering;
 
-#[jolt::provable]
-fn foo() {
+#[jolt::provable(guest_only)]
+fn main() {
     let a: AtomicU16 = core::hint::black_box(AtomicU16::new(5));
     let b: AtomicU16 = core::hint::black_box(AtomicU16::new(7));