jolt: Guest program compilation with stock rust compiler. (#116)

Cargo.lock

@@ -2281,8 +2281,10 @@ dependencies = [
 name = "ere-jolt"
 version = "0.0.11"
 dependencies = [
+ "anyhow",
  "ark-serialize 0.5.0",
  "build-utils",
+ "cargo_metadata 0.19.2",
  "common",
  "jolt",
  "jolt-core",

crates/ere-jolt/Cargo.toml

@@ -6,6 +6,8 @@ rust-version.workspace = true
 license.workspace = true
 
 [dependencies]
+anyhow.workspace = true
+cargo_metadata.workspace = true
 serde.workspace = true
 tempfile.workspace = true
 thiserror.workspace = true

crates/ere-jolt/src/compile_stock_rust.rs

@@ -0,0 +1,134 @@
+use crate::error::CompileError;
+use cargo_metadata::MetadataCommand;
+use std::fs;
+use std::fs::File;
+use std::io::Write;
+use std::path::{Path, PathBuf};
+use std::process::Command;
+use tempfile::TempDir;
+
+static CARGO_ENCODED_RUSTFLAGS_SEPARATOR: &str = "\x1f";
+const TARGET_TRIPLE: &str = "riscv32im-unknown-none-elf";
+// According to https://github.com/a16z/jolt/blob/55b9830a3944dde55d33a55c42522b81dd49f87a/jolt-core/src/host/mod.rs#L95
+const RUSTFLAGS: &[&str] = &[
+    "-C",
+    "passes=lower-atomic",
+    "-C",
+    "panic=abort",
+    "-C",
+    "strip=symbols",
+    "-C",
+    "opt-level=z",
+];
+const CARGO_ARGS: &[&str] = &[
+    "build",
+    "--release",
+    "--features",
+    "guest",
+    "--target",
+    TARGET_TRIPLE,
+    // For bare metal we have to build core and alloc
+    "-Zbuild-std=core,alloc",
+];
+
+pub fn compile_jolt_program_stock_rust(
+    guest_directory: &Path,
+    toolchain: &String,
+) -> Result<Vec<u8>, CompileError> {
+    compile_program_stock_rust(guest_directory, toolchain)
+}
+
+fn compile_program_stock_rust(
+    guest_directory: &Path,
+    toolchain: &String,
+) -> Result<Vec<u8>, CompileError> {
+    let metadata = MetadataCommand::new().current_dir(guest_directory).exec()?;
+    let package = metadata
+        .root_package()
+        .ok_or_else(|| CompileError::MissingPackageName {
+            path: guest_directory.to_path_buf(),
+        })?;
+
+    let plus_toolchain = format!("+{}", toolchain);
+    let mut cargo_args = [plus_toolchain.as_str()].to_vec();
+    cargo_args.append(&mut CARGO_ARGS.to_vec());
+
+    let mut encoded_rust_flags = RUSTFLAGS.to_vec();
+    let temp_output_dir = TempDir::new_in(guest_directory).unwrap();
+    let linker_script_path = make_linker_script(temp_output_dir.path(), &package.name)?;
+    let linker_path = format!("link-arg=-T{}", linker_script_path.display());
+    encoded_rust_flags.append(&mut ["-C", &linker_path].to_vec());
+
+    let encoded_rust_flags_str = encoded_rust_flags.join(CARGO_ENCODED_RUSTFLAGS_SEPARATOR);
+
+    let target_direcotry = guest_directory
+        .join("target")
+        .join(TARGET_TRIPLE)
+        .join("release");
+
+    // Remove target directory.
+    if target_direcotry.exists() {
+        fs::remove_dir_all(&target_direcotry).unwrap();
+    }
+
+    let result = Command::new("cargo")
+        .current_dir(guest_directory)
+        .env("CARGO_ENCODED_RUSTFLAGS", &encoded_rust_flags_str)
+        .args(cargo_args)
+        .stdout(std::process::Stdio::inherit())
+        .stderr(std::process::Stdio::inherit())
+        .status()
+        .map_err(|source| CompileError::BuildFailure {
+            source: source.into(),
+            crate_path: guest_directory.to_path_buf(),
+        });
+
+    if result.is_err() {
+        return Err(result.err().unwrap());
+    }
+
+    let elf_path = target_direcotry.join(&package.name);
+
+    fs::read(&elf_path).map_err(|e| CompileError::ReadFile {
+        path: elf_path,
+        source: e,
+    })
+}
+
+const DEFAULT_MEMORY_SIZE: u64 = 10 * 1024 * 1024;
+const DEFAULT_STACK_SIZE: u64 = 4096;
+const LINKER_SCRIPT_TEMPLATE: &str = include_str!("template.ld");
+
+fn make_linker_script(
+    temp_output_dir_path: &Path,
+    program_name: &String,
+) -> Result<PathBuf, CompileError> {
+    let linker_path = temp_output_dir_path.join(format!("{}.ld", program_name));
+
+    let linker_script = LINKER_SCRIPT_TEMPLATE
+        .replace("{MEMORY_SIZE}", &DEFAULT_MEMORY_SIZE.to_string())
+        .replace("{STACK_SIZE}", &DEFAULT_STACK_SIZE.to_string());
+
+    let mut file = File::create(&linker_path).expect("could not create linker file");
+    file.write_all(linker_script.as_bytes())
+        .expect("could not save linker");
+
+    Ok(linker_path)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::compile_stock_rust::compile_jolt_program_stock_rust;
+    use test_utils::host::testing_guest_directory;
+
+    #[test]
+    fn test_stock_compiler_impl() {
+        let guest_directory = testing_guest_directory("jolt", "stock_nightly_no_std");
+        let result = compile_jolt_program_stock_rust(&guest_directory, &"nightly".to_string());
+        assert!(result.is_ok(), "Jolt guest program compilation failure.");
+        assert!(
+            !result.unwrap().is_empty(),
+            "ELF bytes should not be empty."
+        );
+    }
+}

crates/ere-jolt/src/error.rs

@@ -34,10 +34,26 @@ pub enum CompileError {
     },
     #[error("Failed to build guest")]
     BuildFailed,
+    #[error("`jolt` build failure for {crate_path} failed: {source}")]
+    BuildFailure {
+        #[source]
+        source: anyhow::Error,
+        crate_path: PathBuf,
+    },
     #[error("Failed to read elf at {path}: {source}")]
     ReadElfFailed { source: io::Error, path: PathBuf },
     #[error("Failed to set current directory to {path}: {source}")]
     SetCurrentDirFailed { source: io::Error, path: PathBuf },
+    #[error("`cargo metadata` failed: {0}")]
+    MetadataCommand(#[from] cargo_metadata::Error),
+    #[error("Could not find `[package].name` in guest Cargo.toml at {path}")]
+    MissingPackageName { path: PathBuf },
+    #[error("Failed to read file at {path}: {source}")]
+    ReadFile {
+        path: PathBuf,
+        #[source]
+        source: std::io::Error,
+    },
 }
 
 #[derive(Debug, Error)]

crates/ere-jolt/src/lib.rs

@@ -5,12 +5,14 @@ use crate::{
     utils::package_name_from_manifest,
 };
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
+use compile_stock_rust::compile_jolt_program_stock_rust;
 use jolt::{JoltHyperKZGProof, JoltProverPreprocessing, JoltVerifierPreprocessing};
 use jolt_core::host::Program;
 use jolt_methods::{preprocess_prover, preprocess_verifier, prove_generic, verify_generic};
 use jolt_sdk::host::DEFAULT_TARGET_DIR;
 use serde::de::DeserializeOwned;
 use std::{
+    env,
     env::set_current_dir,
     fs,
     io::{Cursor, Read},
@@ -23,6 +25,7 @@ use zkvm_interface::{
 };
 
 include!(concat!(env!("OUT_DIR"), "/name_and_sdk_version.rs"));
+mod compile_stock_rust;
 mod error;
 mod jolt_methods;
 mod utils;
@@ -35,33 +38,44 @@ impl Compiler for JOLT_TARGET {
 
     type Program = Vec<u8>;
 
-    fn compile(&self, guest_dir: &Path) -> Result<Self::Program, Self::Error> {
-        // Change current directory for `Program::build` to build guest program.
-        set_current_dir(guest_dir).map_err(|source| CompileError::SetCurrentDirFailed {
-            source,
-            path: guest_dir.to_path_buf(),
-        })?;
-
-        let package_name = package_name_from_manifest(Path::new("Cargo.toml"))?;
-
-        // Note that if this fails, it will panic, hence we need to catch it.
-        let elf_path = std::panic::catch_unwind(|| {
-            let mut program = Program::new(&package_name);
-            program.set_std(true);
-            program.build(DEFAULT_TARGET_DIR);
-            program.elf.unwrap()
-        })
-        .map_err(|_| CompileError::BuildFailed)?;
-
-        let elf = fs::read(&elf_path).map_err(|source| CompileError::ReadElfFailed {
-            source,
-            path: elf_path.to_path_buf(),
-        })?;
-
-        Ok(elf)
+    fn compile(&self, guest_directory: &Path) -> Result<Self::Program, Self::Error> {
+        let toolchain = env::var("ERE_GUEST_TOOLCHAIN").unwrap_or_else(|_error| "jolt".into());
+        match toolchain.as_str() {
+            "jolt" => Ok(compile_jolt_program(guest_directory)?),
+            _ => Ok(compile_jolt_program_stock_rust(
+                guest_directory,
+                &toolchain,
+            )?),
+        }
     }
 }
 
+fn compile_jolt_program(guest_directory: &Path) -> Result<Vec<u8>, JoltError> {
+    // Change current directory for `Program::build` to build guest program.
+    set_current_dir(guest_directory).map_err(|source| CompileError::SetCurrentDirFailed {
+        source,
+        path: guest_directory.to_path_buf(),
+    })?;
+
+    let package_name = package_name_from_manifest(Path::new("Cargo.toml"))?;
+
+    // Note that if this fails, it will panic, hence we need to catch it.
+    let elf_path = std::panic::catch_unwind(|| {
+        let mut program = Program::new(&package_name);
+        program.set_std(true);
+        program.build(DEFAULT_TARGET_DIR);
+        program.elf.unwrap()
+    })
+    .map_err(|_| CompileError::BuildFailed)?;
+
+    let elf = fs::read(&elf_path).map_err(|source| CompileError::ReadElfFailed {
+        source,
+        path: elf_path.to_path_buf(),
+    })?;
+
+    Ok(elf)
+}
+
 #[derive(CanonicalSerialize, CanonicalDeserialize)]
 pub struct EreJoltProof {
     proof: JoltHyperKZGProof,
@@ -175,7 +189,7 @@ pub fn program(elf: &[u8]) -> Result<(TempDir, jolt::host::Program), zkVMError>
 mod tests {
     use super::*;
     use std::sync::OnceLock;
-    use test_utils::host::testing_guest_directory;
+    use test_utils::host::{BasicProgramIo, testing_guest_directory};
 
     static BASIC_PRORGAM: OnceLock<Vec<u8>> = OnceLock::new();
 
@@ -194,4 +208,15 @@ mod tests {
         let elf_bytes = basic_program();
         assert!(!elf_bytes.is_empty(), "ELF bytes should not be empty.");
     }
+
+    #[test]
+    fn test_execute_nightly() {
+        let guest_directory = testing_guest_directory("jolt", "stock_nightly_no_std");
+        let program =
+            compile_jolt_program_stock_rust(&guest_directory, &"nightly".to_string()).unwrap();
+        let zkvm = EreJolt::new(program, ProverResourceType::Cpu).unwrap();
+
+        let result = zkvm.execute(&BasicProgramIo::empty());
+        assert!(result.is_ok(), "Jolt execution failure");
+    }
 }

crates/ere-jolt/src/template.ld

@@ -0,0 +1,27 @@
+MEMORY {
+  program (rwx) : ORIGIN = 0x80000000, LENGTH = {MEMORY_SIZE}
+}
+
+SECTIONS {
+  .text.boot : {
+    *(.text.boot)
+  } > program
+
+  .text : {
+    *(.text)
+  } > program
+
+  .data : {
+    *(.data)
+  } > program
+
+  .bss : {
+    *(.bss)
+  } > program
+
+  . = ALIGN(8);
+  . = . + {STACK_SIZE};
+  _STACK_PTR = .;
+  . = ALIGN(8);
+  _HEAP_PTR = .;
+}

docker/jolt/Dockerfile

@@ -5,6 +5,9 @@ FROM ${BASE_IMAGE_TAG}
 # The ere-base image provides Rust, Cargo (with a default nightly), and common tools.
 # We operate as root for SDK installation.
 
+# Add `rust-src` component to enable std build for nightly rust.
+RUN rustup +nightly component add rust-src
+
 # Copy the Jolt SDK (CLI) installer script from the workspace context
 COPY scripts/sdk_installers/install_jolt_sdk.sh /tmp/install_jolt_sdk.sh
 RUN chmod +x /tmp/install_jolt_sdk.sh

tests/jolt/stock_nightly_no_std/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "addition_no_std"
+edition = "2021"
+
+[features]
+guest = []
+
+[dependencies]
+jolt-sdk = { git = "https://github.com/a16z/jolt", rev = "55b9830a3944dde55d33a55c42522b81dd49f87a" }
+
+[workspace]

tests/jolt/stock_nightly_no_std/src/main.rs

@@ -0,0 +1,26 @@
+#![cfg_attr(feature = "guest", 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;
+
+#[jolt::provable]
+fn foo() {
+    let a: AtomicU16 = core::hint::black_box(AtomicU16::new(5));
+    let b: AtomicU16 = core::hint::black_box(AtomicU16::new(7));
+
+    if a.load(Ordering::SeqCst) + b.load(Ordering::SeqCst) != 12 {
+        panic!("Something went wrong!");
+    }
+
+    let mut v: Vec<AtomicU16> = Vec::new();
+    v.push(AtomicU16::new(5));
+    v.push(AtomicU16::new(7));
+
+    if v[0].load(Ordering::SeqCst) + v[1].load(Ordering::SeqCst) != 12 {
+        panic!("Something went wrong!");
+    }
+}