Merge pull request #17 from eth-applied-research-group/kw/add-more-func-risc0

feat: Impl zkVM for Risc0

crates/ere-risczero/Cargo.toml

@@ -7,8 +7,11 @@ license.workspace = true
 
 [dependencies]
 zkvm-interface = { workspace = true }
-anyhow = "1.0"                        #TODO: remove only needed in tests
+anyhow = "1.0"                                               #TODO: remove only needed in tests
 toml = "0.8"
+risc0-zkvm = { version = "^2.0.2", features = ["unstable"] }
+borsh = "1.5.7"
+hex = "*"
 
 tempfile = "3.3"
 serde_json = "1.0"

crates/ere-risczero/src/compile.rs

@@ -1,5 +1,6 @@
 mod file_utils;
 use file_utils::FileRestorer;
+use risc0_zkvm::Digest;
 
 use crate::error::CompileError;
 use serde_json::Value as JsonValue;
@@ -9,13 +10,21 @@ use std::{
     process::Command,
 };
 
+#[derive(Debug, Clone)]
+pub struct Risc0Program {
+    // TODO: Seems like the risc0 compilation is also compiling
+    // TODO: the analogous prover and verifying key
+    pub(crate) elf: Vec<u8>,
+    pub(crate) image_id: Digest,
+}
+
 /// BUILD_SCRIPT_TEMPLATE that we will use to fetch the elf-path
 /// TODO: We might be able to deterministically get the elf path
 /// TODO: But note we also probably want the image id too, so not sure
 /// TODO: we can remove this hack sometime soon.
 const BUILD_SCRIPT_TEMPLATE: &str = include_str!("../build_script_template.rs");
 
-pub(crate) fn compile_risczero_program(path: &Path) -> Result<Vec<u8>, CompileError> {
+pub(crate) fn compile_risczero_program(path: &Path) -> Result<Risc0Program, CompileError> {
     if !path.exists() || !path.is_dir() {
         return Err(CompileError::InvalidMethodsPath(path.to_path_buf()));
     }
@@ -56,9 +65,14 @@ pub(crate) fn compile_risczero_program(path: &Path) -> Result<Vec<u8>, CompileEr
             field: "elf_path",
             file: info_file.clone(),
         })?;
+    let image_id_hex_str = info_json["image_id_hex"].as_str().unwrap();
+    let image_id = hex::decode(image_id_hex_str).unwrap();
+    let image_id = image_id.try_into().unwrap();
 
-    // Return bytes
-    fs::read(&elf_path).map_err(|e| CompileError::io(e, "reading ELF file"))
+    // Return Program
+    fs::read(&elf_path)
+        .map_err(|e| CompileError::io(e, "reading ELF file"))
+        .map(|elf| Risc0Program { elf, image_id })
 }
 
 #[cfg(test)]
@@ -85,7 +99,10 @@ mod tests {
 
             let program =
                 compile_risczero_program(&test_methods_path).expect("risc0 compilation failed");
-            assert!(!program.is_empty(), "Risc0 ELF bytes should not be empty.");
+            assert!(
+                !program.elf.is_empty(),
+                "Risc0 ELF bytes should not be empty."
+            );
         }
     }
 }

crates/ere-risczero/src/lib.rs

@@ -1,7 +1,12 @@
 use compile::compile_risczero_program;
-use zkvm_interface::Compiler;
+use risc0_zkvm::{ExecutorEnv, Receipt, default_executor, default_prover};
+use zkvm_interface::{
+    Compiler, Input, ProgramExecutionReport, ProgramProvingReport, ProverResourceType, zkVM,
+    zkVMError,
+};
 
 mod compile;
+pub use compile::Risc0Program;
 
 mod error;
 use error::RiscZeroError;
@@ -12,9 +17,193 @@ pub struct RV32_IM_RISCZERO_ZKVM_ELF;
 impl Compiler for RV32_IM_RISCZERO_ZKVM_ELF {
     type Error = RiscZeroError;
 
-    type Program = Vec<u8>;
+    type Program = Risc0Program;
 
     fn compile(path_to_program: &std::path::Path) -> Result<Self::Program, Self::Error> {
         compile_risczero_program(path_to_program).map_err(RiscZeroError::from)
     }
 }
+
+impl EreRisc0 {
+    pub fn new(
+        program: <RV32_IM_RISCZERO_ZKVM_ELF as Compiler>::Program,
+        resource_type: ProverResourceType,
+    ) -> Self {
+        Self {
+            program,
+            resource_type,
+        }
+    }
+}
+
+pub struct EreRisc0 {
+    program: <RV32_IM_RISCZERO_ZKVM_ELF as Compiler>::Program,
+    resource_type: ProverResourceType,
+}
+
+impl zkVM for EreRisc0 {
+    fn execute(&self, inputs: &Input) -> Result<ProgramExecutionReport, zkVMError> {
+        let executor = default_executor();
+        let env = ExecutorEnv::builder()
+            .write_slice(inputs.bytes())
+            .build()
+            .map_err(|err| zkVMError::Other(err.into()))?;
+
+        let session_info = executor
+            .execute(env, &self.program.elf)
+            .map_err(|err| zkVMError::Other(err.into()))?;
+        Ok(ProgramExecutionReport {
+            total_num_cycles: session_info.cycles() as u64,
+            region_cycles: Default::default(),
+        })
+    }
+
+    fn prove(&self, inputs: &Input) -> Result<(Vec<u8>, ProgramProvingReport), zkVMError> {
+        let prover = default_prover();
+        let env = ExecutorEnv::builder()
+            .write_slice(inputs.bytes())
+            .build()
+            .map_err(|err| zkVMError::Other(err.into()))?;
+
+        let now = std::time::Instant::now();
+        let prove_info = prover
+            .prove(env, &self.program.elf)
+            .map_err(|err| zkVMError::Other(err.into()))?;
+        let proving_time = now.elapsed();
+
+        let encoded =
+            borsh::to_vec(&prove_info.receipt).map_err(|err| zkVMError::Other(Box::new(err)))?;
+        Ok((encoded, ProgramProvingReport::new(proving_time)))
+    }
+
+    fn verify(&self, proof: &[u8]) -> Result<(), zkVMError> {
+        let decoded: Receipt =
+            borsh::from_slice(&proof).map_err(|err| zkVMError::Other(Box::new(err)))?;
+
+        decoded
+            .verify(self.program.image_id)
+            .map_err(|err| zkVMError::Other(Box::new(err)))
+    }
+}
+
+#[cfg(test)]
+mod prove_tests {
+    use std::path::PathBuf;
+
+    use super::*;
+    use zkvm_interface::Input;
+
+    fn get_prove_test_guest_program_path() -> PathBuf {
+        let workspace_dir = env!("CARGO_WORKSPACE_DIR");
+        PathBuf::from(workspace_dir)
+            .join("tests")
+            .join("risczero")
+            .join("compile")
+            .join("project_structure_build")
+            .canonicalize()
+            .expect("Failed to find or canonicalize test Risc0 methods crate")
+    }
+
+    fn get_compiled_test_r0_elf_for_prove() -> Result<Risc0Program, RiscZeroError> {
+        let test_guest_path = get_prove_test_guest_program_path();
+        RV32_IM_RISCZERO_ZKVM_ELF::compile(&test_guest_path)
+    }
+
+    #[test]
+    fn test_prove_r0_dummy_input() {
+        let program = get_compiled_test_r0_elf_for_prove().unwrap();
+
+        let mut input_builder = Input::new();
+        let n: u32 = 42;
+        let a: u16 = 42;
+        input_builder.write(&n).unwrap();
+        input_builder.write(&a).unwrap();
+
+        let zkvm = EreRisc0::new(program, ProverResourceType::Cpu);
+
+        let proof_bytes = match zkvm.prove(&input_builder) {
+            Ok((prove_result, _)) => prove_result,
+            Err(err) => {
+                panic!("Proving error in test: {:?}", err);
+            }
+        };
+
+        assert!(!proof_bytes.is_empty(), "Proof bytes should not be empty.");
+
+        let verify_results = zkvm.verify(&proof_bytes).is_ok();
+        assert!(verify_results);
+
+        // TODO: Check public inputs
+    }
+
+    #[test]
+    // TODO: Note: SP1 will panic here
+    // #[should_panic]
+    fn test_prove_r0_fails_on_bad_input_causing_execution_failure() {
+        let elf_bytes = get_compiled_test_r0_elf_for_prove().unwrap();
+
+        let empty_input = Input::new();
+
+        let zkvm = EreRisc0::new(elf_bytes, ProverResourceType::Cpu);
+        let prove_result = zkvm.prove(&empty_input);
+        assert!(prove_result.is_err());
+    }
+}
+
+#[cfg(test)]
+mod execute_tests {
+    use std::path::PathBuf;
+
+    use super::*;
+    use zkvm_interface::Input;
+
+    fn get_compiled_test_r0_elf() -> Result<Risc0Program, RiscZeroError> {
+        let test_guest_path = get_execute_test_guest_program_path();
+        RV32_IM_RISCZERO_ZKVM_ELF::compile(&test_guest_path)
+    }
+
+    fn get_execute_test_guest_program_path() -> PathBuf {
+        let workspace_dir = env!("CARGO_WORKSPACE_DIR");
+        PathBuf::from(workspace_dir)
+            .join("tests")
+            .join("risczero")
+            .join("compile")
+            .join("project_structure_build")
+            .canonicalize()
+            .expect("Failed to find or canonicalize test Risc0 methods crate")
+    }
+
+    #[test]
+    fn test_execute_r0_dummy_input() {
+        let program = get_compiled_test_r0_elf().unwrap();
+
+        let mut input_builder = Input::new();
+        let n: u32 = 42;
+        let a: u16 = 42;
+        input_builder.write(&n).unwrap();
+        input_builder.write(&a).unwrap();
+
+        let zkvm = EreRisc0::new(program, ProverResourceType::Cpu);
+
+        let result = zkvm.execute(&input_builder);
+
+        if let Err(e) = &result {
+            panic!("Execution error: {:?}", e);
+        }
+    }
+
+    #[test]
+    fn test_execute_r0_no_input_for_guest_expecting_input() {
+        let program = get_compiled_test_r0_elf().unwrap();
+
+        let empty_input = Input::new();
+
+        let zkvm = EreRisc0::new(program, ProverResourceType::Cpu);
+        let result = zkvm.execute(&empty_input);
+
+        assert!(
+            result.is_err(),
+            "execute should fail if guest expects input but none is provided."
+        );
+    }
+}