diff --git a/Makefile b/Makefile index 4e64af1..8ec7430 100644 --- a/Makefile +++ b/Makefile @@ -36,6 +36,10 @@ install_risc0: @rzup install @cargo risczero --version +install_jolt: + @echo "Installing Jolt from source" + @cargo install --git https://github.com/a16z/jolt jolt-cli + all: install __EXAMPLES__: @@ -105,6 +109,16 @@ prove_sp1_iseven: prove_sp1_bubble_sort: cargo run --release -p prooflab -- prove-sp1 examples/bubble_sort +# JOLT +prove_jolt_fibonacci: + cargo run --release -p prooflab -- prove-jolt examples/fibonacci + +prove_jolt_is_even: + cargo run --release -p prooflab -- prove-jolt examples/is_even + +prove_jolt_bubble_sort: + cargo run --release -p prooflab -- prove-jolt examples/bubble_sort + # Benchmark Commands benchmark_sp1_fibonacci: cargo run --release -p prooflab -- prove-sp1 examples/fibonacci --enable-telemetry diff --git a/README.md b/README.md index 65ae4bb..68e76f7 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # prooflab-rs -prooflab-rs is a CLI tool that simplifies developing zero-knowledge applications in Rust using zkVMs like SP1 or Risc0. +prooflab-rs is a CLI tool that simplifies developing zero-knowledge applications in Rust using zkVMs like SP1, Risc0, or Jolt. ## Project Structure @@ -199,6 +199,11 @@ To generate a proof of your code's execution, run one of the following commands: cargo run --release -- prove-risc0 ``` +- **Using Jolt**: + ```sh + cargo run --release -- prove-jolt + ``` + > **Note:** Aligned currently supports verification of [Risc0](https://dev.risczero.com/api/zkvm/quickstart#1-install-the-risc-zero-toolchain) proofs from release version `v1.0.1`. ### Submitting Proofs to Aligned @@ -257,16 +262,18 @@ For help with prooflab-rs or questions about implementation, please join our [Te After installing prooflab-rs, run any of the following example commands. You can choose either Risc0 or SP1 as your zkVM backend: -| Example | Risc0 | SP1 | -|---------|-------|-----| -| **Fibonacci** | `make prove_risc0_fibonacci` | `make prove_sp1_fibonacci` | -| **RSA** | `make prove_risc0_rsa` | `make prove_sp1_rsa` | -| **ECDSA** | `make prove_risc0_ecdsa` | `make prove_sp1_ecdsa` | -| **Blockchain State Diff** | `make prove_risc0_json` | `make prove_sp1_json` | -| **Regex** | `make prove_risc0_regex` | `make prove_sp1_regex` | -| **SHA256** | `make prove_risc0_sha` | `make prove_sp1_sha` | -| **Tendermint** | `make prove_risc0_tendermint` | `make prove_sp1_tendermint` | -| **ZK Quiz** | `make prove_risc0_zkquiz` | `make prove_sp1_zkquiz` | +| Example | Risc0 | SP1 | Jolt | +|---------|-------|-----|------| +| **Fibonacci** | `make prove_risc0_fibonacci` | `make prove_sp1_fibonacci` | `make prove_jolt_fibonacci` | +| **RSA** | `make prove_risc0_rsa` | `make prove_sp1_rsa` | - | +| **ECDSA** | `make prove_risc0_ecdsa` | `make prove_sp1_ecdsa` | - | +| **Blockchain State Diff** | `make prove_risc0_json` | `make prove_sp1_json` | - | +| **Regex** | `make prove_risc0_regex` | `make prove_sp1_regex` | - | +| **SHA256** | `make prove_risc0_sha` | `make prove_sp1_sha` | - | +| **Tendermint** | `make prove_risc0_tendermint` | `make prove_sp1_tendermint` | - | +| **ZK Quiz** | `make prove_risc0_zkquiz` | `make prove_sp1_zkquiz` | - | +| **Is Even** | - | `make prove_sp1_iseven` | `make prove_jolt_is_even` | +| **Bubble Sort** | `make prove_risc0_bubble_sort` | `make prove_sp1_bubble_sort` | `make prove_jolt_bubble_sort` | ## Acknowledgments @@ -278,6 +285,7 @@ We extend our deepest gratitude to: - The [RISC0](https://github.com/risc0/risc0) team for their pioneering work on RISC-V based zero-knowledge virtual machines - The [SP1](https://github.com/succinctlabs/sp1) team for their innovative approach to high-performance zero-knowledge computation +- The [Jolt](https://github.com/a16z/jolt) team for their novel Rust-native zero-knowledge VM - The entire zero-knowledge community for pushing the boundaries of what's possible with this transformative technology ProofLab is committed to advancing the state of zero-knowledge application development by providing open tools, standardized metrics, and a collaborative platform for the broader ecosystem. diff --git a/crates/prooflab/src/jolt.rs b/crates/prooflab/src/jolt.rs new file mode 100644 index 0000000..b638fa0 --- /dev/null +++ b/crates/prooflab/src/jolt.rs @@ -0,0 +1,116 @@ +use serde::{Deserialize, Serialize}; +use std::{ + fs, + io::{self, Write}, + path::PathBuf, + process::{Command, ExitStatus}, + time::Duration, +}; + +use crate::utils; + +#[derive(Default, Serialize, Deserialize)] +pub struct JoltMetrics { + pub cycles: u64, + pub num_segments: usize, + pub core_proof_size: usize, + pub recursive_proof_size: usize, + pub core_prove_duration: Duration, + pub core_verify_duration: Duration, + pub compress_prove_duration: Duration, + pub compress_verify_duration: Duration, +} + +/// Jolt workspace directories +pub const JOLT_SCRIPT_DIR: &str = "workspaces/jolt/script"; +pub const JOLT_SRC_DIR: &str = "workspaces/jolt/program"; +pub const JOLT_GUEST_MAIN: &str = "workspaces/jolt/program/src/lib.rs"; +pub const JOLT_HOST_MAIN: &str = "workspaces/jolt/script/src/main.rs"; +pub const JOLT_BASE_GUEST_CARGO_TOML: &str = "workspaces/base_files/jolt/cargo_guest"; +pub const JOLT_BASE_HOST_CARGO_TOML: &str = "workspaces/base_files/jolt/cargo_host"; +pub const JOLT_BASE_HOST: &str = "workspaces/base_files/jolt/host"; +pub const JOLT_BASE_HOST_FILE: &str = "workspaces/base_files/jolt/host"; +pub const JOLT_GUEST_CARGO_TOML: &str = "workspaces/jolt/program/Cargo.toml"; + +// Proof data generation paths +pub const JOLT_ELF_PATH: &str = "./proof_data/jolt/jolt.elf"; +pub const JOLT_PROOF_PATH: &str = "./proof_data/jolt/jolt.proof"; +pub const JOLT_PUB_INPUT_PATH: &str = "./proof_data/jolt/jolt.pub"; +pub const JOLT_METRICS_PATH: &str = "./proof_data/jolt/jolt_metrics.json"; + +/// Jolt header added to programs for generating proofs of their execution +pub const JOLT_GUEST_PROGRAM_HEADER: &str = "#![cfg_attr(feature = \"guest\", no_std)]\n#![no_main]\n"; + +/// Jolt Cargo patch for optimization (if needed in the future) +pub const JOLT_ACCELERATION_IMPORT: &str = "\n[patch.crates-io]\n# Jolt-specific optimized crates would go here\n"; + +/// Jolt User I/O +// Host +pub const JOLT_HOST_WRITE: &str = "input"; +pub const JOLT_HOST_READ: &str = "output"; + +// Guest +pub const JOLT_IO_READ: &str = "jolt::io::read();"; +pub const JOLT_IO_COMMIT: &str = "jolt::io::commit"; + +pub fn prepare_host( + input: &str, + output: &str, + imports: &str, + host_dir: &PathBuf, + host_main: &PathBuf, +) -> io::Result<()> { + let mut host_program = imports.to_string(); + let contents = fs::read_to_string(host_dir)?; + + host_program.push_str(&contents); + + // Insert input body + let host_program = host_program.replace(utils::HOST_INPUT, input); + // Insert output body + let host_program = host_program.replace(utils::HOST_OUTPUT, output); + + // replace prooflab_io::write + let host_program = host_program.replace(utils::IO_WRITE, JOLT_HOST_WRITE); + // replace prooflab_io::out() + let host_program = host_program.replace(utils::IO_OUT, JOLT_HOST_READ); + + // Write to host + let mut file = fs::File::create(host_main)?; + file.write_all(host_program.as_bytes())?; + Ok(()) +} + +/// Build the Jolt program +pub fn build_jolt_program(script_dir: &PathBuf) -> io::Result { + Command::new("cargo") + .arg("build") + .arg("--release") + .current_dir(script_dir) + .status() +} + +/// Generates Jolt proof +pub fn generate_jolt_proof( + script_dir: &PathBuf, + current_dir: &PathBuf, + use_gpu: bool, +) -> io::Result { + let mut cmd = Command::new("cargo"); + cmd.arg("run").arg("--release"); + + if use_gpu { + cmd.arg("--features").arg("gpu"); + cmd.env("JOLT_GPU", "1"); + } + + cmd.arg("--") + .arg(current_dir) + .current_dir(script_dir) + .status() +} + +pub fn read_metrics() -> io::Result { + let metrics_str = fs::read_to_string(JOLT_METRICS_PATH)?; + serde_json::from_str(&metrics_str).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)) +} \ No newline at end of file diff --git a/crates/prooflab/src/lib.rs b/crates/prooflab/src/lib.rs index f941e43..bf33bfb 100644 --- a/crates/prooflab/src/lib.rs +++ b/crates/prooflab/src/lib.rs @@ -21,9 +21,17 @@ use ethers::signers::LocalWallet; pub mod risc0; pub mod sp1; +pub mod jolt; pub mod telemetry; pub mod utils; +// Add Jolt to ProvingSystemId enum +// Note: This is a temporary solution. In a proper implementation, +// this would be added to the aligned_sdk crate. +impl ProvingSystemId { + pub const JOLT: Self = Self(2); // Assuming RISC0=0, SP1=1 +} + // Make proof_data path optional // Make keystore unneeded #[derive(Args, Debug)] diff --git a/crates/prooflab/src/main.rs b/crates/prooflab/src/main.rs index 174b7fa..9d31786 100644 --- a/crates/prooflab/src/main.rs +++ b/crates/prooflab/src/main.rs @@ -27,6 +27,8 @@ enum Commands { ProveSp1(ProofArgs), #[clap(about = "Generate a proof of execution of a program using RISC0")] ProveRisc0(ProofArgs), + #[clap(about = "Generate a proof of execution of a program using Jolt")] + ProveJolt(ProofArgs), } #[tokio::main] @@ -35,6 +37,261 @@ async fn main() -> io::Result<()> { let cli = Cli::parse(); match &cli.command { + Commands::ProveJolt(args) => { + info!("Proving with Jolt, program in: {}", args.guest_path); + + let telemetry = TelemetryCollector::new( + "JOLT", + args.precompiles, + args.gpu, + args.enable_telemetry, + &args.guest_path, + ); + let workspace_start = Instant::now(); + + // Perform sanitation checks on directory + let proof_data_dir = PathBuf::from(&args.proof_data_directory_path); + if !proof_data_dir.exists() { + info!("Saving Proofs to: {:?}", &args.proof_data_directory_path); + std::fs::create_dir_all(proof_data_dir)?; + } + if utils::validate_directory_structure(&args.guest_path) { + let Some(home_dir) = dirs::home_dir() else { + error!("Failed to locate home directory"); + return Ok(()); + }; + let Ok(current_dir) = std::env::current_dir() else { + error!("Failed to locate current directory"); + return Ok(()); + }; + let home_dir = home_dir.join(".prooflab"); + utils::prepare_workspace( + &PathBuf::from(&args.guest_path), + &home_dir.join(jolt::JOLT_SRC_DIR), + &home_dir.join(jolt::JOLT_GUEST_CARGO_TOML), + &home_dir.join("workspaces/jolt/script"), + &home_dir.join("workspaces/jolt/script/Cargo.toml"), + &home_dir.join(jolt::JOLT_BASE_HOST_CARGO_TOML), + &home_dir.join(jolt::JOLT_BASE_GUEST_CARGO_TOML), + )?; + + telemetry.record_workspace_setup(workspace_start.elapsed()); + + let compilation_start = Instant::now(); + let Ok(imports) = utils::get_imports(&home_dir.join(jolt::JOLT_GUEST_MAIN)) else { + error!("Failed to extract imports"); + return Ok(()); + }; + + let main_path = home_dir.join(jolt::JOLT_GUEST_MAIN); + let Ok(function_bodies) = utils::extract_function_bodies( + &main_path, + vec![ + "fn main()".to_string(), + "fn input()".to_string(), + "fn output()".to_string(), + ], + ) else { + error!("Failed to extract function bodies"); + return Ok(()); + }; + + utils::prepare_guest( + &imports, + &function_bodies[0], + jolt::JOLT_GUEST_PROGRAM_HEADER, + jolt::JOLT_IO_READ, + jolt::JOLT_IO_COMMIT, + &home_dir.join(jolt::JOLT_GUEST_MAIN), + )?; + jolt::prepare_host( + &function_bodies[1], + &function_bodies[2], + &imports, + &home_dir.join(jolt::JOLT_BASE_HOST), + &home_dir.join(jolt::JOLT_HOST_MAIN), + )?; + + if args.precompiles { + let mut toml_file = OpenOptions::new() + .append(true) + .open(home_dir.join(jolt::JOLT_GUEST_CARGO_TOML))?; + + writeln!(toml_file, "{}", jolt::JOLT_ACCELERATION_IMPORT)?; + } + + let script_dir = home_dir.join(jolt::JOLT_SCRIPT_DIR); + + // Build the program first + let build_result = jolt::build_jolt_program(&script_dir)?; + if !build_result.success() { + error!("Jolt program build failed"); + return Ok(()); + } + info!("Jolt program built successfully"); + telemetry.record_compilation(compilation_start.elapsed()); + + // Record compiled program size + if let Ok(metadata) = fs::metadata(home_dir.join( + "workspaces/jolt/program/target/release/guest", + )) { + telemetry.record_program_size(metadata.len()); + info!("Recorded Jolt program size: {} bytes", metadata.len()); + } else { + error!("Failed to read Jolt program size"); + } + + let proof_gen_start = Instant::now(); + + // Start resource sampling in a separate thread + let tx = telemetry.start_resource_monitoring(); + + let result = jolt::generate_jolt_proof(&script_dir, ¤t_dir, args.gpu)?; + + // Stop resource sampling + let _ = tx.send(()); + + telemetry.record_proof_generation(proof_gen_start.elapsed()); + + if result.success() { + info!("Jolt proof generated"); + + // Read and record Jolt metrics + if let Ok(jolt_metrics) = jolt::read_metrics() { + telemetry.record_zk_metrics( + Some(jolt_metrics.cycles), + Some(jolt_metrics.num_segments), + Some(jolt_metrics.core_proof_size), + Some(jolt_metrics.recursive_proof_size), + ); + telemetry.record_proof_timings( + jolt_metrics.core_prove_duration, + jolt_metrics.core_verify_duration, + Some(jolt_metrics.compress_prove_duration), + Some(jolt_metrics.compress_verify_duration), + ); + } + + utils::replace( + &home_dir.join(jolt::JOLT_GUEST_CARGO_TOML), + jolt::JOLT_ACCELERATION_IMPORT, + "", + )?; + + // Submit to aligned + if args.submit_to_aligned { + submit_proof_to_aligned( + jolt::JOLT_PROOF_PATH, + jolt::JOLT_ELF_PATH, + Some(jolt::JOLT_PUB_INPUT_PATH), + args, + ProvingSystemId::JOLT, + ) + .await + .map_err(|e| { + error!("Proof not submitted to Aligned"); + io::Error::other(e.to_string()) + })?; + info!("Jolt proof submitted and verified on Aligned"); + } + + // Save telemetry data if enabled + if let Some(telemetry_data) = telemetry.finalize() { + if args.enable_telemetry { + fs::create_dir_all(&args.telemetry_output_path)?; + let timestamp = chrono::Local::now().format("%Y%m%d_%H%M%S"); + let package_name = telemetry_data + .program + .guest_metadata + .package_name + .as_deref() + .unwrap_or("unknown"); + let instance_type = telemetry_data + .system_info + .ec2_instance_type + .as_deref() + .unwrap_or("local"); + let telemetry_file = format!( + "{}/jolt_telemetry_{}_{}_{}_{}.json", + args.telemetry_output_path, + package_name, + instance_type, + timestamp, + if result.success() { + "success" + } else { + "failed" + } + ); + fs::write( + &telemetry_file, + serde_json::to_string_pretty(&telemetry_data)?, + )?; + info!("Telemetry data saved to: {}", telemetry_file); + } + } + + std::fs::copy( + home_dir.join(jolt::JOLT_BASE_HOST_FILE), + home_dir.join(jolt::JOLT_HOST_MAIN), + ) + .inspect_err(|_e| { + error!("Failed to clear Jolt host file"); + })?; + return Ok(()); + } + error!( + "Jolt proof generation failed with exit code: {}", + result.code().unwrap_or(-1) + ); + + // Save telemetry data even on failure + if let Some(telemetry_data) = telemetry.finalize() { + if args.enable_telemetry { + fs::create_dir_all(&args.telemetry_output_path)?; + let timestamp = chrono::Local::now().format("%Y%m%d_%H%M%S"); + let package_name = telemetry_data + .program + .guest_metadata + .package_name + .as_deref() + .unwrap_or("unknown"); + let instance_type = telemetry_data + .system_info + .ec2_instance_type + .as_deref() + .unwrap_or("local"); + let telemetry_file = format!( + "{}/jolt_telemetry_{}_{}_{}_{}.json", + args.telemetry_output_path, + package_name, + instance_type, + timestamp, + if result.success() { + "success" + } else { + "failed" + } + ); + fs::write( + &telemetry_file, + serde_json::to_string_pretty(&telemetry_data)?, + )?; + info!("Telemetry data saved to: {}", telemetry_file); + } + } + + // Clear host + std::fs::copy( + home_dir.join(jolt::JOLT_BASE_HOST_FILE), + home_dir.join(jolt::JOLT_HOST_MAIN), + )?; + return Ok(()); + } else { + error!("Invalid directory structure. Please ensure your program directory meets the minimum requirements."); + return Ok(()); + } + }, Commands::ProveSp1(args) => { info!("Proving with SP1, program in: {}", args.guest_path); diff --git a/workspaces/base_files/jolt/cargo_guest b/workspaces/base_files/jolt/cargo_guest new file mode 100644 index 0000000..f65c75d --- /dev/null +++ b/workspaces/base_files/jolt/cargo_guest @@ -0,0 +1,10 @@ +[package] +version = "0.1.0" +name = "guest" +edition = "2021" + +[features] +guest = [] + +[dependencies] +jolt = { package = "jolt-sdk", git = "https://github.com/a16z/jolt" } \ No newline at end of file diff --git a/workspaces/base_files/jolt/cargo_host b/workspaces/base_files/jolt/cargo_host new file mode 100644 index 0000000..8d10feb --- /dev/null +++ b/workspaces/base_files/jolt/cargo_host @@ -0,0 +1,18 @@ +[package] +version = "0.1.0" +name = "host" +edition = "2021" + +[features] +default = [] +gpu = ["jolt-sdk/gpu"] + +[dependencies] +guest = { path = "../program" } +jolt-sdk = { git = "https://github.com/a16z/jolt", features = ["host"] } + +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +bincode = "1.3" +tracing = "0.1" +tracing-subscriber = "0.3" \ No newline at end of file diff --git a/workspaces/base_files/jolt/host b/workspaces/base_files/jolt/host new file mode 100644 index 0000000..9f65401 --- /dev/null +++ b/workspaces/base_files/jolt/host @@ -0,0 +1,71 @@ +use std::time::Instant; +use std::fs; +use std::path::Path; +use serde_json; +use guest::*; +mod metrics; +use metrics::{MetricsCollector, JoltMetrics}; + +fn main() { + let args: Vec = std::env::args().collect(); + let current_dir = std::path::PathBuf::from(args[1].clone()); + + // Create metrics collection + let mut metrics = JoltMetrics::default(); + let mut core_timer = MetricsCollector::new(); + let mut compress_timer = MetricsCollector::new(); + + // Get the build function from guest + let (prove_main, verify_main) = build_main(); + + // Get program analysis if available + if let Ok(program_summary) = analyze_main() { + // Estimate cycles + metrics.cycles = program_summary.cycle_estimate.unwrap_or(0); + metrics.num_segments = 1; // Jolt doesn't have segments concept + } + + // INPUT // + + // Generate proof + core_timer.start_timing(); + let (output, proof) = prove_main(); + metrics.core_prove_duration = core_timer.elapsed().unwrap(); + + // Get proof size + let serialized_proof = bincode::serialize(&proof).unwrap(); + metrics.core_proof_size = serialized_proof.len(); + + // Verify proof + core_timer.start_timing(); + let is_valid = verify_main(proof.clone()); + metrics.core_verify_duration = core_timer.elapsed().unwrap(); + + if !is_valid { + panic!("Proof verification failed"); + } + + // OUTPUT // + + // Save proof artifacts + fs::create_dir_all(current_dir.join("proof_data/jolt")) + .expect("Failed to create proof_data/jolt"); + + fs::write( + current_dir.join("proof_data/jolt/jolt.proof"), + &serialized_proof, + ) + .expect("Failed to save Jolt Proof file"); + + // Save output + let serialized_output = bincode::serialize(&output).unwrap(); + fs::write( + current_dir.join("proof_data/jolt/jolt.pub"), + &serialized_output, + ) + .expect("Failed to save Jolt public output"); + + // Save metrics + metrics::write_metrics(&metrics, current_dir.join("proof_data/jolt").as_path()) + .expect("Failed to save metrics"); +} \ No newline at end of file diff --git a/workspaces/base_files/jolt/metrics.rs b/workspaces/base_files/jolt/metrics.rs new file mode 100644 index 0000000..743cdca --- /dev/null +++ b/workspaces/base_files/jolt/metrics.rs @@ -0,0 +1,38 @@ +use serde::{Deserialize, Serialize}; +use std::path::Path; +use std::time::{Duration, Instant}; + +#[derive(Default, Serialize, Deserialize)] +pub struct JoltMetrics { + pub cycles: u64, + pub num_segments: usize, + pub core_proof_size: usize, + pub recursive_proof_size: usize, + pub core_prove_duration: Duration, + pub core_verify_duration: Duration, + pub compress_prove_duration: Duration, + pub compress_verify_duration: Duration, +} + +pub struct MetricsCollector { + start_time: Option, +} + +impl MetricsCollector { + pub fn new() -> Self { + Self { start_time: None } + } + + pub fn start_timing(&mut self) { + self.start_time = Some(Instant::now()); + } + + pub fn elapsed(&self) -> Option { + self.start_time.map(|start| start.elapsed()) + } +} + +pub fn write_metrics(metrics: &JoltMetrics, path: &Path) -> std::io::Result<()> { + let json = serde_json::to_string_pretty(metrics)?; + std::fs::write(path.join("jolt_metrics.json"), json) +} \ No newline at end of file