Support precompile

Cargo.toml

@@ -4,6 +4,7 @@ members = [
    "host-program/sha2-rust",
    "host-program/revme",
    "host-program/mem-alloc-vec",
+   "host-program/sha2-precompile",
    "sdk"
   ]
 resolver = "2"

guest-program/sha2-precompile/Cargo.toml

@@ -0,0 +1,9 @@
+[workspace]
+[package]
+version = "0.1.0"
+name = "sha2-precompile"
+edition = "2021"
+
+[dependencies]
+zkm-runtime = { git = "https://github.com/zkMIPS/zkm.git"}
+sha2 = { version = "0.10.8", default-features = false }

guest-program/sha2-precompile/src/main.rs

@@ -0,0 +1,24 @@
+#![no_std]
+#![no_main]
+
+use sha2::{Digest, Sha256};
+extern crate alloc;
+use alloc::vec::Vec;
+
+zkm_runtime::entrypoint!(main);
+
+pub fn main() {
+    let public_input: Vec<u8> = zkm_runtime::io::read();
+    let input: [u8; 32] = zkm_runtime::io::read();
+    let elf_id: Vec<u8> = zkm_runtime::io::read();
+
+    zkm_runtime::io::verify(elf_id, &input);
+    let mut hasher = Sha256::new();
+    hasher.update(input.to_vec());
+    let result = hasher.finalize();
+
+    let output: [u8; 32] = result.into();
+    assert_eq!(output.to_vec(), public_input);
+
+    zkm_runtime::io::commit::<[u8; 32]>(&output);
+}

host-program/mem-alloc-vec/src/main.rs

@@ -59,6 +59,9 @@ async fn main() -> Result<()> {
         private_inputstream: vec![],
         seg_size,
         execute_only,
+        precompile: false,
+        receipt_inputs: vec![],
+        receipts: vec![],
     };
 
     //If the guest program does't have inputs, it does't need the set_guest_input().

host-program/revme/src/main.rs

@@ -60,6 +60,9 @@ async fn main() -> Result<()> {
         private_inputstream: vec![],
         seg_size,
         execute_only,
+        precompile: false,
+        receipt_inputs: vec![],
+        receipts: vec![],
     };
 
     //If the guest program does't have inputs, it does't need the setting.

host-program/run-proving.sh

@@ -34,6 +34,8 @@ if [ "$program" == "sha2-rust" ];then
     export SEG_SIZE=65536
 elif [ "$program" == "mem-alloc-vec" ];then
      export SEG_SIZE=65536
+elif [ "$program" == "sha2-precompile" ];then
+     export SEG_SIZE=16384
 fi
 
 echo "SEG_SIZE:$SEG_SIZE"

host-program/sha2-go/src/main.rs

@@ -66,6 +66,9 @@ async fn main() -> Result<()> {
         private_inputstream: vec![],
         seg_size,
         execute_only,
+        precompile: false,
+        receipt_inputs: vec![],
+        receipts: vec![],
     };
 
     //If the guest program does't have inputs, it does't need the setting.

host-program/sha2-precompile/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "sha2-precompile"
+version = "0.1.0"
+edition = "2021"
+
+
+[dependencies]
+zkm-sdk = { path = "../../sdk", features = ["snark"] }
+bincode = "1.3.3"
+sha2 = { version = "0.10.8", default-features = false }
+
+tokio = { version = "1.21.0", features = ["macros", "rt-multi-thread", "signal"] }
+#ethers = "2.0.14"
+
+log = { version = "0.4.14", default-features = false }
+anyhow = "1.0.75"
+env_logger = "0.10.0"
+
+[build-dependencies]
+tonic-build = "0.8.0"
+zkm-build = { git = "https://github.com/zkMIPS/zkm", branch = "main", default-features = false }
+
+[features]
+test = []

host-program/sha2-precompile/build.rs

@@ -0,0 +1,30 @@
+fn main() {
+    let pre_guest_path = format!(
+        "{}/../../guest-program/sha2-rust",
+        env!("CARGO_MANIFEST_DIR")
+    );
+    zkm_build::build_program(&pre_guest_path);
+    let pre_guest_target_path = format!(
+        "{}/{}/{}",
+        pre_guest_path,
+        zkm_build::DEFAULT_OUTPUT_DIR,
+        zkm_build::BUILD_TARGET
+    );
+    println!(
+        "cargo:rustc-env=PRE_GUEST_TARGET_PATH={}",
+        pre_guest_target_path
+    );
+
+    let guest_path = format!(
+        "{}/../../guest-program/sha2-precompile",
+        env!("CARGO_MANIFEST_DIR")
+    );
+    zkm_build::build_program(&guest_path);
+    let guest_target_path = format!(
+        "{}/{}/{}",
+        guest_path,
+        zkm_build::DEFAULT_OUTPUT_DIR,
+        zkm_build::BUILD_TARGET
+    );
+    println!("cargo:rustc-env=GUEST_TARGET_PATH={}", guest_target_path);
+}

host-program/sha2-precompile/src/main.rs

@@ -0,0 +1,195 @@
+use anyhow::bail;
+use anyhow::Result;
+use sha2::{Digest, Sha256};
+use std::env;
+use std::fs::read;
+use std::time::Instant;
+use zkm_sdk::{prover::ClientCfg, prover::ProverInput, ProverClient};
+
+pub const DEFAULT_PROVER_NETWORK_RPC: &str = "https://152.32.186.45:20002";
+pub const DEFALUT_PROVER_NETWORK_DOMAIN: &str = "stage";
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    env_logger::try_init().unwrap_or_default();
+    let seg_size = env::var("SEG_SIZE")
+        .ok()
+        .and_then(|seg| seg.parse::<u32>().ok())
+        .unwrap_or(65536);
+
+    let execute_only = env::var("EXECUTE_ONLY")
+        .ok()
+        .and_then(|seg| seg.parse::<bool>().ok())
+        .unwrap_or(false);
+
+    let elf_path = env::var("ELF_PATH").unwrap_or(env!("GUEST_TARGET_PATH").to_string());
+    let pre_elf_path =
+        env::var("PRE_ELF_PATH").unwrap_or(env!("PRE_GUEST_TARGET_PATH").to_string());
+    let proof_results_path = env::var("PROOF_RESULTS_PATH").unwrap_or("../contracts".to_string());
+    let vk_path = env::var("VERIFYING_KEY_PATH").unwrap_or("/tmp/input".to_string());
+
+    //network proving
+    let endpoint = env::var("ENDPOINT").unwrap_or(DEFAULT_PROVER_NETWORK_RPC.to_string());
+    let ca_cert_path = env::var("CA_CERT_PATH").unwrap_or("".to_string());
+    let cert_path = env::var("CERT_PATH").unwrap_or("".to_string());
+    let key_path = env::var("KEY_PATH").unwrap_or("".to_string());
+    let domain_name = env::var("DOMAIN_NAME").unwrap_or(DEFALUT_PROVER_NETWORK_DOMAIN.to_string());
+    let private_key = env::var("PRIVATE_KEY").unwrap_or("".to_string());
+    let zkm_prover_type = env::var("ZKM_PROVER").expect("ZKM PROVER is missing");
+
+    let client_config: ClientCfg = ClientCfg {
+        zkm_prover: zkm_prover_type.to_owned(),
+        endpoint: Some(endpoint),
+        ca_cert_path: Some(ca_cert_path),
+        cert_path: Some(cert_path),
+        key_path: Some(key_path),
+        domain_name: Some(domain_name),
+        private_key: Some(private_key),
+        vk_path: vk_path.to_owned(),
+    };
+
+    log::info!("new prover client:");
+    let prover_client = ProverClient::new(&client_config).await;
+    log::info!("new prover client,ok.");
+
+    let mut prover_input = ProverInput {
+        elf: read(pre_elf_path).unwrap(),
+        public_inputstream: vec![],
+        private_inputstream: vec![],
+        seg_size: 0,
+        execute_only: false,
+        precompile: true,
+        receipt_inputs: vec![],
+        receipts: vec![],
+    };
+
+    set_pre_guest_input(&mut prover_input, None);
+
+    let start = Instant::now();
+    let proving_result = prover_client.prover.prove(&prover_input, None).await;
+    let mut receipts = vec![];
+    let pre_elf_id: Vec<u8>;
+    match proving_result {
+        Ok(Some(prover_result)) => {
+            prover_client
+                .print_guest_execution_output(true, &prover_result)
+                .expect("print pre guest program excution's output false.");
+            receipts.push(prover_result.receipt);
+            pre_elf_id = prover_result.elf_id;
+            log::info!("pre_elf_id: {:?}", pre_elf_id);
+        }
+        Ok(None) => {
+            log::info!("Failed to generate proof.The result is None.");
+            bail!("Failed to generate proof.");
+        }
+        Err(e) => {
+            log::info!("Failed to generate proof. error: {}", e);
+            bail!("Failed to generate proof.");
+        }
+    }
+
+    let end = Instant::now();
+    let elapsed = end.duration_since(start);
+    log::info!("Elapsed time: {:?} secs", elapsed.as_secs());
+
+    let mut prover_input = ProverInput {
+        elf: read(elf_path).unwrap(),
+        public_inputstream: vec![],
+        private_inputstream: vec![],
+        seg_size,
+        execute_only,
+        precompile: false,
+        receipt_inputs: vec![],
+        receipts,
+    };
+
+    set_guest_input(&mut prover_input, &pre_elf_id);
+
+    let start = Instant::now();
+    let proving_result = prover_client.prover.prove(&prover_input, None).await;
+    match proving_result {
+        Ok(Some(prover_result)) => {
+            if !execute_only {
+                //excute the guest program and generate the proof
+                prover_client
+                    .process_proof_results(
+                        &prover_result,
+                        &prover_input,
+                        &proof_results_path,
+                        &zkm_prover_type,
+                    )
+                    .expect("process proof results error");
+            } else {
+                //only excute the guest program without generating the proof.
+                //the sha2-rust guest program has outputs messages, which are basic type.
+                prover_client
+                    .print_guest_execution_output(true, &prover_result)
+                    .expect("print guest program excution's output false.");
+            }
+        }
+        Ok(None) => {
+            log::info!("Failed to generate proof.The result is None.");
+            bail!("Failed to generate proof.");
+        }
+        Err(e) => {
+            log::info!("Failed to generate proof. error: {}", e);
+            bail!("Failed to generate proof.");
+        }
+    }
+
+    let end = Instant::now();
+    let elapsed = end.duration_since(start);
+    log::info!("Elapsed time: {:?} secs", elapsed.as_secs());
+    Ok(())
+}
+
+fn set_pre_guest_input(input: &mut ProverInput, _param: Option<&str>) {
+    let num_bytes: usize = 1024;
+    let pri_input = vec![5u8; num_bytes];
+    let mut hasher = Sha256::new();
+    hasher.update(&pri_input);
+    let result = hasher.finalize();
+    let output: [u8; 32] = result.into();
+
+    // assume the  arg[0] = hash(public input), and the arg[1] = public input.
+    let public_input = output.to_vec();
+    let mut pub_buf = Vec::new();
+    bincode::serialize_into(&mut pub_buf, &public_input)
+        .expect("public_input serialization failed");
+
+    let mut pri_buf = Vec::new();
+    bincode::serialize_into(&mut pri_buf, &pri_input).expect("private_input serialization failed");
+
+    input.public_inputstream = pub_buf;
+    input.private_inputstream = pri_buf;
+}
+
+fn set_guest_input(input: &mut ProverInput, elf_id: &Vec<u8>) {
+    let num_bytes: usize = 1024;
+    let pri_input = vec![5u8; num_bytes];
+    let mut hasher = Sha256::new();
+    hasher.update(&pri_input);
+    let result = hasher.finalize();
+    let output: [u8; 32] = result.into();
+
+    // assume the  arg[0] = hash(public input), and the arg[1] = public input.
+    let mut pre_pub_buf = Vec::new();
+    bincode::serialize_into(&mut pre_pub_buf, &output).expect("public_input serialization failed");
+
+    let mut hasher = Sha256::new();
+    hasher.update(output);
+    let result = hasher.finalize();
+    let output: [u8; 32] = result.into();
+
+    let public_input = output.to_vec();
+    let mut pub_buf = Vec::new();
+    bincode::serialize_into(&mut pub_buf, &public_input)
+        .expect("public_input serialization failed");
+
+    let mut elf_id_buf = Vec::new();
+    bincode::serialize_into(&mut elf_id_buf, elf_id).expect("elf_id serialization failed");
+
+    input.public_inputstream = pub_buf;
+    input.private_inputstream = pre_pub_buf;
+    input.receipt_inputs.push(elf_id_buf);
+}

host-program/sha2-rust/src/main.rs

@@ -61,6 +61,9 @@ async fn main() -> Result<()> {
         private_inputstream: vec![],
         seg_size,
         execute_only,
+        precompile: false,
+        receipt_inputs: vec![],
+        receipts: vec![],
     };
 
     //If the guest program does't have inputs, it does't need the setting.

sdk/src/local/prover.rs

@@ -36,7 +36,7 @@ impl ProverTask {
         let outputdir = format!("/tmp/{}/output", self.proof_id);
         fs::create_dir_all(&inputdir).unwrap();
         fs::create_dir_all(&outputdir).unwrap();
-        let should_agg =
+        let (should_agg, receipt, elf_id) =
             crate::local::stark::prove_stark(&self.input, &inputdir, &mut result).unwrap();
         if self.input.execute_only {
             result.proof_with_public_inputs = vec![];
@@ -47,7 +47,7 @@ impl ProverTask {
                 "There is only one segment with segment size {}, will skip the aggregation!",
                 self.input.seg_size
             );
-        } else {
+        } else if !self.input.precompile {
             match crate::local::snark::prove_snark(&vk_path, &inputdir, &outputdir) {
                 Ok(()) => {
                     result.stark_proof =
@@ -66,6 +66,10 @@ impl ProverTask {
                 }
             }
         }
+        if let Some(receipt) = receipt {
+            result.receipt.clone_from(&receipt);
+            result.elf_id.clone_from(&elf_id.unwrap());
+        }
         self.result = Some(result);
         self.is_done = true;
     }
@@ -141,7 +145,8 @@ impl Prover for LocalProver {
         }
         fs::create_dir_all(vk_path).unwrap();
 
-        let should_agg = crate::local::stark::prove_stark(input, vk_path, &mut result).unwrap();
+        let (should_agg, _, _) =
+            crate::local::stark::prove_stark(input, vk_path, &mut result).unwrap();
         if !should_agg {
             log::info!("Setup: generating the stark proof false, please check the SEG_SIZE or other parameters.");
             bail!("Setup: generating the stark proof false, please check the SEG_SIZE or other parameters!");