feat(zkvms/zkm): Implement proper handling of public and private inputs

4 files changed, 108 insertions, 6 deletions

diff --git a/zkvms/zkm/guest/Cargo.lock b/zkvms/zkm/guest/Cargo.lock
index 848d2cc..f8879d4 100644
--- a/zkvms/zkm/guest/Cargo.lock
+++ b/zkvms/zkm/guest/Cargo.lock
@@ -62,6 +62,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "equivalent"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5"
+
+[[package]]
 name = "generic-array"
 version = "0.14.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -91,6 +97,22 @@ dependencies = [
 ]
 
 [[package]]
+name = "hashbrown"
+version = "0.15.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289"
+
+[[package]]
+name = "indexmap"
+version = "2.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8c9c992b02b5b4c94ea26e32fe5bccb7aa7d9f390ab5c1221ff895bc7ea8b652"
+dependencies = [
+ "equivalent",
+ "hashbrown",
+]
+
+[[package]]
 name = "lazy_static"
 version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -109,6 +131,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa"
 
 [[package]]
+name = "memchr"
+version = "2.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
+
+[[package]]
 name = "ppv-lite86"
 version = "0.2.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -186,6 +214,15 @@ dependencies = [
 ]
 
 [[package]]
+name = "serde_spanned"
+version = "0.6.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "87607cb1398ed59d48732e575a4c28a7a8ebf2454b964fe3f224f2afc07909e1"
+dependencies = [
+ "serde",
+]
+
+[[package]]
 name = "sha2"
 version = "0.10.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -198,9 +235,9 @@ dependencies = [
 
 [[package]]
 name = "syn"
-version = "2.0.96"
+version = "2.0.98"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d5d0adab1ae378d7f53bdebc67a39f1f151407ef230f0ce2883572f5d8985c80"
+checksum = "36147f1a48ae0ec2b5b3bc5b537d267457555a10dc06f3dbc8cb11ba3006d3b1"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -208,6 +245,40 @@ dependencies = [
 ]
 
 [[package]]
+name = "toml"
+version = "0.8.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a1ed1f98e3fdc28d6d910e6737ae6ab1a93bf1985935a1193e68f93eeb68d24e"
+dependencies = [
+ "serde",
+ "serde_spanned",
+ "toml_datetime",
+ "toml_edit",
+]
+
+[[package]]
+name = "toml_datetime"
+version = "0.6.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0dd7358ecb8fc2f8d014bf86f6f638ce72ba252a2c3a2572f2a795f1d23efb41"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.22.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "02a8b472d1a3d7c18e2d61a489aee3453fd9031c33e4f55bd533f4a7adca1bee"
+dependencies = [
+ "indexmap",
+ "serde",
+ "serde_spanned",
+ "toml_datetime",
+ "winnow",
+]
+
+[[package]]
 name = "typenum"
 version = "1.17.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -232,8 +303,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
 
 [[package]]
+name = "winnow"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7e49d2d35d3fad69b39b94139037ecfb4f359f08958b9c11e7315ce770462419"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
 name = "wrapper_macro"
 version = "0.1.0"
+dependencies = [
+ "toml",
+]
 
 [[package]]
 name = "zerocopy"
diff --git a/zkvms/zkm/host/src/main.rs b/zkvms/zkm/host/src/main.rs
index 24b4e92..936e981 100644
--- a/zkvms/zkm/host/src/main.rs
+++ b/zkvms/zkm/host/src/main.rs
@@ -93,13 +93,18 @@ async fn main() -> Result<()> {
 
     // Set input
     let mut public_inputstream = Vec::new();
-    bincode::serialize_into(&mut public_inputstream, &run_info.input)
+    bincode::serialize_into(&mut public_inputstream, &run_info.public_input)
+        .expect("public_input serialization failed");
+
+    let mut private_inputstream = Vec::new();
+    bincode::serialize_into(&mut private_inputstream, &run_info.private_input)
         .expect("private_input serialization failed");
 
     let mut prover_input = ProverInput {
         elf: read(elf_path).unwrap(),
         seg_size,
         public_inputstream,
+        private_inputstream,
         ..Default::default()
     };
 
diff --git a/zkvms/zkm/wrapper_macro/Cargo.toml b/zkvms/zkm/wrapper_macro/Cargo.toml
index 14348e5..6223ed8 100644
--- a/zkvms/zkm/wrapper_macro/Cargo.toml
+++ b/zkvms/zkm/wrapper_macro/Cargo.toml
@@ -5,3 +5,6 @@ edition = "2021"
 
 [lib]
 proc-macro = true
+
+[dependencies]
+toml = "0.8.19"
diff --git a/zkvms/zkm/wrapper_macro/src/lib.rs b/zkvms/zkm/wrapper_macro/src/lib.rs
index 57b9d24..4bb580c 100644
--- a/zkvms/zkm/wrapper_macro/src/lib.rs
+++ b/zkvms/zkm/wrapper_macro/src/lib.rs
@@ -2,17 +2,28 @@ use proc_macro::TokenStream;
 
 #[path = "../../../../guests_macro/src/parse_fn.rs"]
 mod parse_fn;
-use crate::parse_fn::{ split_fn, args_split, args_divide_grouped };
+use crate::parse_fn::{ split_fn, args_split, args_divide_public, args_divide_grouped, group_streams };
 
 #[proc_macro]
 pub fn make_wrapper(item: TokenStream) -> TokenStream {
     let (name, args, ret) = split_fn(&item);
 
-    let (ts_patterns, ts_types) = args_divide_grouped(&args);
+    let (ts_patterns, _) = args_divide_grouped(&args);
+
+    let public_inputs = toml::from_str::<toml::Table>(
+            include_str!(concat!(env!("INPUTS_DIR"), "/default_public_input.toml"))
+        )
+        .unwrap();
+    let ((pub_pat, pub_typ), (prv_pat, prv_typ)) = args_divide_public(&args, &public_inputs.keys().collect());
+    let ((pub_pat, pub_typ), (prv_pat, prv_typ)) = (
+        (group_streams(&pub_pat), group_streams(&pub_typ)),
+        (group_streams(&prv_pat), group_streams(&prv_typ)));
 
     let mut out = TokenStream::new();
     // NOTE: The first read returns public data, the second returns private
-    out.extend(format!("let {} : {} = read();", ts_patterns, ts_types).parse::<TokenStream>());
+    out.extend(format!("let {} : {} = read();", pub_pat, pub_typ).parse::<TokenStream>());
+    out.extend(format!("let {} : {} = read();", prv_pat, prv_typ).parse::<TokenStream>());
+
     out.extend(format!("commit::<{}>(&zkp::{}{});", ret, name, ts_patterns).parse::<TokenStream>());
 
     let mut block = TokenStream::new();