From 0c03b5bb3a5b25d5d2ca238fc51d12bb9aaaeafc Mon Sep 17 00:00:00 2001 From: "Xin, Gao" Date: Fri, 5 Jan 2024 19:50:54 +0800 Subject: [PATCH] More friendly host apis (#224) * add wasm_witness * support standard env in cli * change wasm_witness_inject to wasm_witness_insert * support wasm_dbg_char * move witness api into external foreign host * fix: add missing files * use better merkle circuit from zkwasm-host-circuits * cargo fmt * remove dryrun service from cli * merge dryrun and run by specify dryrun flag * fix dry_run * fix tests * add wasm_witness * support standard env in cli * change wasm_witness_inject to wasm_witness_insert * support wasm_dbg_char * move witness api into external foreign host * fix: add missing files * use better merkle circuit from zkwasm-host-circuits * cargo fmt * add keccak host support * support keccak * bump zkwasm-host-circuits for zkwasm * bump wasmi version * fix uniform circuit verify * bump zkwasm-host-circuits * bump zkwasm-host-circuits * split merkle and datahash * cleanup merkle host to support hash leaf only * support datacache * bump zkwasm host circuits to main * get statics for host-ops * add name for plugin context * fix tests compilation * support tracer in ExecEnv and can be visited in host api * polish the host api to get trace count * reset keccak256 cursor in keccak_new() * fix tests * using enum for host mode * feature: support indexed witness in witness host * bump wasmi version * remove useless Sequence * adjust circuit_with_witness and put println! outside in cli * add comment for picking the first instance data * add default implement for ForeignContext * fix tests * cargo fmt code * deduce the type of Arg for HostEnvBuilder * fix examples * fix index for witness host api * cargo fmt after merge * expose witness host context to arg * add zkwasm_indexed_witness_push * use local cache * fix counter of merkle * bugfix: witness push and insert --------- Co-authored-by: Zhang Junyu Co-authored-by: Qi Zhou --- Cargo.lock | 21 ++- README.md | 8 +- crates/cli/src/app_builder.rs | 166 ++++++++++++----- crates/cli/src/args.rs | 24 +++ crates/cli/src/exec.rs | 84 ++++++--- crates/cli/test_cli.sh | 8 +- .../src/host/ecc_helper/bls381/pair/mod.rs | 22 ++- .../src/host/ecc_helper/bls381/sum/mod.rs | 19 +- .../src/host/ecc_helper/bn254/pair/mod.rs | 21 ++- .../host/src/host/ecc_helper/bn254/sum/mod.rs | 25 ++- .../src/host/ecc_helper/jubjub/sum/mod.rs | 23 ++- crates/host/src/host/hash_helper/keccak256.rs | 131 +++++++++++++ crates/host/src/host/hash_helper/mod.rs | 1 + crates/host/src/host/hash_helper/poseidon.rs | 25 ++- crates/host/src/host/hash_helper/sha256.rs | 15 +- .../host/src/host/merkle_helper/datacache.rs | 176 ++++++++++++++++++ crates/host/src/host/merkle_helper/merkle.rs | 150 ++++----------- crates/host/src/host/merkle_helper/mod.rs | 1 + crates/host/src/host/mod.rs | 1 + crates/host/src/host/witness_helper/mod.rs | 176 ++++++++++++++++++ crates/host/src/lib.rs | 78 ++++++-- crates/playground/Cargo.lock | 15 +- crates/playground/examples/binary_search.rs | 8 +- crates/playground/examples/context.rs | 7 +- crates/playground/examples/fibonacci.rs | 8 +- crates/playground/examples/phantom.rs | 5 +- crates/zkwasm/Cargo.toml | 2 +- crates/zkwasm/src/foreign/context/runtime.rs | 26 ++- crates/zkwasm/src/foreign/log_helper/mod.rs | 22 ++- .../zkwasm/src/foreign/require_helper/mod.rs | 5 +- .../src/foreign/wasm_input_helper/runtime.rs | 14 +- crates/zkwasm/src/loader/mod.rs | 48 +++-- crates/zkwasm/src/runtime/host/default_env.rs | 21 +-- .../runtime/host/external_circuit_plugin.rs | 23 ++- crates/zkwasm/src/runtime/host/host_env.rs | 22 ++- .../runtime/host/internal_circuit_plugin.rs | 13 +- crates/zkwasm/src/runtime/host/mod.rs | 32 ++-- crates/zkwasm/src/runtime/mod.rs | 4 + .../zkwasm/src/runtime/wasmi_interpreter.rs | 21 ++- crates/zkwasm/src/test/mod.rs | 4 +- crates/zkwasm/src/test/test_rlp.rs | 21 ++- crates/zkwasm/src/test/test_start.rs | 18 +- .../test_wasm_instructions/op_call_host.rs | 12 +- 43 files changed, 1169 insertions(+), 357 deletions(-) create mode 100644 crates/host/src/host/hash_helper/keccak256.rs create mode 100644 crates/host/src/host/merkle_helper/datacache.rs create mode 100644 crates/host/src/host/witness_helper/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 8439b13c5..a3b0c01cc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -313,7 +313,7 @@ dependencies = [ [[package]] name = "circuits-batcher" version = "0.1.0" -source = "git+https://github.com/DelphinusLab/continuation-batcher.git#7eccb6358a3b68d13be901791bd0661f0545ea06" +source = "git+https://github.com/DelphinusLab/continuation-batcher.git#8376e7de916b165c8e4965b9b8933e0f3095b50a" dependencies = [ "anyhow", "ark-std", @@ -328,6 +328,7 @@ dependencies = [ "hex", "lazy_static", "log", + "lru", "md5", "num", "num-bigint", @@ -1071,7 +1072,7 @@ dependencies = [ [[package]] name = "halo2_proofs" version = "0.1.0-beta.1" -source = "git+https://github.com/DelphinusLab/halo2-gpu-specific.git#64ccffddc97faff24a00d248270fe84f6693c130" +source = "git+https://github.com/DelphinusLab/halo2-gpu-specific.git#9ac108049b488df640eecdd5eb18751062442898" dependencies = [ "ark-std", "blake2b_simd", @@ -1099,7 +1100,7 @@ dependencies = [ [[package]] name = "halo2aggregator-s" version = "0.1.0" -source = "git+https://github.com/DelphinusLab/halo2aggregator-s.git?branch=main#953bbbd50ee05211e2395e69c9760095b80f6408" +source = "git+https://github.com/DelphinusLab/halo2aggregator-s.git?branch=main#4fc143e5a10a77cb24f31c0b792ec9a01d1bfd34" dependencies = [ "ark-std", "blake2b_simd", @@ -1338,6 +1339,15 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.9" @@ -3306,21 +3316,24 @@ dependencies = [ [[package]] name = "zkwasm-host-circuits" version = "0.1.0" -source = "git+https://github.com/DelphinusLab/zkWasm-host-circuits.git?branch=main#e4dbc533b04963c0e3eb2db5101c0b096a0eea8b" +source = "git+https://github.com/DelphinusLab/zkWasm-host-circuits.git?branch=main#e3a2eff4583b2fd8be7fc3e54f2789cbfbfd72d4" dependencies = [ "ark-std", + "cfg-if 1.0.0", "circuits-batcher", "clap", "ff", "halo2_proofs", "halo2ecc-s", "hex", + "itertools", "lazy_static", "lru", "mongodb", "num-bigint", "poseidon", "rand", + "rand_core", "ripemd", "serde", "serde_json", diff --git a/README.md b/README.md index 27b4ab0ec..9f5923198 100644 --- a/README.md +++ b/README.md @@ -44,18 +44,20 @@ simulation of wasm execution of target wasm bytecode with particular inputs are # Command line: ## Setup via WASM image: ``` -cargo run --release -- --function --wasm setup [OPTIONS] +cargo run --release -- --host default --function --wasm setup [OPTIONS] ``` ## Single prove and verify: ``` -cargo run --release -- --function --wasm single-prove [OPTIONS] -cargo run --release -- --function --wasm single-verify [OPTIONS] +cargo run --release -- --host default --function --wasm single-prove [OPTIONS] +cargo run --release -- --host default --function --wasm single-verify [OPTIONS] ``` with OPTIONS: ``` -h, --help Print help information + --host ... + Which host env you would like to run your binary. [possible values: default, standard] -k [...] Circuit Size K diff --git a/crates/cli/src/app_builder.rs b/crates/cli/src/app_builder.rs index 02dc3a073..dfb7857bb 100644 --- a/crates/cli/src/app_builder.rs +++ b/crates/cli/src/app_builder.rs @@ -1,16 +1,24 @@ use anyhow::Result; use clap::App; use clap::AppSettings; +use delphinus_host::ExecutionArg as StandardArg; +use delphinus_host::HostEnvConfig; +use delphinus_host::StandardHostEnvBuilder as StandardEnvBuilder; use delphinus_zkwasm::circuits::config::MIN_K; use delphinus_zkwasm::runtime::host::default_env::DefaultHostEnvBuilder; use delphinus_zkwasm::runtime::host::default_env::ExecutionArg; + use log::info; +use std::cell::RefCell; +use std::collections::HashMap; use std::fs; use std::io::Write; use std::path::PathBuf; +use std::rc::Rc; use std::sync::Arc; use std::sync::Mutex; +use crate::args::HostMode; use crate::exec::exec_dry_run; use super::command::CommandBuilder; @@ -69,7 +77,8 @@ pub trait AppBuilder: CommandBuilder { .arg(Self::param_path_arg()) .arg(Self::function_name_arg()) .arg(Self::phantom_functions_arg()) - .arg(Self::zkwasm_file_arg()); + .arg(Self::zkwasm_file_arg()) + .arg(Self::host_mode_arg()); let app = Self::append_setup_subcommand(app); let app = Self::append_dry_run_subcommand(app); @@ -113,22 +122,49 @@ pub trait AppBuilder: CommandBuilder { fs::create_dir_all(&output_dir)?; fs::create_dir_all(¶m_dir)?; + let host_mode = Self::parse_host_mode(&top_matches); + match top_matches.subcommand() { - Some(("setup", _)) => exec_setup::( - zkwasm_k, - Self::AGGREGATE_K, - Self::NAME, - wasm_binary, - phantom_functions, - &output_dir, - ¶m_dir, - ), - Some(("checksum", _)) => exec_image_checksum::( - zkwasm_k, - wasm_binary, - phantom_functions, - &output_dir, - ), + Some(("setup", _)) => match host_mode { + HostMode::DEFAULT => exec_setup::( + zkwasm_k, + Self::AGGREGATE_K, + Self::NAME, + wasm_binary, + phantom_functions, + (), + &output_dir, + ¶m_dir, + ), + HostMode::STANDARD => exec_setup::( + zkwasm_k, + Self::AGGREGATE_K, + Self::NAME, + wasm_binary, + phantom_functions, + HostEnvConfig::default(), + &output_dir, + ¶m_dir, + ), + }, + + Some(("checksum", _)) => match host_mode { + HostMode::DEFAULT => exec_image_checksum::( + zkwasm_k, + wasm_binary, + (), + phantom_functions, + &output_dir, + ), + HostMode::STANDARD => exec_image_checksum::( + zkwasm_k, + wasm_binary, + HostEnvConfig::default(), + phantom_functions, + &output_dir, + ), + }, + Some(("dry-run", sub_matches)) => { let public_inputs: Vec = Self::parse_single_public_arg(&sub_matches); let private_inputs: Vec = Self::parse_single_private_arg(&sub_matches); @@ -139,22 +175,43 @@ pub trait AppBuilder: CommandBuilder { let context_output = Arc::new(Mutex::new(vec![])); - exec_dry_run::( - zkwasm_k, - wasm_binary, - phantom_functions, - ExecutionArg { - public_inputs, - private_inputs, - context_inputs: context_in, - context_outputs: context_output.clone(), - }, - )?; + match host_mode { + HostMode::DEFAULT => { + exec_dry_run::( + zkwasm_k, + wasm_binary, + phantom_functions, + ExecutionArg { + public_inputs, + private_inputs, + context_inputs: context_in, + context_outputs: context_output.clone(), + }, + (), + )?; + } + HostMode::STANDARD => { + exec_dry_run::( + zkwasm_k, + wasm_binary, + phantom_functions, + StandardArg { + public_inputs, + private_inputs, + context_inputs: context_in, + context_outputs: context_output.clone(), + indexed_witness: Rc::new(RefCell::new(HashMap::new())), + tree_db: None, + }, + HostEnvConfig::default(), + )?; + } + }; write_context_output(&context_output.lock().unwrap(), context_out_path)?; - Ok(()) } + Some(("single-prove", sub_matches)) => { let public_inputs: Vec = Self::parse_single_public_arg(&sub_matches); let private_inputs: Vec = Self::parse_single_private_arg(&sub_matches); @@ -165,21 +222,44 @@ pub trait AppBuilder: CommandBuilder { let context_out = Arc::new(Mutex::new(vec![])); assert!(public_inputs.len() <= Self::MAX_PUBLIC_INPUT_SIZE); - - exec_create_proof::( - Self::NAME, - zkwasm_k, - wasm_binary, - phantom_functions, - &output_dir, - ¶m_dir, - ExecutionArg { - public_inputs, - private_inputs, - context_inputs: context_in, - context_outputs: context_out.clone(), - }, - )?; + match host_mode { + HostMode::DEFAULT => { + exec_create_proof::( + Self::NAME, + zkwasm_k, + wasm_binary, + phantom_functions, + &output_dir, + ¶m_dir, + ExecutionArg { + public_inputs, + private_inputs, + context_inputs: context_in, + context_outputs: context_out.clone(), + }, + (), + )?; + } + HostMode::STANDARD => { + exec_create_proof::( + Self::NAME, + zkwasm_k, + wasm_binary, + phantom_functions, + &output_dir, + ¶m_dir, + StandardArg { + public_inputs, + private_inputs, + context_inputs: context_in, + context_outputs: context_out.clone(), + indexed_witness: Rc::new(RefCell::new(HashMap::new())), + tree_db: None, + }, + HostEnvConfig::default(), + )?; + } + }; write_context_output(&context_out.lock().unwrap(), context_out_path)?; diff --git a/crates/cli/src/args.rs b/crates/cli/src/args.rs index e933070f0..39b7e1985 100644 --- a/crates/cli/src/args.rs +++ b/crates/cli/src/args.rs @@ -6,6 +6,12 @@ use clap::ArgMatches; use specs::args::parse_args; use std::path::PathBuf; +#[derive(clap::ArgEnum, Clone, Debug)] +pub enum HostMode { + DEFAULT, + STANDARD, +} + pub trait ArgBuilder { fn zkwasm_k_arg<'a>() -> Arg<'a> { arg!( @@ -23,6 +29,7 @@ pub trait ArgBuilder { ) .value_parser(value_parser!(PathBuf)) } + fn parse_zkwasm_file_arg(matches: &ArgMatches) -> PathBuf { matches .get_one::("wasm") @@ -35,6 +42,7 @@ pub trait ArgBuilder { -f --function "Function you would like to run." ) } + fn parse_function_name(matches: &ArgMatches) -> String { matches .get_one::("function") @@ -42,6 +50,22 @@ pub trait ArgBuilder { .to_string() } + fn host_mode_arg<'a>() -> Arg<'a> { + Arg::new("host") + .long("host") + .value_parser(value_parser!(HostMode)) + .action(ArgAction::Set) + .help("Specify host functions set.") + .min_values(0) + .max_values(1) + } + + fn parse_host_mode(matches: &ArgMatches) -> HostMode { + matches + .get_one::("host") + .map_or(HostMode::DEFAULT, |v| v.clone()) + } + fn phantom_functions_arg<'a>() -> Arg<'a> { Arg::new("phantom") .long("phantom") diff --git a/crates/cli/src/exec.rs b/crates/cli/src/exec.rs index f8db80654..690a660da 100644 --- a/crates/cli/src/exec.rs +++ b/crates/cli/src/exec.rs @@ -1,7 +1,9 @@ use anyhow::Result; use circuits_batcher::proof::CircuitInfo; +use circuits_batcher::proof::ParamsCache; use circuits_batcher::proof::ProofInfo; use circuits_batcher::proof::ProofLoadInfo; +use circuits_batcher::proof::ProvingKeyCache; use delphinus_zkwasm::circuits::TestCircuit; use delphinus_zkwasm::loader::ZkWasmLoader; use delphinus_zkwasm::runtime::host::HostEnvBuilder; @@ -15,18 +17,16 @@ use log::info; use std::io::Write; use std::path::PathBuf; -pub fn exec_setup( +pub fn exec_setup( zkwasm_k: u32, aggregate_k: u32, prefix: &str, wasm_binary: Vec, phantom_functions: Vec, + envconfig: Builder::HostConfig, _output_dir: &PathBuf, param_dir: &PathBuf, -) -> Result<()> -where - Builder: HostEnvBuilder, -{ +) -> Result<()> { info!("Setup Params and VerifyingKey"); macro_rules! prepare_params { @@ -54,10 +54,13 @@ where info!("Found Verifying at {:?}", vk_path); } else { info!("Create Verifying to {:?}", vk_path); - let loader = - ZkWasmLoader::::new(zkwasm_k, wasm_binary, phantom_functions)?; + let loader = ZkWasmLoader::::new( + zkwasm_k, + wasm_binary, + phantom_functions, + )?; - let vkey = loader.create_vkey(¶ms)?; + let vkey = loader.create_vkey(¶ms, envconfig)?; let mut fd = std::fs::File::create(&vk_path)?; vkey.write(&mut fd)?; @@ -67,24 +70,28 @@ where Ok(()) } -pub fn exec_image_checksum( +pub fn exec_image_checksum( zkwasm_k: u32, wasm_binary: Vec, + hostenv: Builder::HostConfig, phantom_functions: Vec, output_dir: &PathBuf, ) -> Result<()> where - Builder: HostEnvBuilder, + Builder: HostEnvBuilder, { - let loader = - ZkWasmLoader::::new(zkwasm_k, wasm_binary, phantom_functions)?; + let loader = ZkWasmLoader::::new( + zkwasm_k, + wasm_binary, + phantom_functions, + )?; let params = load_or_build_unsafe_params::( zkwasm_k, Some(&output_dir.join(format!("K{}.params", zkwasm_k))), ); - let checksum = loader.checksum(¶ms)?; + let checksum = loader.checksum(¶ms, hostenv)?; assert_eq!(checksum.len(), 1); let checksum = checksum[0]; @@ -98,31 +105,50 @@ where Ok(()) } -pub fn exec_dry_run>( +pub fn exec_dry_run( zkwasm_k: u32, wasm_binary: Vec, phantom_functions: Vec, - arg: Arg, + arg: Builder::Arg, + config: Builder::HostConfig, ) -> Result<()> { - let loader = - ZkWasmLoader::::new(zkwasm_k, wasm_binary, phantom_functions)?; - loader.run(arg, true, false)?; + let loader = ZkWasmLoader::::new( + zkwasm_k, + wasm_binary, + phantom_functions, + )?; + let result = loader.run(arg, config, true, false)?; + println!("total guest instructions used {:?}", result.guest_statics); + println!("total host api used {:?}", result.host_statics); Ok(()) } -pub fn exec_create_proof>( +pub fn exec_create_proof( prefix: &'static str, zkwasm_k: u32, wasm_binary: Vec, phantom_functions: Vec, output_dir: &PathBuf, param_dir: &PathBuf, - arg: Arg, + arg: Builder::Arg, + config: Builder::HostConfig, ) -> Result<()> { - let loader = - ZkWasmLoader::::new(zkwasm_k, wasm_binary, phantom_functions)?; + let loader = ZkWasmLoader::::new( + zkwasm_k, + wasm_binary, + phantom_functions, + )?; - let (circuit, instances, _) = loader.circuit_with_witness(arg)?; + let execution_result = loader.run(arg, config, false, true)?; + + println!( + "total guest instructions used {:?}", + execution_result.guest_statics + ); + println!("total host api used {:?}", execution_result.host_statics); + println!("application outout {:?}", execution_result.outputs); + + let (circuit, instances) = loader.circuit_with_witness(execution_result)?; if false { info!("Mock test..."); @@ -137,8 +163,18 @@ pub fn exec_create_proof>( zkwasm_k as usize, circuits_batcher::args::HashType::Poseidon, ); + + // save the proof load info for the zkwasm circuit circuit.proofloadinfo.save(output_dir); - circuit.exec_create_proof(output_dir, param_dir, 0); + + // Cli saves zkwasm.0.instance.data as the + // first instance file for .loadinfo + // Thus we provide arg index = 0 to generate a + // proof with the first instance file + let mut param_cache = ParamsCache::new(5); + let mut pkey_cache = ProvingKeyCache::new(5); + + circuit.exec_create_proof(output_dir, param_dir, &mut pkey_cache, 0, &mut param_cache); info!("Proof has been created."); diff --git a/crates/cli/test_cli.sh b/crates/cli/test_cli.sh index bdf25a072..8df568320 100755 --- a/crates/cli/test_cli.sh +++ b/crates/cli/test_cli.sh @@ -6,8 +6,8 @@ set -x rm -rf output/*.data # Single test -RUST_LOG=info cargo run --release --features cuda -- -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm setup -RUST_LOG=info cargo run --release --features cuda -- -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm checksum +RUST_LOG=info cargo run --release --features cuda -- --host default -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm setup +RUST_LOG=info cargo run --release --features cuda -- --host default -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm checksum -RUST_LOG=info cargo run --release --features cuda -- -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm single-prove --public 133:i64 --public 2:i64 -RUST_LOG=info cargo run --release --features cuda -- -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm single-verify +RUST_LOG=info cargo run --release --features cuda -- --host default -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm single-prove --public 133:i64 --public 2:i64 +RUST_LOG=info cargo run --release --features cuda -- --host default -k 18 --function zkmain --param ./params --output ./output --wasm ../zkwasm/wasm/wasm_output.wasm single-verify diff --git a/crates/host/src/host/ecc_helper/bls381/pair/mod.rs b/crates/host/src/host/ecc_helper/bls381/pair/mod.rs index b1e6c6d0c..4e3e6920a 100644 --- a/crates/host/src/host/ecc_helper/bls381/pair/mod.rs +++ b/crates/host/src/host/ecc_helper/bls381/pair/mod.rs @@ -1,14 +1,19 @@ use super::bls381_fq_to_limbs; use super::fetch_fq; use super::fetch_fq2; +use delphinus_zkwasm::circuits::config::zkwasm_k; use delphinus_zkwasm::runtime::host::host_env::HostEnv; use delphinus_zkwasm::runtime::host::ForeignContext; +use delphinus_zkwasm::runtime::host::ForeignStatics; use halo2_proofs::arithmetic::CurveAffine; use halo2_proofs::pairing::bls12_381::pairing; use halo2_proofs::pairing::bls12_381::G1Affine; use halo2_proofs::pairing::bls12_381::G2Affine; use halo2_proofs::pairing::bls12_381::Gt as Bls381Gt; use std::rc::Rc; +use wasmi::tracer::Observer; +use zkwasm_host_circuits::circuits::bls::Bls381PairChip; +use zkwasm_host_circuits::circuits::host::HostOpSelector; use zkwasm_host_circuits::host::ForeignInst; #[derive(Default)] @@ -19,6 +24,7 @@ struct BlsPairContext { pub result_limbs: Vec, pub result_cursor: usize, pub input_cursor: usize, + pub used_round: usize, } impl BlsPairContext { @@ -50,7 +56,14 @@ impl BlsPairContext { } } -impl ForeignContext for BlsPairContext {} +impl ForeignContext for BlsPairContext { + fn get_statics(&self) -> Option { + Some(ForeignStatics { + used_round: self.used_round, + max_round: Bls381PairChip::max_rounds(zkwasm_k() as usize), + }) + } +} use specs::external_host_call_table::ExternalHostCallSignature; pub fn register_blspair_foreign(env: &mut HostEnv) { @@ -64,7 +77,7 @@ pub fn register_blspair_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_blspair_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); if context.input_cursor == 16 { let t: u64 = args.nth(0); @@ -84,7 +97,7 @@ pub fn register_blspair_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_blspair_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); if context.input_cursor == 32 { let t: u64 = args.nth(0); @@ -107,6 +120,7 @@ pub fn register_blspair_foreign(env: &mut HostEnv) { let ab = pairing(&g1, &g2); log::debug!("gt {:?}", ab); context.bls381_gt_to_limbs(ab); + context.used_round += 1; } else { context.limbs.push(args.nth(0)); context.input_cursor += 1; @@ -122,7 +136,7 @@ pub fn register_blspair_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Return, foreign_blspair_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); let ret = Some(wasmi::RuntimeValue::I64( context.result_limbs[context.result_cursor] as i64, diff --git a/crates/host/src/host/ecc_helper/bls381/sum/mod.rs b/crates/host/src/host/ecc_helper/bls381/sum/mod.rs index 4679c5938..71b199c16 100644 --- a/crates/host/src/host/ecc_helper/bls381/sum/mod.rs +++ b/crates/host/src/host/ecc_helper/bls381/sum/mod.rs @@ -1,12 +1,17 @@ +use delphinus_zkwasm::circuits::config::zkwasm_k; use delphinus_zkwasm::runtime::host::host_env::HostEnv; use delphinus_zkwasm::runtime::host::ForeignContext; +use delphinus_zkwasm::runtime::host::ForeignStatics; use halo2_proofs::pairing::bls12_381::G1Affine; use std::ops::Add; use std::rc::Rc; +use wasmi::tracer::Observer; use super::bls381_fq_to_limbs; use super::fetch_g1; +use zkwasm_host_circuits::circuits::bls::Bls381SumChip; +use zkwasm_host_circuits::circuits::host::HostOpSelector; use zkwasm_host_circuits::host::ForeignInst; #[derive(Default)] @@ -16,6 +21,7 @@ struct BlsSumContext { pub result_limbs: Option>, pub result_cursor: usize, pub input_cursor: usize, + pub used_round: usize, } impl BlsSumContext { @@ -32,7 +38,14 @@ impl BlsSumContext { } } -impl ForeignContext for BlsSumContext {} +impl ForeignContext for BlsSumContext { + fn get_statics(&self) -> Option { + Some(ForeignStatics { + used_round: self.used_round, + max_round: Bls381SumChip::max_rounds(zkwasm_k() as usize), + }) + } +} use specs::external_host_call_table::ExternalHostCallSignature; pub fn register_blssum_foreign(env: &mut HostEnv) { @@ -46,7 +59,7 @@ pub fn register_blssum_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_blssum_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); if context.input_cursor == 16 { let t: u64 = args.nth(0); @@ -67,7 +80,7 @@ pub fn register_blssum_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Return, foreign_blssum_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.result_limbs.clone().map_or_else( || { diff --git a/crates/host/src/host/ecc_helper/bn254/pair/mod.rs b/crates/host/src/host/ecc_helper/bn254/pair/mod.rs index 6b00b32b8..5f64e01a7 100644 --- a/crates/host/src/host/ecc_helper/bn254/pair/mod.rs +++ b/crates/host/src/host/ecc_helper/bn254/pair/mod.rs @@ -1,5 +1,7 @@ +use delphinus_zkwasm::circuits::config::zkwasm_k; use delphinus_zkwasm::runtime::host::host_env::HostEnv; use delphinus_zkwasm::runtime::host::ForeignContext; +use delphinus_zkwasm::runtime::host::ForeignStatics; use halo2_proofs::arithmetic::CurveAffine; use halo2_proofs::pairing::bn256::pairing; use halo2_proofs::pairing::bn256::G1Affine; @@ -7,12 +9,15 @@ use halo2_proofs::pairing::bn256::G2Affine; use halo2_proofs::pairing::bn256::Gt as BN254Gt; use halo2_proofs::pairing::group::prime::PrimeCurveAffine; use std::rc::Rc; +use wasmi::tracer::Observer; use super::bn254_fq_to_limbs; use super::fetch_fq; use super::fetch_fq2; use super::LIMBNB; +use zkwasm_host_circuits::circuits::bn256::Bn256PairChip; +use zkwasm_host_circuits::circuits::host::HostOpSelector; use zkwasm_host_circuits::host::ForeignInst::Bn254PairG1; use zkwasm_host_circuits::host::ForeignInst::Bn254PairG2; use zkwasm_host_circuits::host::ForeignInst::Bn254PairG3; @@ -26,6 +31,7 @@ struct BN254PairContext { pub result_limbs: Vec, pub result_cursor: usize, pub input_cursor: usize, + pub used_round: usize, } impl BN254PairContext { @@ -57,7 +63,14 @@ impl BN254PairContext { } } -impl ForeignContext for BN254PairContext {} +impl ForeignContext for BN254PairContext { + fn get_statics(&self) -> Option { + Some(ForeignStatics { + used_round: self.used_round, + max_round: Bn256PairChip::max_rounds(zkwasm_k() as usize), + }) + } +} use specs::external_host_call_table::ExternalHostCallSignature; pub fn register_bn254pair_foreign(env: &mut HostEnv) { @@ -71,7 +84,7 @@ pub fn register_bn254pair_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_blspair_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); if context.input_cursor == LIMBNB * 2 { let t: u64 = args.nth(0); @@ -91,7 +104,7 @@ pub fn register_bn254pair_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_blspair_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); if context.input_cursor == LIMBNB * 4 { let t: u64 = args.nth(0); @@ -136,7 +149,7 @@ pub fn register_bn254pair_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Return, foreign_blspair_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); if context.result_cursor == 0 { let gt = context.gt.unwrap(); diff --git a/crates/host/src/host/ecc_helper/bn254/sum/mod.rs b/crates/host/src/host/ecc_helper/bn254/sum/mod.rs index 2ae7a8301..a63c16ef9 100644 --- a/crates/host/src/host/ecc_helper/bn254/sum/mod.rs +++ b/crates/host/src/host/ecc_helper/bn254/sum/mod.rs @@ -1,9 +1,14 @@ +use delphinus_zkwasm::circuits::config::zkwasm_k; use delphinus_zkwasm::runtime::host::host_env::HostEnv; use delphinus_zkwasm::runtime::host::ForeignContext; +use delphinus_zkwasm::runtime::host::ForeignStatics; use halo2_proofs::pairing::bn256::G1Affine; use halo2_proofs::pairing::group::prime::PrimeCurveAffine; use std::ops::Add; use std::rc::Rc; +use wasmi::tracer::Observer; +use zkwasm_host_circuits::circuits::bn256::Bn256SumChip; +use zkwasm_host_circuits::circuits::host::HostOpSelector; use zkwasm_host_circuits::host::ForeignInst::Bn254SumG1; use zkwasm_host_circuits::host::ForeignInst::Bn254SumNew; use zkwasm_host_circuits::host::ForeignInst::Bn254SumResult; @@ -19,6 +24,7 @@ struct BN254SumContext { pub coeffs: Vec, pub result_limbs: Option>, pub result_cursor: usize, + pub used_round: usize, } impl BN254SumContext { @@ -41,6 +47,7 @@ impl BN254SumContext { coeffs: vec![], result_limbs: None, result_cursor: 0, + used_round: 0, } } @@ -53,6 +60,7 @@ impl BN254SumContext { if new != 0 { G1Affine::identity(); } + self.used_round += 1; } fn bn254_sum_push_scalar(&mut self, v: u64) { @@ -66,7 +74,14 @@ impl BN254SumContext { } } -impl ForeignContext for BN254SumContext {} +impl ForeignContext for BN254SumContext { + fn get_statics(&self) -> Option { + Some(ForeignStatics { + used_round: self.used_round, + max_round: Bn256SumChip::max_rounds(zkwasm_k() as usize), + }) + } +} /* * ForeignInst::Bn254SumNew @@ -87,7 +102,7 @@ pub fn register_bn254sum_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_bn254sum_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.bn254_sum_new(args.nth::(0) as usize); None @@ -101,7 +116,7 @@ pub fn register_bn254sum_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_bn254sum_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.bn254_sum_push_scalar(args.nth::(0)); None @@ -115,7 +130,7 @@ pub fn register_bn254sum_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_bn254sum_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.bn254_sum_push_limb(args.nth::(0)); None @@ -129,7 +144,7 @@ pub fn register_bn254sum_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Return, foreign_bn254sum_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); log::debug!("calculate finalize"); context.result_limbs.clone().map_or_else( diff --git a/crates/host/src/host/ecc_helper/jubjub/sum/mod.rs b/crates/host/src/host/ecc_helper/jubjub/sum/mod.rs index 45b7fe65f..e61d6a3c2 100644 --- a/crates/host/src/host/ecc_helper/jubjub/sum/mod.rs +++ b/crates/host/src/host/ecc_helper/jubjub/sum/mod.rs @@ -1,8 +1,13 @@ +use delphinus_zkwasm::circuits::config::zkwasm_k; use delphinus_zkwasm::runtime::host::host_env::HostEnv; use delphinus_zkwasm::runtime::host::ForeignContext; +use delphinus_zkwasm::runtime::host::ForeignStatics; use num_bigint::BigUint; use std::rc::Rc; +use wasmi::tracer::Observer; +use zkwasm_host_circuits::circuits::babyjub::AltJubChip; +use zkwasm_host_circuits::circuits::host::HostOpSelector; use zkwasm_host_circuits::host::ForeignInst::JubjubSumNew; use zkwasm_host_circuits::host::ForeignInst::JubjubSumPush; use zkwasm_host_circuits::host::ForeignInst::JubjubSumResult; @@ -31,6 +36,7 @@ pub struct BabyJubjubSumContext { pub result_limbs: Option>, pub result_cursor: usize, pub input_cursor: usize, + pub used_round: usize, } impl BabyJubjubSumContext { @@ -42,6 +48,7 @@ impl BabyJubjubSumContext { result_limbs: None, result_cursor: 0, input_cursor: 0, + used_round: 0, } } @@ -51,6 +58,7 @@ impl BabyJubjubSumContext { self.limbs = vec![]; self.input_cursor = 0; self.coeffs = vec![]; + self.used_round += 1; if new != 0 { self.acc = jubjub::Point::identity(); } @@ -105,7 +113,14 @@ impl BabyJubjubSumContext { } } -impl ForeignContext for BabyJubjubSumContext {} +impl ForeignContext for BabyJubjubSumContext { + fn get_statics(&self) -> Option { + Some(ForeignStatics { + used_round: self.used_round, + max_round: AltJubChip::max_rounds(zkwasm_k() as usize), + }) + } +} use specs::external_host_call_table::ExternalHostCallSignature; pub fn register_babyjubjubsum_foreign(env: &mut HostEnv) { @@ -120,7 +135,7 @@ pub fn register_babyjubjubsum_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_babyjubjubsum_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.babyjubjub_sum_new(args.nth::(0) as usize); None @@ -134,7 +149,7 @@ pub fn register_babyjubjubsum_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_babyjubjubsum_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.babyjubjub_sum_push(args.nth(0)); None @@ -148,7 +163,7 @@ pub fn register_babyjubjubsum_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Return, foreign_babyjubjubsum_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); let ret = Some(wasmi::RuntimeValue::I64( context.babyjubjub_sum_finalize() as i64 diff --git a/crates/host/src/host/hash_helper/keccak256.rs b/crates/host/src/host/hash_helper/keccak256.rs new file mode 100644 index 000000000..a0e6e5f22 --- /dev/null +++ b/crates/host/src/host/hash_helper/keccak256.rs @@ -0,0 +1,131 @@ +use delphinus_zkwasm::circuits::config::zkwasm_k; +use delphinus_zkwasm::runtime::host::host_env::HostEnv; +use delphinus_zkwasm::runtime::host::ForeignContext; +use delphinus_zkwasm::runtime::host::ForeignStatics; +use std::rc::Rc; +use wasmi::tracer::Observer; +use zkwasm_host_circuits::circuits::host::HostOpSelector; +use zkwasm_host_circuits::circuits::keccak256::KeccakChip; +use zkwasm_host_circuits::host::keccak256::Keccak; +use zkwasm_host_circuits::host::ForeignInst::Keccak256Finalize; +use zkwasm_host_circuits::host::ForeignInst::Keccak256New; +use zkwasm_host_circuits::host::ForeignInst::Keccak256Push; + +pub use zkwasm_host_circuits::host::keccak256::KECCAK_HASHER; + +struct Generator { + pub cursor: usize, + pub values: Vec, +} + +impl Generator { + fn gen(&mut self) -> u64 { + let r = self.values[self.cursor]; + self.cursor += 1; + r + } +} + +struct Keccak256Context { + pub hasher: Option, + pub generator: Generator, + pub buf: Vec, + pub used_round: usize, +} + +impl Keccak256Context { + fn default() -> Self { + Keccak256Context { + hasher: None, + generator: Generator { + cursor: 0, + values: vec![], + }, + buf: vec![], + used_round: 0, + } + } + + pub fn keccak_new(&mut self, new: usize) { + self.buf = vec![]; + self.generator.cursor = 0; + if new != 0 { + self.hasher = Some(KECCAK_HASHER.clone()); + self.used_round += 1; + } + } + + pub fn keccak_push(&mut self, v: u64) { + self.buf.push(v); + } + + pub fn keccak_finalize(&mut self) -> u64 { + assert!(self.buf.len() == 17); + if self.generator.cursor == 0 { + self.hasher.as_mut().map(|s| { + log::debug!("perform hash with {:?}", self.buf); + let r = s.update_exact(&self.buf.clone().try_into().unwrap()); + self.generator.values = r.to_vec(); + }); + } + self.generator.gen() + } +} + +impl ForeignContext for Keccak256Context { + fn get_statics(&self) -> Option { + Some(ForeignStatics { + used_round: self.used_round, + max_round: KeccakChip::max_rounds(zkwasm_k() as usize), + }) + } +} + +use specs::external_host_call_table::ExternalHostCallSignature; +pub fn register_keccak_foreign(env: &mut HostEnv) { + let foreign_keccak_plugin = env + .external_env + .register_plugin("foreign_keccak", Box::new(Keccak256Context::default())); + + env.external_env.register_function( + "keccak_new", + Keccak256New as usize, + ExternalHostCallSignature::Argument, + foreign_keccak_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + log::debug!("buf len is {}", context.buf.len()); + context.keccak_new(args.nth::(0) as usize); + None + }, + ), + ); + + env.external_env.register_function( + "keccak_push", + Keccak256Push as usize, + ExternalHostCallSignature::Argument, + foreign_keccak_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + context.keccak_push(args.nth::(0) as u64); + None + }, + ), + ); + + env.external_env.register_function( + "keccak_finalize", + Keccak256Finalize as usize, + ExternalHostCallSignature::Return, + foreign_keccak_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + Some(wasmi::RuntimeValue::I64(context.keccak_finalize() as i64)) + }, + ), + ); +} diff --git a/crates/host/src/host/hash_helper/mod.rs b/crates/host/src/host/hash_helper/mod.rs index dbfe2a9d7..3a78e996b 100644 --- a/crates/host/src/host/hash_helper/mod.rs +++ b/crates/host/src/host/hash_helper/mod.rs @@ -1,2 +1,3 @@ +pub mod keccak256; pub mod poseidon; pub mod sha256; diff --git a/crates/host/src/host/hash_helper/poseidon.rs b/crates/host/src/host/hash_helper/poseidon.rs index d924f7414..f98a7e9fc 100644 --- a/crates/host/src/host/hash_helper/poseidon.rs +++ b/crates/host/src/host/hash_helper/poseidon.rs @@ -1,14 +1,19 @@ +use delphinus_zkwasm::circuits::config::zkwasm_k; use delphinus_zkwasm::runtime::host::host_env::HostEnv; use delphinus_zkwasm::runtime::host::ForeignContext; +use delphinus_zkwasm::runtime::host::ForeignStatics; use ff::PrimeField; use halo2_proofs::pairing::bn256::Fr; use poseidon::Poseidon; use std::rc::Rc; +use wasmi::tracer::Observer; pub use zkwasm_host_circuits::host::poseidon::POSEIDON_HASHER; use zkwasm_host_circuits::host::Reduce; use zkwasm_host_circuits::host::ReduceRule; +use zkwasm_host_circuits::circuits::host::HostOpSelector; +use zkwasm_host_circuits::circuits::poseidon::PoseidonChip; use zkwasm_host_circuits::host::ForeignInst::PoseidonFinalize; use zkwasm_host_circuits::host::ForeignInst::PoseidonNew; use zkwasm_host_circuits::host::ForeignInst::PoseidonPush; @@ -58,6 +63,7 @@ pub struct PoseidonContext { pub generator: Generator, pub buf: Vec, pub fieldreducer: Reduce, + pub used_round: usize, } impl PoseidonContext { @@ -70,6 +76,7 @@ impl PoseidonContext { cursor: 0, values: vec![], }, + used_round: 0, } } @@ -77,6 +84,7 @@ impl PoseidonContext { self.buf = vec![]; if new != 0 { self.hasher = Some(POSEIDON_HASHER.clone()); + self.used_round += 1; } } @@ -105,13 +113,20 @@ impl PoseidonContext { } } -impl ForeignContext for PoseidonContext {} +impl ForeignContext for PoseidonContext { + fn get_statics(&self) -> Option { + Some(ForeignStatics { + used_round: self.used_round, + max_round: PoseidonChip::max_rounds(zkwasm_k() as usize), + }) + } +} use specs::external_host_call_table::ExternalHostCallSignature; pub fn register_poseidon_foreign(env: &mut HostEnv) { let foreign_poseidon_plugin = env .external_env - .register_plugin("foreign_sh256", Box::new(PoseidonContext::default())); + .register_plugin("foreign_poseidon", Box::new(PoseidonContext::default())); env.external_env.register_function( "poseidon_new", @@ -119,7 +134,7 @@ pub fn register_poseidon_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_poseidon_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); log::debug!("buf len is {}", context.buf.len()); context.poseidon_new(args.nth::(0) as usize); @@ -134,7 +149,7 @@ pub fn register_poseidon_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_poseidon_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.poseidon_push(args.nth::(0) as u64); None @@ -148,7 +163,7 @@ pub fn register_poseidon_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Return, foreign_poseidon_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); Some(wasmi::RuntimeValue::I64(context.poseidon_finalize() as i64)) }, diff --git a/crates/host/src/host/hash_helper/sha256.rs b/crates/host/src/host/hash_helper/sha256.rs index 4fe54d886..6925fe8f9 100644 --- a/crates/host/src/host/hash_helper/sha256.rs +++ b/crates/host/src/host/hash_helper/sha256.rs @@ -1,7 +1,9 @@ use delphinus_zkwasm::runtime::host::host_env::HostEnv; use delphinus_zkwasm::runtime::host::ForeignContext; +use delphinus_zkwasm::runtime::host::ForeignStatics; use sha2::Digest; use std::rc::Rc; +use wasmi::tracer::Observer; use zkwasm_host_circuits::host::ForeignInst::SHA256Finalize; use zkwasm_host_circuits::host::ForeignInst::SHA256New; use zkwasm_host_circuits::host::ForeignInst::SHA256Push; @@ -66,7 +68,12 @@ impl Sha256Context { } } -impl ForeignContext for Sha256Context {} +impl ForeignContext for Sha256Context { + fn get_statics(&self) -> Option { + // we did not support full sha256 as host yet + None + } +} use specs::external_host_call_table::ExternalHostCallSignature; pub fn register_sha256_foreign(env: &mut HostEnv) { @@ -80,7 +87,7 @@ pub fn register_sha256_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_sha256_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); let hasher = context .hasher @@ -101,7 +108,7 @@ pub fn register_sha256_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Argument, foreign_sha256_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.hasher.as_mut().map(|s| { let sz = if context.size > 8 { @@ -127,7 +134,7 @@ pub fn register_sha256_foreign(env: &mut HostEnv) { ExternalHostCallSignature::Return, foreign_sha256_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.hasher.as_ref().map(|s| { let dwords: Vec = s.clone().finalize()[..].to_vec(); diff --git a/crates/host/src/host/merkle_helper/datacache.rs b/crates/host/src/host/merkle_helper/datacache.rs new file mode 100644 index 000000000..aea78d3b8 --- /dev/null +++ b/crates/host/src/host/merkle_helper/datacache.rs @@ -0,0 +1,176 @@ +use delphinus_zkwasm::runtime::host::host_env::HostEnv; +use delphinus_zkwasm::runtime::host::ForeignContext; +use delphinus_zkwasm::runtime::host::ForeignStatics; +use halo2_proofs::pairing::bn256::Fr; +use std::cell::RefCell; +use std::rc::Rc; +use wasmi::tracer::Observer; +use zkwasm_host_circuits::host::datahash as datahelper; +use zkwasm_host_circuits::host::datahash::DataHashRecord; +use zkwasm_host_circuits::host::db::TreeDB; +use zkwasm_host_circuits::host::ForeignInst::CacheFetchData; +use zkwasm_host_circuits::host::ForeignInst::CacheSetHash; +use zkwasm_host_circuits::host::ForeignInst::CacheSetMode; +use zkwasm_host_circuits::host::ForeignInst::CacheStoreData; +use zkwasm_host_circuits::host::Reduce; +use zkwasm_host_circuits::host::ReduceRule; + +const FETCH_MODE: u64 = 0; +const STORE_MODE: u64 = 1; + +pub struct CacheContext { + pub mode: u64, + pub hash: Reduce, + pub data: Vec, + pub fetch: bool, + pub mongo_datahash: datahelper::MongoDataHash, + pub tree_db: Option>>, +} + +fn new_reduce(rules: Vec>) -> Reduce { + Reduce { cursor: 0, rules } +} + +impl CacheContext { + pub fn new(tree_db: Option>>) -> Self { + CacheContext { + mode: 0, + hash: new_reduce(vec![ReduceRule::Bytes(vec![], 4)]), + fetch: false, + data: vec![], + mongo_datahash: datahelper::MongoDataHash::construct([0; 32], tree_db.clone()), + tree_db, + } + } + + pub fn set_mode(&mut self, v: u64) { + self.mode = v; + self.data = vec![]; + } + + pub fn set_data_hash(&mut self, v: u64) { + self.hash.reduce(v); + if self.hash.cursor == 0 { + let hash: [u8; 32] = self.hash.rules[0] + .bytes_value() + .unwrap() + .try_into() + .unwrap(); + if self.mode == FETCH_MODE { + let datahashrecord = self.mongo_datahash.get_record(&hash).unwrap(); + self.data = datahashrecord.map_or(vec![], |r| { + r.data + .chunks_exact(8) + .into_iter() + .into_iter() + .map(|x| u64::from_le_bytes(x.try_into().unwrap())) + .collect::>() + }); + self.fetch = false; + } else if self.mode == STORE_MODE { + // put data and hash into mongo_datahash + if !self.data.is_empty() { + self.mongo_datahash + .update_record({ + DataHashRecord { + hash, + data: self + .data + .iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect::>(), + } + }) + .unwrap(); + } + } + } + } + + pub fn fetch_data(&mut self) -> u64 { + if self.fetch == false { + self.fetch = true; + self.data.reverse(); + self.data.len() as u64 + } else { + self.data.pop().unwrap() + } + } + + pub fn store_data(&mut self, v: u64) { + self.data.push(v); + } +} + +impl CacheContext {} + +impl ForeignContext for CacheContext { + fn get_statics(&self) -> Option { + // pure witness function + None + } +} + +use specs::external_host_call_table::ExternalHostCallSignature; +pub fn register_datacache_foreign(env: &mut HostEnv, tree_db: Option>>) { + let foreign_merkle_plugin = env + .external_env + .register_plugin("foreign_cache", Box::new(CacheContext::new(tree_db))); + + env.external_env.register_function( + "cache_set_mode", + CacheSetMode as usize, + ExternalHostCallSignature::Argument, + foreign_merkle_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + context.set_mode(args.nth(0)); + None + }, + ), + ); + + env.external_env.register_function( + "cache_set_hash", + CacheSetHash as usize, + ExternalHostCallSignature::Argument, + foreign_merkle_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + context.set_data_hash(args.nth(0)); + None + }, + ), + ); + + env.external_env.register_function( + "cache_store_data", + CacheStoreData as usize, + ExternalHostCallSignature::Argument, + foreign_merkle_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + context.store_data(args.nth(0)); + None + }, + ), + ); + + env.external_env.register_function( + "cache_fetch_data", + CacheFetchData as usize, + ExternalHostCallSignature::Return, + foreign_merkle_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + let ret = Some(wasmi::RuntimeValue::I64(context.fetch_data() as i64)); + ret + }, + ), + ); +} diff --git a/crates/host/src/host/merkle_helper/merkle.rs b/crates/host/src/host/merkle_helper/merkle.rs index cd075c41c..ebfe3a5ac 100644 --- a/crates/host/src/host/merkle_helper/merkle.rs +++ b/crates/host/src/host/merkle_helper/merkle.rs @@ -1,18 +1,20 @@ +use delphinus_zkwasm::circuits::config::zkwasm_k; use delphinus_zkwasm::runtime::host::host_env::HostEnv; use delphinus_zkwasm::runtime::host::ForeignContext; +use delphinus_zkwasm::runtime::host::ForeignStatics; use halo2_proofs::pairing::bn256::Fr; use std::cell::RefCell; use std::rc::Rc; +use wasmi::tracer::Observer; +use zkwasm_host_circuits::circuits::host::HostOpSelector; +use zkwasm_host_circuits::circuits::merkle::MerkleChip; use zkwasm_host_circuits::host::datahash as datahelper; -use zkwasm_host_circuits::host::datahash::DataHashRecord; use zkwasm_host_circuits::host::db::TreeDB; use zkwasm_host_circuits::host::merkle::MerkleTree; use zkwasm_host_circuits::host::mongomerkle as merklehelper; use zkwasm_host_circuits::host::ForeignInst::MerkleAddress; -use zkwasm_host_circuits::host::ForeignInst::MerkleFetchData; use zkwasm_host_circuits::host::ForeignInst::MerkleGet; use zkwasm_host_circuits::host::ForeignInst::MerkleGetRoot; -use zkwasm_host_circuits::host::ForeignInst::MerklePutData; use zkwasm_host_circuits::host::ForeignInst::MerkleSet; use zkwasm_host_circuits::host::ForeignInst::MerkleSetRoot; use zkwasm_host_circuits::host::Reduce; @@ -25,13 +27,13 @@ pub struct MerkleContext { pub get_root: Reduce, pub address: Reduce, pub set: Reduce, - pub get: Reduce, - pub data: Vec, + pub data: [u64; 4], pub data_cursor: usize, pub fetch: bool, pub mongo_merkle: Option>, pub mongo_datahash: datahelper::MongoDataHash, pub tree_db: Option>>, + pub used_round: usize, } fn new_reduce(rules: Vec>) -> Reduce { @@ -45,18 +47,13 @@ impl MerkleContext { get_root: new_reduce(vec![ReduceRule::Bytes(vec![], 4)]), address: new_reduce(vec![ReduceRule::U64(0)]), set: new_reduce(vec![ReduceRule::Bytes(vec![], 4)]), - get: new_reduce(vec![ - ReduceRule::U64(0), - ReduceRule::U64(0), - ReduceRule::U64(0), - ReduceRule::U64(0), - ]), fetch: false, - data: vec![], + data: [0; 4], data_cursor: 0, mongo_merkle: None, mongo_datahash: datahelper::MongoDataHash::construct([0; 32], tree_db.clone()), tree_db, + used_round: 0, } } @@ -92,9 +89,14 @@ impl MerkleContext { values[cursor] } + /// reset the address of merkle op together with the data and data_cursor pub fn merkle_address(&mut self, v: u64) { - self.data = vec![]; + if self.address.cursor == 0 { + self.used_round += 1; + } + self.data = [0; 4]; self.fetch = false; + self.data_cursor = 0; self.address.reduce(v); } @@ -110,22 +112,6 @@ impl MerkleContext { let hash = self.set.rules[0].bytes_value().unwrap(); mt.update_leaf_data_with_proof(index, &hash) .expect("Unexpected failure: update leaf with proof fail"); - // put data and hash into mongo_datahash if the data is binded to the merkle tree leaf - if !self.data.is_empty() { - self.mongo_datahash - .update_record({ - DataHashRecord { - hash: hash.try_into().unwrap(), - data: self - .data - .iter() - .map(|x| x.to_le_bytes()) - .flatten() - .collect::>(), - } - }) - .unwrap(); - } } } @@ -139,69 +125,26 @@ impl MerkleContext { let (leaf, _) = mt .get_leaf_with_proof(index) .expect("Unexpected failure: get leaf fail"); - let cursor = self.get.cursor; let values = leaf.data_as_u64(); - self.get.reduce(values[self.get.cursor]); - // fetch data if we get the target hash - if self.get.cursor == 0 { - let hash: [u8; 32] = vec![ - self.get.rules[0] - .u64_value() - .unwrap() - .to_le_bytes() - .to_vec(), - self.get.rules[1] - .u64_value() - .unwrap() - .to_le_bytes() - .to_vec(), - self.get.rules[2] - .u64_value() - .unwrap() - .to_le_bytes() - .to_vec(), - self.get.rules[3] - .u64_value() - .unwrap() - .to_le_bytes() - .to_vec(), - ] - .into_iter() - .flatten() - .collect::>() - .try_into() - .unwrap(); - let datahashrecord = self.mongo_datahash.get_record(&hash).unwrap(); - self.data = datahashrecord.map_or(vec![], |r| { - r.data - .chunks_exact(8) - .into_iter() - .into_iter() - .map(|x| u64::from_le_bytes(x.try_into().unwrap())) - .collect::>() - }); + if self.data_cursor == 0 { + self.data = values; } - values[cursor] - } - - pub fn merkle_fetch_data(&mut self) -> u64 { - if self.fetch == false { - self.fetch = true; - self.data.reverse(); - self.data.len() as u64 - } else { - self.data.pop().unwrap() - } - } - - pub fn merkle_put_data(&mut self, v: u64) { - self.data.push(v); + let v = values[self.data_cursor]; + self.data_cursor += 1; + return v; } } impl MerkleContext {} -impl ForeignContext for MerkleContext {} +impl ForeignContext for MerkleContext { + fn get_statics(&self) -> Option { + Some(ForeignStatics { + used_round: self.used_round, + max_round: MerkleChip::::max_rounds(zkwasm_k() as usize), + }) + } +} use specs::external_host_call_table::ExternalHostCallSignature; pub fn register_merkle_foreign(env: &mut HostEnv, tree_db: Option>>) { @@ -215,7 +158,7 @@ pub fn register_merkle_foreign(env: &mut HostEnv, tree_db: Option().unwrap(); context.merkle_setroot(args.nth(0)); None @@ -229,33 +172,20 @@ pub fn register_merkle_foreign(env: &mut HostEnv, tree_db: Option().unwrap(); Some(wasmi::RuntimeValue::I64(context.merkle_getroot() as i64)) }, ), ); - env.external_env.register_function( - "merkle_fetch_data", - MerkleFetchData as usize, - ExternalHostCallSignature::Return, - foreign_merkle_plugin.clone(), - Rc::new( - |context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { - let context = context.downcast_mut::().unwrap(); - Some(wasmi::RuntimeValue::I64(context.merkle_fetch_data() as i64)) - }, - ), - ); - env.external_env.register_function( "merkle_address", MerkleAddress as usize, ExternalHostCallSignature::Argument, foreign_merkle_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.merkle_address(args.nth(0)); None @@ -263,27 +193,13 @@ pub fn register_merkle_foreign(env: &mut HostEnv, tree_db: Option().unwrap(); - context.merkle_put_data(args.nth(0)); - None - }, - ), - ); - env.external_env.register_function( "merkle_set", MerkleSet as usize, ExternalHostCallSignature::Argument, foreign_merkle_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); context.merkle_set(args.nth(0)); None @@ -297,7 +213,7 @@ pub fn register_merkle_foreign(env: &mut HostEnv, tree_db: Option().unwrap(); let ret = Some(wasmi::RuntimeValue::I64(context.merkle_get() as i64)); ret diff --git a/crates/host/src/host/merkle_helper/mod.rs b/crates/host/src/host/merkle_helper/mod.rs index bdf9eb8f3..d6b71bb10 100644 --- a/crates/host/src/host/merkle_helper/mod.rs +++ b/crates/host/src/host/merkle_helper/mod.rs @@ -1 +1,2 @@ +pub mod datacache; pub mod merkle; diff --git a/crates/host/src/host/mod.rs b/crates/host/src/host/mod.rs index 9b9534c06..49254c63e 100644 --- a/crates/host/src/host/mod.rs +++ b/crates/host/src/host/mod.rs @@ -1,3 +1,4 @@ pub mod ecc_helper; pub mod hash_helper; pub mod merkle_helper; +pub mod witness_helper; diff --git a/crates/host/src/host/witness_helper/mod.rs b/crates/host/src/host/witness_helper/mod.rs new file mode 100644 index 000000000..61487aa1d --- /dev/null +++ b/crates/host/src/host/witness_helper/mod.rs @@ -0,0 +1,176 @@ +use delphinus_zkwasm::runtime::host::ForeignContext; +use std::cell::RefCell; +use std::collections::HashMap; +use std::rc::Rc; +use wasmi::tracer::Observer; + +use crate::HostEnv; +use zkwasm_host_circuits::host::ForeignInst::WitnessIndexedInsert; +use zkwasm_host_circuits::host::ForeignInst::WitnessIndexedPop; +use zkwasm_host_circuits::host::ForeignInst::WitnessIndexedPush; +use zkwasm_host_circuits::host::ForeignInst::WitnessInsert; +use zkwasm_host_circuits::host::ForeignInst::WitnessPop; +use zkwasm_host_circuits::host::ForeignInst::WitnessSetIndex; +use zkwasm_host_circuits::host::ForeignInst::WitnessTraceSize; + +#[derive(Default)] +pub struct WitnessContext { + pub buf: Vec, + pub indexed_buf: Rc>>>, + pub focus: u64, +} + +impl WitnessContext { + fn new(indexed_map: Rc>>>) -> Self { + WitnessContext { + buf: vec![], + indexed_buf: indexed_map, + focus: 0, + } + } +} + +impl WitnessContext { + pub fn witness_insert(&mut self, new: u64) { + self.buf.insert(0, new); + } + + pub fn witness_pop(&mut self) -> u64 { + self.buf.pop().unwrap() + } + + pub fn witness_set_index(&mut self, index: u64) { + self.focus = index; + } + + pub fn witness_indexed_insert(&mut self, new: u64) { + let mut bind = self.indexed_buf.borrow_mut(); + let buf = bind.get_mut(&self.focus); + if let Some(vec) = buf { + vec.insert(0, new); + } else { + bind.insert(self.focus, vec![new]); + } + } + + pub fn witness_indexed_push(&mut self, new: u64) { + let mut bind = self.indexed_buf.borrow_mut(); + let buf = bind.get_mut(&self.focus); + if let Some(vec) = buf { + vec.push(new); + } else { + bind.insert(self.focus, vec![new]); + } + } + + pub fn witness_indexed_pop(&mut self) -> u64 { + let mut bind = self.indexed_buf.borrow_mut(); + let buf = bind.get_mut(&self.focus).unwrap(); + buf.pop().unwrap() + } +} + +impl ForeignContext for WitnessContext {} + +use specs::external_host_call_table::ExternalHostCallSignature; +pub fn register_witness_foreign(env: &mut HostEnv, index_map: Rc>>>) { + let foreign_witness_plugin = env + .external_env + .register_plugin("foreign_witness", Box::new(WitnessContext::new(index_map))); + + env.external_env.register_function( + "wasm_witness_insert", + WitnessInsert as usize, + ExternalHostCallSignature::Argument, + foreign_witness_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + context.witness_insert(args.nth::(0) as u64); + None + }, + ), + ); + + env.external_env.register_function( + "wasm_witness_set_index", + WitnessSetIndex as usize, + ExternalHostCallSignature::Argument, + foreign_witness_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + context.witness_set_index(args.nth::(0) as u64); + None + }, + ), + ); + + env.external_env.register_function( + "wasm_witness_indexed_insert", + WitnessIndexedInsert as usize, + ExternalHostCallSignature::Argument, + foreign_witness_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + context.witness_indexed_insert(args.nth::(0) as u64); + None + }, + ), + ); + + env.external_env.register_function( + "wasm_witness_indexed_push", + WitnessIndexedPush as usize, + ExternalHostCallSignature::Argument, + foreign_witness_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + context.witness_indexed_push(args.nth::(0) as u64); + None + }, + ), + ); + + env.external_env.register_function( + "wasm_witness_pop", + WitnessPop as usize, + ExternalHostCallSignature::Return, + foreign_witness_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + Some(wasmi::RuntimeValue::I64(context.witness_pop() as i64)) + }, + ), + ); + + env.external_env.register_function( + "wasm_witness_indexed_pop", + WitnessIndexedPop as usize, + ExternalHostCallSignature::Return, + foreign_witness_plugin.clone(), + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); + Some(wasmi::RuntimeValue::I64( + context.witness_indexed_pop() as i64 + )) + }, + ), + ); + + env.external_env.register_function( + "wasm_trace_size", + WitnessTraceSize as usize, + ExternalHostCallSignature::Return, + foreign_witness_plugin.clone(), + Rc::new( + |obs: &Observer, _context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + Some(wasmi::RuntimeValue::I64(obs.counter as i64)) + }, + ), + ); +} diff --git a/crates/host/src/lib.rs b/crates/host/src/lib.rs index 735f47f79..dc5ea6db7 100644 --- a/crates/host/src/lib.rs +++ b/crates/host/src/lib.rs @@ -9,10 +9,15 @@ use delphinus_zkwasm::foreign::wasm_input_helper::runtime::register_wasm_input_f use delphinus_zkwasm::runtime::wasmi_interpreter::WasmRuntimeIO; use delphinus_zkwasm::runtime::host::host_env::HostEnv; +use delphinus_zkwasm::runtime::host::ContextOutput; use delphinus_zkwasm::runtime::host::HostEnvBuilder; +use serde::Deserialize; +use serde::Serialize; +use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; use zkwasm_host_circuits::host::db::TreeDB; +use zkwasm_host_circuits::proof::OpType; pub struct ExecutionArg { /// Public inputs for `wasm_input(1)` @@ -23,43 +28,92 @@ pub struct ExecutionArg { pub context_inputs: Vec, /// Context outputs for `wasm_write_context()` pub context_outputs: Arc>>, + /// indexed witness context + pub indexed_witness: Rc>>>, /// db src pub tree_db: Option>>, } +impl ContextOutput for ExecutionArg { + fn get_context_outputs(&self) -> Arc>> { + self.context_outputs.clone() + } +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct HostEnvConfig { + pub ops: Vec, +} + +impl Default for HostEnvConfig { + fn default() -> Self { + HostEnvConfig { + ops: vec![ + OpType::POSEIDONHASH, + OpType::MERKLE, + OpType::JUBJUBSUM, + OpType::KECCAKHASH, + OpType::BN256SUM, + ], + } + } +} + +impl HostEnvConfig { + fn register_op(op: &OpType, env: &mut HostEnv) { + match op { + OpType::BLS381PAIR => host::ecc_helper::bls381::pair::register_blspair_foreign(env), + OpType::BLS381SUM => host::ecc_helper::bls381::sum::register_blssum_foreign(env), + OpType::BN256PAIR => host::ecc_helper::bn254::pair::register_bn254pair_foreign(env), + OpType::BN256SUM => host::ecc_helper::bn254::sum::register_bn254sum_foreign(env), + OpType::POSEIDONHASH => host::hash_helper::poseidon::register_poseidon_foreign(env), + OpType::MERKLE => { + host::merkle_helper::merkle::register_merkle_foreign(env, None); + host::merkle_helper::datacache::register_datacache_foreign(env, None); + } + OpType::JUBJUBSUM => host::ecc_helper::jubjub::sum::register_babyjubjubsum_foreign(env), + OpType::KECCAKHASH => host::hash_helper::keccak256::register_keccak_foreign(env), + } + } + + fn register_ops(&self, env: &mut HostEnv) { + for op in &self.ops { + Self::register_op(op, env); + } + } +} + pub struct StandardHostEnvBuilder; impl HostEnvBuilder for StandardHostEnvBuilder { type Arg = ExecutionArg; + type HostConfig = HostEnvConfig; - fn create_env_without_value() -> (HostEnv, WasmRuntimeIO) { + fn create_env_without_value(envconfig: Self::HostConfig) -> (HostEnv, WasmRuntimeIO) { let mut env = HostEnv::new(); let wasm_runtime_io = register_wasm_input_foreign(&mut env, vec![], vec![]); register_require_foreign(&mut env); register_log_foreign(&mut env); register_context_foreign(&mut env, vec![], Arc::new(Mutex::new(vec![]))); - host::hash_helper::poseidon::register_poseidon_foreign(&mut env); - host::merkle_helper::merkle::register_merkle_foreign(&mut env, None); - host::ecc_helper::bn254::sum::register_bn254sum_foreign(&mut env); - host::ecc_helper::bn254::pair::register_bn254pair_foreign(&mut env); - host::ecc_helper::jubjub::sum::register_babyjubjubsum_foreign(&mut env); + envconfig.register_ops(&mut env); + host::witness_helper::register_witness_foreign( + &mut env, + Rc::new(RefCell::new(HashMap::new())), + ); env.finalize(); (env, wasm_runtime_io) } - fn create_env(arg: Self::Arg) -> (HostEnv, WasmRuntimeIO) { + fn create_env(arg: Self::Arg, envconfig: Self::HostConfig) -> (HostEnv, WasmRuntimeIO) { let mut env = HostEnv::new(); let wasm_runtime_io = register_wasm_input_foreign(&mut env, arg.public_inputs, arg.private_inputs); register_require_foreign(&mut env); register_log_foreign(&mut env); register_context_foreign(&mut env, arg.context_inputs, arg.context_outputs); - host::hash_helper::poseidon::register_poseidon_foreign(&mut env); - host::merkle_helper::merkle::register_merkle_foreign(&mut env, arg.tree_db); - host::ecc_helper::bn254::sum::register_bn254sum_foreign(&mut env); - host::ecc_helper::bn254::pair::register_bn254pair_foreign(&mut env); - host::ecc_helper::jubjub::sum::register_babyjubjubsum_foreign(&mut env); + host::witness_helper::register_witness_foreign(&mut env, arg.indexed_witness); + envconfig.register_ops(&mut env); env.finalize(); (env, wasm_runtime_io) diff --git a/crates/playground/Cargo.lock b/crates/playground/Cargo.lock index 08892f8e2..88716ebba 100644 --- a/crates/playground/Cargo.lock +++ b/crates/playground/Cargo.lock @@ -1228,6 +1228,15 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.9" @@ -1473,6 +1482,7 @@ dependencies = [ "num-integer", "num-traits", "rand", + "serde", ] [[package]] @@ -3083,21 +3093,24 @@ dependencies = [ [[package]] name = "zkwasm-host-circuits" version = "0.1.0" -source = "git+https://github.com/DelphinusLab/zkWasm-host-circuits.git?branch=main#997af6f26acd12be6ec988c7a14be29a5f1039c4" +source = "git+https://github.com/DelphinusLab/zkWasm-host-circuits.git?branch=main#61c52e13667291b360783139e85aa499458eb06d" dependencies = [ "ark-std", + "cfg-if 1.0.0", "circuits-batcher", "clap", "ff", "halo2_proofs", "halo2ecc-s", "hex", + "itertools", "lazy_static", "lru", "mongodb", "num-bigint", "poseidon", "rand", + "rand_core", "ripemd", "serde", "serde_json", diff --git a/crates/playground/examples/binary_search.rs b/crates/playground/examples/binary_search.rs index ceb628073..6d8bfe75d 100644 --- a/crates/playground/examples/binary_search.rs +++ b/crates/playground/examples/binary_search.rs @@ -11,12 +11,14 @@ fn main() -> Result<()> { let wasm = std::fs::read("wasm/binary_search.wasm")?; let loader = ZkWasmLoader::::new(18, wasm, vec![])?; - - let (circuit, instances, _) = loader.circuit_with_witness(ExecutionArg { + let result = loader.run(ExecutionArg { public_inputs: vec![0], private_inputs: vec![], context_inputs: vec![], context_outputs: Arc::new(Mutex::new(vec![])), - })?; + }, (), false, true)?; + + + let (circuit, instances) = loader.circuit_with_witness(result)?; loader.mock_test(&circuit, &instances) } diff --git a/crates/playground/examples/context.rs b/crates/playground/examples/context.rs index 69d7ed0b5..f9589fa9b 100644 --- a/crates/playground/examples/context.rs +++ b/crates/playground/examples/context.rs @@ -21,7 +21,8 @@ fn main() -> Result<()> { context_outputs: context_outputs.clone(), }; - let (circuit, instances, _) = loader.circuit_with_witness(arg)?; + let result = loader.run(arg, (), false, true)?; + let (circuit, instances) = loader.circuit_with_witness(result)?; loader.mock_test(&circuit, &instances)?; let arg = ExecutionArg { @@ -31,7 +32,9 @@ fn main() -> Result<()> { context_outputs: Arc::new(Mutex::new(vec![])), }; - let (circuit, instances, _) = loader.circuit_with_witness(arg)?; + let result = loader.run(arg, (), false, true)?; + + let (circuit, instances) = loader.circuit_with_witness(result)?; loader.mock_test(&circuit, &instances)?; Ok(()) diff --git a/crates/playground/examples/fibonacci.rs b/crates/playground/examples/fibonacci.rs index f46e2adab..6e4733839 100644 --- a/crates/playground/examples/fibonacci.rs +++ b/crates/playground/examples/fibonacci.rs @@ -11,12 +11,14 @@ fn main() -> Result<()> { let wasm = std::fs::read("wasm/fibonacci.wasm")?; let loader = ZkWasmLoader::::new(18, wasm, vec![])?; - - let (circuit, instances, _) = loader.circuit_with_witness(ExecutionArg { + let result = loader.run(ExecutionArg { public_inputs: vec![5], private_inputs: vec![], context_inputs: vec![], context_outputs: Arc::new(Mutex::new(vec![])), - })?; + }, (), false, true)?; + + + let (circuit, instances) = loader.circuit_with_witness(result)?; loader.mock_test(&circuit, &instances) } diff --git a/crates/playground/examples/phantom.rs b/crates/playground/examples/phantom.rs index a8f62be0c..5ce2f5603 100644 --- a/crates/playground/examples/phantom.rs +++ b/crates/playground/examples/phantom.rs @@ -16,11 +16,12 @@ fn main() -> Result<()> { vec!["search".to_owned()], )?; - let (circuit, instances, _) = loader.circuit_with_witness(ExecutionArg { + let result = loader.run(ExecutionArg { public_inputs: vec![2], private_inputs: vec![], context_inputs: vec![], context_outputs: Arc::new(Mutex::new(vec![])), - })?; + }, (), false, true)?; + let (circuit, instances) = loader.circuit_with_witness(result)?; loader.mock_test(&circuit, &instances) } diff --git a/crates/zkwasm/Cargo.toml b/crates/zkwasm/Cargo.toml index 194f38123..2f79adb85 100644 --- a/crates/zkwasm/Cargo.toml +++ b/crates/zkwasm/Cargo.toml @@ -43,4 +43,4 @@ rusty-fork = "0.3.0" [features] default = [] cuda = ["halo2_proofs/cuda", "specs/cuda"] -uniform-circuit = [] \ No newline at end of file +uniform-circuit = [] diff --git a/crates/zkwasm/src/foreign/context/runtime.rs b/crates/zkwasm/src/foreign/context/runtime.rs index 1a7ff2751..829b61b2c 100644 --- a/crates/zkwasm/src/foreign/context/runtime.rs +++ b/crates/zkwasm/src/foreign/context/runtime.rs @@ -5,6 +5,7 @@ use std::sync::Mutex; use specs::host_function::HostPlugin; use specs::host_function::Signature; use specs::types::ValueType; +use wasmi::tracer::Observer; use wasmi::RuntimeArgs; use crate::runtime::host::host_env::HostEnv; @@ -47,6 +48,7 @@ pub fn register_context_foreign( context_output: Arc>>, ) { env.internal_env.register_plugin( + "context plugin", HostPlugin::Context, Box::new(Context::new(context_input, context_output)), ); @@ -59,11 +61,13 @@ pub fn register_context_foreign( }, HostPlugin::Context, Op::ReadContext as usize, - Rc::new(|context: &mut dyn ForeignContext, _args: RuntimeArgs| { - let context = context.downcast_mut::().unwrap(); + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, _args: RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); - Some(wasmi::RuntimeValue::I64(context.read_context() as i64)) - }), + Some(wasmi::RuntimeValue::I64(context.read_context() as i64)) + }, + ), ); env.internal_env.register_function( @@ -74,13 +78,15 @@ pub fn register_context_foreign( }, HostPlugin::Context, Op::WriteContext as usize, - Rc::new(|context: &mut dyn ForeignContext, args: RuntimeArgs| { - let context = context.downcast_mut::().unwrap(); + Rc::new( + |_obs: &Observer, context: &mut dyn ForeignContext, args: RuntimeArgs| { + let context = context.downcast_mut::().unwrap(); - let value: i64 = args.nth(0); - context.write_context(value as u64); + let value: i64 = args.nth(0); + context.write_context(value as u64); - None - }), + None + }, + ), ); } diff --git a/crates/zkwasm/src/foreign/log_helper/mod.rs b/crates/zkwasm/src/foreign/log_helper/mod.rs index adc62834c..34bdd155f 100644 --- a/crates/zkwasm/src/foreign/log_helper/mod.rs +++ b/crates/zkwasm/src/foreign/log_helper/mod.rs @@ -1,10 +1,12 @@ use std::rc::Rc; use specs::external_host_call_table::ExternalHostCallSignature; +use wasmi::tracer::Observer; use crate::runtime::host::host_env::HostEnv; use crate::runtime::host::ForeignContext; use zkwasm_host_circuits::host::ForeignInst::Log; +use zkwasm_host_circuits::host::ForeignInst::LogChar; struct Context; impl ForeignContext for Context {} @@ -15,11 +17,17 @@ pub fn register_log_foreign(env: &mut HostEnv) { .register_plugin("foreign_print", Box::new(Context)); let print = Rc::new( - |_context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_observer: &Observer, _context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let value: u64 = args.nth(0); - println!("{}", value); + None + }, + ); + let printchar = Rc::new( + |_observer: &Observer, _context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + let value: u64 = args.nth(0); + print!("{}", value as u8 as char); None }, ); @@ -28,7 +36,15 @@ pub fn register_log_foreign(env: &mut HostEnv) { "wasm_dbg", Log as usize, ExternalHostCallSignature::Argument, - foreign_log_plugin, + foreign_log_plugin.clone(), print, ); + + env.external_env.register_function( + "wasm_dbg_char", + LogChar as usize, + ExternalHostCallSignature::Argument, + foreign_log_plugin, + printchar, + ); } diff --git a/crates/zkwasm/src/foreign/require_helper/mod.rs b/crates/zkwasm/src/foreign/require_helper/mod.rs index 9752269c0..cdfc426ac 100644 --- a/crates/zkwasm/src/foreign/require_helper/mod.rs +++ b/crates/zkwasm/src/foreign/require_helper/mod.rs @@ -2,6 +2,7 @@ use std::rc::Rc; use specs::host_function::HostPlugin; use specs::types::ValueType; +use wasmi::tracer::Observer; use crate::runtime::host::host_env::HostEnv; use crate::runtime::host::ForeignContext; @@ -13,7 +14,7 @@ impl ForeignContext for Context {} pub fn register_require_foreign(env: &mut HostEnv) { let require = Rc::new( - |_context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_observer: &Observer, _context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let cond: u32 = args.nth(0); if cond == 0 { @@ -29,7 +30,7 @@ pub fn register_require_foreign(env: &mut HostEnv) { ); env.internal_env - .register_plugin(HostPlugin::Require, Box::new(Context)); + .register_plugin("require plugin", HostPlugin::Require, Box::new(Context)); env.internal_env.register_function( "require", diff --git a/crates/zkwasm/src/foreign/wasm_input_helper/runtime.rs b/crates/zkwasm/src/foreign/wasm_input_helper/runtime.rs index d170cbfa1..c740c7b49 100644 --- a/crates/zkwasm/src/foreign/wasm_input_helper/runtime.rs +++ b/crates/zkwasm/src/foreign/wasm_input_helper/runtime.rs @@ -4,9 +4,11 @@ use std::rc::Rc; use specs::host_function::HostPlugin; use specs::types::ValueType; +use wasmi::tracer::Observer; use crate::runtime::host::host_env::HostEnv; use crate::runtime::host::ForeignContext; +use crate::runtime::host::ForeignStatics; use crate::runtime::wasmi_interpreter::WasmRuntimeIO; use super::Op; @@ -79,7 +81,11 @@ impl Context { } } -impl ForeignContext for Context {} +impl ForeignContext for Context { + fn get_statics(&self) -> Option { + None + } +} pub fn register_wasm_input_foreign( env: &mut HostEnv, @@ -90,9 +96,8 @@ pub fn register_wasm_input_foreign( let outputs = Rc::new(RefCell::new(vec![])); let wasm_input = Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_observer: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); - let arg: i32 = args.nth(0); let input = context.wasm_input(arg); @@ -101,7 +106,7 @@ pub fn register_wasm_input_foreign( ); let wasm_output = Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_observer: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); let value: i64 = args.nth(0); @@ -112,6 +117,7 @@ pub fn register_wasm_input_foreign( ); env.internal_env.register_plugin( + "wasm input plugin", HostPlugin::HostInput, Box::new(Context::new( public_inputs, diff --git a/crates/zkwasm/src/loader/mod.rs b/crates/zkwasm/src/loader/mod.rs index e84364b06..953a9a089 100644 --- a/crates/zkwasm/src/loader/mod.rs +++ b/crates/zkwasm/src/loader/mod.rs @@ -101,8 +101,11 @@ impl> ZkWasmLoader Result> { - let (env, wasm_runtime_io) = EnvBuilder::create_env_without_value(); + fn circuit_without_witness( + &self, + envconfig: EnvBuilder::HostConfig, + ) -> Result> { + let (env, wasm_runtime_io) = EnvBuilder::create_env_without_value(envconfig); let compiled_module = self.compile(&env, true)?; @@ -145,14 +148,22 @@ impl> ZkWasmLoader) -> Result> { - let circuit = self.circuit_without_witness()?; + pub fn create_vkey( + &self, + params: &Params, + envconfig: EnvBuilder::HostConfig, + ) -> Result> { + let circuit = self.circuit_without_witness(envconfig)?; Ok(keygen_vk(¶ms, &circuit).unwrap()) } - pub fn checksum(&self, params: &Params) -> Result> { - let (env, _) = EnvBuilder::create_env_without_value(); + pub fn checksum( + &self, + params: &Params, + envconfig: EnvBuilder::HostConfig, + ) -> Result> { + let (env, _) = EnvBuilder::create_env_without_value(envconfig); let compiled = self.compile(&env, true)?; let table_with_params = CompilationTableWithParams { @@ -168,14 +179,13 @@ impl> ZkWasmLoader Result> { - let (mut env, wasm_runtime_io) = EnvBuilder::create_env(arg); + let (env, wasm_runtime_io) = EnvBuilder::create_env(arg, config); let compiled_module = self.compile(&env, dryrun)?; - - let result = compiled_module.run(&mut env, dryrun, wasm_runtime_io)?; - + let result = compiled_module.run(env, dryrun, wasm_runtime_io)?; if !dryrun { result.tables.profile_tables(); @@ -189,9 +199,9 @@ impl> ZkWasmLoader Result<(TestCircuit, Vec, Vec)> { - let execution_result = self.run(arg, false, true)?; + execution_result: ExecutionResult, + ) -> Result<(TestCircuit, Vec)> { + //let execution_result = self.run(arg, config, false, true)?; let instance: Vec = execution_result .public_inputs_and_outputs .clone() @@ -201,13 +211,10 @@ impl> ZkWasmLoader> ZkWasmLoader, instances: Vec, proof: Vec, + #[cfg(feature = "uniform-circuit")] config: EnvBuilder::HostConfig, ) -> Result<()> { let params_verifier: ParamsVerifier = params.verifier(instances.len()).unwrap(); let strategy = SingleVerifier::new(¶ms_verifier); @@ -282,7 +290,7 @@ impl> ZkWasmLoader for ExecutionArg { - fn from(seq: super::Sequence) -> ExecutionArg { - let private_inputs = parse_args(seq.private_inputs.iter().map(|s| s.as_str()).collect()); - let public_inputs = parse_args(seq.public_inputs.iter().map(|s| s.as_str()).collect()); - let context_inputs = parse_args(seq.context_input.iter().map(|s| s.as_str()).collect()); - let context_outputs = Arc::new(Mutex::new(vec![])); - ExecutionArg { - private_inputs, - public_inputs, - context_inputs, - context_outputs, - } - } -} - pub struct DefaultHostEnvBuilder; impl HostEnvBuilder for DefaultHostEnvBuilder { type Arg = ExecutionArg; + type HostConfig = (); - fn create_env_without_value() -> (HostEnv, WasmRuntimeIO) { + fn create_env_without_value(_config: Self::HostConfig) -> (HostEnv, WasmRuntimeIO) { let mut env = HostEnv::new(); let wasm_runtime_io = register_wasm_input_foreign(&mut env, vec![], vec![]); register_require_foreign(&mut env); @@ -59,7 +44,7 @@ impl HostEnvBuilder for DefaultHostEnvBuilder { (env, wasm_runtime_io) } - fn create_env(arg: Self::Arg) -> (HostEnv, WasmRuntimeIO) { + fn create_env(arg: Self::Arg, _config: Self::HostConfig) -> (HostEnv, WasmRuntimeIO) { let mut env = HostEnv::new(); let wasm_runtime_io = register_wasm_input_foreign(&mut env, arg.public_inputs, arg.private_inputs); diff --git a/crates/zkwasm/src/runtime/host/external_circuit_plugin.rs b/crates/zkwasm/src/runtime/host/external_circuit_plugin.rs index 7eade3bb0..65b69c534 100644 --- a/crates/zkwasm/src/runtime/host/external_circuit_plugin.rs +++ b/crates/zkwasm/src/runtime/host/external_circuit_plugin.rs @@ -2,6 +2,7 @@ use specs::external_host_call_table::ExternalHostCallSignature; use std::cell::RefCell; use std::collections::HashMap; use std::rc::Rc; +use wasmi::tracer::Observer; use wasmi::FuncInstance; use wasmi::ModuleImportResolver; use wasmi::RuntimeArgs; @@ -9,13 +10,14 @@ use wasmi::RuntimeValue; use super::ForeignContext; use super::ForeignPlugin; +use super::ForeignStatics; use super::MatchForeignOpSignature; pub(super) struct ForeignOp { pub op_index: usize, pub sig: ExternalHostCallSignature, pub plugin: Rc, - pub cb: Rc Option>, + pub cb: Rc Option>, } pub struct ExternalCircuitEnv { @@ -34,10 +36,11 @@ impl ExternalCircuitEnv { /// Register a plugin without circuit pub fn register_plugin( &mut self, - _name: &str, + name: &str, ctx: Box, ) -> Rc { Rc::new(ForeignPlugin { + name: name.to_string(), ctx: Rc::new(RefCell::new(ctx)), }) } @@ -49,7 +52,7 @@ impl ExternalCircuitEnv { op_index: usize, sig: ExternalHostCallSignature, plugin: Rc, - cb: Rc Option>, + cb: Rc Option>, ) { assert!(!*self.finalized.borrow()); @@ -63,6 +66,20 @@ impl ExternalCircuitEnv { }, ); } + + pub fn get_statics(&self) -> HashMap { + let mut m = HashMap::new(); + for (_, v) in &self.functions { + let plugin_name = &v.plugin.name; + + if !m.contains_key(plugin_name) { + if let Some(stat) = (v.plugin.ctx).as_ref().borrow().get_statics() { + m.insert(plugin_name.to_string(), stat); + } + } + } + m + } } impl ModuleImportResolver for ExternalCircuitEnv { diff --git a/crates/zkwasm/src/runtime/host/host_env.rs b/crates/zkwasm/src/runtime/host/host_env.rs index cd23d0452..0a087b37b 100644 --- a/crates/zkwasm/src/runtime/host/host_env.rs +++ b/crates/zkwasm/src/runtime/host/host_env.rs @@ -6,6 +6,7 @@ use std::time::Instant; use log::debug; use specs::host_function::HostFunctionDesc; +use wasmi::tracer::Tracer; use wasmi::Externals; use wasmi::ModuleImportResolver; use wasmi::RuntimeArgs; @@ -148,13 +149,25 @@ impl ModuleImportResolver for HostEnv { } } -impl Externals for HostEnv { +pub struct ExecEnv { + pub host_env: HostEnv, + pub tracer: Rc>, +} + +impl Externals for ExecEnv { fn invoke_index( &mut self, index: usize, args: RuntimeArgs, ) -> Result, Trap> { - match self.cached_lookup.as_ref().unwrap().get(&index).clone() { + match self + .host_env + .cached_lookup + .as_ref() + .unwrap() + .get(&index) + .clone() + { Some(HostFunction { desc, execution_env: HostFunctionExecutionEnv { ctx, cb }, @@ -163,10 +176,11 @@ impl Externals for HostEnv { let ctx = ctx.as_mut(); let start = Instant::now(); - let r = cb(ctx, args); + let r = cb(&self.tracer.borrow().observer, ctx, args); let duration = start.elapsed(); - self.time_profile + self.host_env + .time_profile .entry(desc.name().to_string()) .and_modify(|d| *d += duration.as_millis()) .or_insert(duration.as_millis()); diff --git a/crates/zkwasm/src/runtime/host/internal_circuit_plugin.rs b/crates/zkwasm/src/runtime/host/internal_circuit_plugin.rs index 344cb2a78..67de479b0 100644 --- a/crates/zkwasm/src/runtime/host/internal_circuit_plugin.rs +++ b/crates/zkwasm/src/runtime/host/internal_circuit_plugin.rs @@ -3,6 +3,7 @@ use specs::host_function::Signature; use std::cell::RefCell; use std::collections::HashMap; use std::rc::Rc; +use wasmi::tracer::Observer; use wasmi::FuncInstance; use wasmi::ModuleImportResolver; use wasmi::RuntimeArgs; @@ -16,7 +17,7 @@ pub(super) struct ForeignOp { pub index_within_plugin: usize, pub sig: Signature, pub plugin: HostPlugin, - pub cb: Rc Option>, + pub cb: Rc Option>, } pub struct InternalCircuitEnv { @@ -34,10 +35,16 @@ impl InternalCircuitEnv { } } - pub fn register_plugin(&mut self, plugin: HostPlugin, context: Box) { + pub fn register_plugin( + &mut self, + name: &str, + plugin: HostPlugin, + context: Box, + ) { let _ = self.plugins.insert( plugin, ForeignPlugin { + name: name.to_string(), ctx: Rc::new(RefCell::new(context)), }, ); @@ -49,7 +56,7 @@ impl InternalCircuitEnv { sig: Signature, plugin: HostPlugin, index_within_plugin: usize, - cb: Rc Option>, + cb: Rc Option>, ) { assert!(!*self.finalized.borrow()); diff --git a/crates/zkwasm/src/runtime/host/mod.rs b/crates/zkwasm/src/runtime/host/mod.rs index 32ac94d1b..388b609c5 100644 --- a/crates/zkwasm/src/runtime/host/mod.rs +++ b/crates/zkwasm/src/runtime/host/mod.rs @@ -2,15 +2,13 @@ use self::host_env::HostEnv; use super::wasmi_interpreter::WasmRuntimeIO; use downcast_rs::impl_downcast; use downcast_rs::Downcast; -use serde::Deserialize; -use serde::Serialize; use specs::external_host_call_table::ExternalHostCallSignature; use specs::host_function::HostFunctionDesc; use std::cell::RefCell; -use std::path::PathBuf; use std::rc::Rc; use std::sync::Arc; use std::sync::Mutex; +use wasmi::tracer::Observer; use wasmi::RuntimeArgs; use wasmi::RuntimeValue; use wasmi::Signature; @@ -19,14 +17,6 @@ pub trait ContextOutput { fn get_context_outputs(&self) -> Arc>>; } -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct Sequence { - private_inputs: Vec, - public_inputs: Vec, - context_input: Vec, - pub context_output: Option, -} - pub mod default_env; pub mod external_circuit_plugin; @@ -56,6 +46,12 @@ impl MatchForeignOpSignature for ExternalHostCallSignature { } } +#[derive(Clone, Debug)] +pub struct ForeignStatics { + pub used_round: usize, + pub max_round: usize, +} + /// Context of the plugin. /// /// # Examples @@ -69,17 +65,22 @@ impl MatchForeignOpSignature for ExternalHostCallSignature { /// impl ForeignContext for Context { /// } /// ``` -pub trait ForeignContext: Downcast {} +pub trait ForeignContext: Downcast { + fn get_statics(&self) -> Option { + None + } +} impl_downcast!(ForeignContext); pub struct ForeignPlugin { + pub name: String, ctx: Rc>>, } #[derive(Clone)] struct HostFunctionExecutionEnv { ctx: Rc>>, - cb: Rc Option>, + cb: Rc Option>, } #[derive(Clone)] @@ -92,8 +93,9 @@ struct HostFunction { pub trait HostEnvBuilder { /// Argument type type Arg; + type HostConfig: Default; /// Create an empty env without value, this is used by compiling, computing hash - fn create_env_without_value() -> (HostEnv, WasmRuntimeIO); + fn create_env_without_value(config: Self::HostConfig) -> (HostEnv, WasmRuntimeIO); /// Create an env with execution parameters, this is used by dry-run, run - fn create_env(env: Self::Arg) -> (HostEnv, WasmRuntimeIO); + fn create_env(env: Self::Arg, config: Self::HostConfig) -> (HostEnv, WasmRuntimeIO); } diff --git a/crates/zkwasm/src/runtime/mod.rs b/crates/zkwasm/src/runtime/mod.rs index afb43d21c..ad4bc2b59 100644 --- a/crates/zkwasm/src/runtime/mod.rs +++ b/crates/zkwasm/src/runtime/mod.rs @@ -1,4 +1,5 @@ use std::cell::RefCell; +use std::collections::HashMap; use std::rc::Rc; use specs::etable::EventTableEntry; @@ -11,6 +12,7 @@ use specs::step::StepInfo; use specs::CompilationTable; use specs::Tables; +use self::host::ForeignStatics; use self::wasmi_interpreter::WasmiRuntime; pub mod host; @@ -28,6 +30,8 @@ pub struct ExecutionResult { pub tables: Tables, pub result: Option, pub public_inputs_and_outputs: Vec, + pub host_statics: HashMap, + pub guest_statics: usize, // total instructions used in guest circuits pub outputs: Vec, } diff --git a/crates/zkwasm/src/runtime/wasmi_interpreter.rs b/crates/zkwasm/src/runtime/wasmi_interpreter.rs index 6690a1886..edb6d01aa 100644 --- a/crates/zkwasm/src/runtime/wasmi_interpreter.rs +++ b/crates/zkwasm/src/runtime/wasmi_interpreter.rs @@ -11,11 +11,12 @@ use specs::mtable::MTable; use specs::CompilationTable; use specs::ExecutionTable; use specs::Tables; -use wasmi::Externals; use wasmi::ImportResolver; use wasmi::ModuleInstance; use wasmi::RuntimeValue; +use super::host::host_env::ExecEnv; +use super::host::host_env::HostEnv; use super::CompiledImage; use super::ExecutionResult; @@ -34,9 +35,9 @@ impl WasmRuntimeIO { } pub trait Execution { - fn run( + fn run( self, - externals: &mut E, + externals: HostEnv, dryrun: bool, wasm_io: WasmRuntimeIO, ) -> Result>; @@ -45,19 +46,23 @@ pub trait Execution { impl Execution for CompiledImage, wasmi::tracer::Tracer> { - fn run( + fn run( self, - externals: &mut E, + externals: HostEnv, dryrun: bool, wasm_io: WasmRuntimeIO, ) -> Result> { + let mut exec_env = ExecEnv { + host_env: externals, + tracer: self.tracer.clone(), + }; let instance = self .instance - .run_start_tracer(externals, self.tracer.clone()) + .run_start_tracer(&mut exec_env, self.tracer.clone()) .unwrap(); let result = - instance.invoke_export_trace(&self.entry, &[], externals, self.tracer.clone())?; + instance.invoke_export_trace(&self.entry, &[], &mut exec_env, self.tracer.clone())?; let execution_tables = if !dryrun { let tracer = self.tracer.borrow(); @@ -89,6 +94,8 @@ impl Execution execution_tables, }, result, + host_statics: exec_env.host_env.external_env.get_statics(), + guest_statics: self.tracer.borrow().observer.counter, public_inputs_and_outputs: wasm_io.public_inputs_and_outputs.borrow().clone(), outputs: wasm_io.outputs.borrow().clone(), }) diff --git a/crates/zkwasm/src/test/mod.rs b/crates/zkwasm/src/test/mod.rs index ce73ea5e4..1f9d52033 100644 --- a/crates/zkwasm/src/test/mod.rs +++ b/crates/zkwasm/src/test/mod.rs @@ -64,7 +64,7 @@ fn test_circuit_mock( /// Run function and generate trace. fn compile_then_execute_wasm( - mut env: HostEnv, + env: HostEnv, wasm_runtime_io: WasmRuntimeIO, wasm: Vec, function_name: &str, @@ -83,7 +83,7 @@ fn compile_then_execute_wasm( ) .unwrap(); - let execution_result = compiled_module.run(&mut env, false, wasm_runtime_io)?; + let execution_result = compiled_module.run(env, false, wasm_runtime_io)?; Ok(execution_result) } diff --git a/crates/zkwasm/src/test/test_rlp.rs b/crates/zkwasm/src/test/test_rlp.rs index 75d0e8d66..e57efd3d4 100644 --- a/crates/zkwasm/src/test/test_rlp.rs +++ b/crates/zkwasm/src/test/test_rlp.rs @@ -155,14 +155,21 @@ fn build_circuit() -> Result<( let wasm = std::fs::read("wasm/rlp.wasm").unwrap(); let loader = ZkWasmLoader::::new(20, wasm, vec![])?; + let result = loader + .run( + ExecutionArg { + public_inputs, + private_inputs, + context_inputs: vec![], + context_outputs: Arc::new(Mutex::new(vec![])), + }, + (), + false, + true, + ) + .unwrap(); - let (circuit, instances, _) = loader.circuit_with_witness(ExecutionArg { - public_inputs, - private_inputs, - context_inputs: vec![], - context_outputs: Arc::new(Mutex::new(vec![])), - })?; - + let (circuit, instances) = loader.circuit_with_witness(result)?; Ok((loader, circuit, instances)) } diff --git a/crates/zkwasm/src/test/test_start.rs b/crates/zkwasm/src/test/test_start.rs index 1199dff90..8534cb3cb 100644 --- a/crates/zkwasm/src/test/test_start.rs +++ b/crates/zkwasm/src/test/test_start.rs @@ -39,14 +39,16 @@ mod tests { ZkWasmLoader::::new(18, wasm, vec![]) .unwrap(); - let (circuit, instances, _) = loader - .circuit_with_witness(ExecutionArg { - public_inputs: vec![], - private_inputs: vec![], - context_inputs: vec![], - context_outputs: Arc::new(Mutex::new(vec![])), - }) - .unwrap(); + let arg = ExecutionArg { + public_inputs: vec![], + private_inputs: vec![], + context_inputs: vec![], + context_outputs: Arc::new(Mutex::new(vec![])), + }; + + let result = loader.run(arg, (), false, true).unwrap(); + + let (circuit, instances) = loader.circuit_with_witness(result).unwrap(); loader.mock_test(&circuit, &instances).unwrap() } diff --git a/crates/zkwasm/src/test/test_wasm_instructions/op_call_host.rs b/crates/zkwasm/src/test/test_wasm_instructions/op_call_host.rs index 85720034d..abb426ace 100644 --- a/crates/zkwasm/src/test/test_wasm_instructions/op_call_host.rs +++ b/crates/zkwasm/src/test/test_wasm_instructions/op_call_host.rs @@ -1,8 +1,10 @@ use specs::external_host_call_table::ExternalHostCallSignature; use std::rc::Rc; +use wasmi::tracer::Observer; use crate::runtime::host::host_env::HostEnv; use crate::runtime::host::ForeignContext; +use crate::runtime::host::ForeignStatics; use crate::runtime::wasmi_interpreter::WasmRuntimeIO; use crate::test::test_circuit_with_env; @@ -10,7 +12,11 @@ use crate::test::test_circuit_with_env; struct Context { acc: u64, } -impl ForeignContext for Context {} +impl ForeignContext for Context { + fn get_statics(&self) -> Option { + None + } +} #[test] fn test_call_host_external() { @@ -45,7 +51,7 @@ fn test_call_host_external() { ExternalHostCallSignature::Argument, foreign_playground_plugin.clone(), Rc::new( - |context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); let value: u64 = args.nth(0); @@ -61,7 +67,7 @@ fn test_call_host_external() { ExternalHostCallSignature::Return, foreign_playground_plugin, Rc::new( - |context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { + |_obs: &Observer, context: &mut dyn ForeignContext, _args: wasmi::RuntimeArgs| { let context = context.downcast_mut::().unwrap(); Some(wasmi::RuntimeValue::I64(context.acc as i64))