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

5 files changed, 174 insertions, 27 deletions

diff --git a/zkvms/zkwasm/guest/Cargo.lock b/zkvms/zkwasm/guest/Cargo.lock
index 8418182..7dac1bd 100644
--- a/zkvms/zkwasm/guest/Cargo.lock
+++ b/zkvms/zkwasm/guest/Cargo.lock
@@ -16,9 +16,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
 
 [[package]]
 name = "bytes"
-version = "1.9.0"
+version = "1.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "325918d6fe32f23b19878fe4b34794ae41fc19ddbe53b10571a4874d44ffd39b"
+checksum = "f61dac84819c6588b558454b194026eb1f09c293b9036ae9b159e74e73ab6cf9"
 
 [[package]]
 name = "cfg-if"
@@ -33,6 +33,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "43da5946c66ffcc7745f48db692ffbb10a83bfe0afd96235c5c2a4fb23994929"
 
 [[package]]
+name = "equivalent"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5"
+
+[[package]]
 name = "fixed-hash"
 version = "0.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -52,18 +58,40 @@ dependencies = [
 ]
 
 [[package]]
+name = "hashbrown"
+version = "0.15.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289"
+
+[[package]]
 name = "hex"
 version = "0.4.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
 
 [[package]]
+name = "indexmap"
+version = "2.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8c9c992b02b5b4c94ea26e32fe5bccb7aa7d9f390ab5c1221ff895bc7ea8b652"
+dependencies = [
+ "equivalent",
+ "hashbrown",
+]
+
+[[package]]
 name = "log"
 version = "0.4.25"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "04cbf5b083de1c7e0222a7a51dbfdba1cbe1c6ab0b15e29fff3f6c077fd9cd9f"
 
 [[package]]
+name = "memchr"
+version = "2.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
+
+[[package]]
 name = "once_cell"
 version = "1.20.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -114,6 +142,35 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3e75f6a532d0fd9f7f13144f392b6ad56a32696bfcd9c78f797f16bbb6f072d6"
 
 [[package]]
+name = "serde"
+version = "1.0.217"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "02fc4265df13d6fa1d00ecff087228cc0a2b5f3c0e87e258d8b94a156e984c70"
+dependencies = [
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_derive"
+version = "1.0.217"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "serde_spanned"
+version = "0.6.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "87607cb1398ed59d48732e575a4c28a7a8ebf2454b964fe3f224f2afc07909e1"
+dependencies = [
+ "serde",
+]
+
+[[package]]
 name = "static_assertions"
 version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -121,9 +178,9 @@ checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
 
 [[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",
@@ -131,6 +188,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 = "uint"
 version = "0.9.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -204,8 +295,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "65fc09f10666a9f147042251e0dda9c18f166ff7de300607007e96bdebc1068d"
 
 [[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 = "zkwasm-rust-sdk"
diff --git a/zkvms/zkwasm/guest/src/lib.rs b/zkvms/zkwasm/guest/src/lib.rs
index 5c00fc0..eba34b0 100644
--- a/zkvms/zkwasm/guest/src/lib.rs
+++ b/zkvms/zkwasm/guest/src/lib.rs
@@ -7,6 +7,10 @@ fn read_private() -> u64 {
     unsafe { wasm_input(0) }
 }
 
+fn read_public() -> u64 {
+    unsafe { wasm_input(1) }
+}
+
 fn assert(cond: bool) {
     unsafe { require(cond); }
 }
@@ -19,11 +23,11 @@ static VERTICES: u64 = 10;
 
 macro_rules! read {
     // HACK for graph_coloring
-    (Vec , u32) => {
+    (Vec , u32 $readfn:tt) => {
         {
             let mut ret = Vec::new();
             for _ in 0..2 {
-                ret.push(read!(u32));
+                ret.push(read!(u32 $readfn));
             }
             ret
         }
@@ -39,12 +43,12 @@ macro_rules! read {
             ret
         }
     };
-    (bool) => {
-        (read_private() != 0)
+    (bool $readfn:tt) => {
+        ($readfn() != 0)
     };
     // Has to be primitive!
-    ($type:ty) => {
-        (read_private() as $type)
+    ($type:tt $readfn:tt) => {
+        ($readfn() as $type)
     };
 }
 
diff --git a/zkvms/zkwasm/host/src/main.rs b/zkvms/zkwasm/host/src/main.rs
index 16caea1..d6c5381 100644
--- a/zkvms/zkwasm/host/src/main.rs
+++ b/zkvms/zkwasm/host/src/main.rs
@@ -1,13 +1,36 @@
-use zkvms_host_io::{Input, foreach_input_field, read_args, RunType::{Execute, Prove, Verify}};
+use zkvms_host_io::{PublicInput, PrivateInput, foreach_public_input_field, foreach_private_input_field, read_args, RunType::{Execute, Prove, Verify}};
 use std::io::{self, Write};
 use std::process::{Command, Stdio};
 use regex::Regex;
 
-fn build_input(input: &Input) -> String {
+fn build_public_input(input: &PublicInput) -> String {
     let numreg: Regex = Regex::new("(?:^|[^A-Za-z])([0-9]+)").unwrap();
 
     let mut ret = String::new();
-    foreach_input_field!{
+    foreach_public_input_field!{
+        let flat = format!("{:?}", input.yield)
+            .replace("false", "0")
+            .replace("true",  "1");
+
+        let numbers: Vec<&str> = numreg
+            .captures_iter(&flat)
+            .map(|cap| cap.get(1).unwrap().as_str())
+            .collect();
+
+        for num in numbers {
+            ret.push_str(num);
+            ret.push_str(":i64,");
+        }
+    }
+    ret.pop(); // removes trailing comma
+    ret
+}
+
+fn build_private_input(input: &PrivateInput) -> String {
+    let numreg: Regex = Regex::new("(?:^|[^A-Za-z])([0-9]+)").unwrap();
+
+    let mut ret = String::new();
+    foreach_private_input_field!{
         let flat = format!("{:?}", input.yield)
             .replace("false", "0")
             .replace("true",  "1");
@@ -58,7 +81,8 @@ fn main() {
         .arg("-k").arg(k)
         .arg("--scheme").arg(scheme));
 
-    let input = build_input(&run_info.input);
+    let public_input = build_public_input(&run_info.public_input);
+    let private_input = build_private_input(&run_info.private_input);
 
     let output = run_info
         .env_or(
@@ -75,17 +99,20 @@ fn main() {
     match run_info.run_type {
         Execute => {
             run(zkwasm_command("dry-run")
-                .arg("--private").arg(input)
+                .arg("--public").arg(public_input)
+                .arg("--private").arg(private_input)
                 .arg("--output").arg(output));
         },
         Prove => {
             run(zkwasm_command("prove")
-                .arg("--private").arg(input)
+                .arg("--public").arg(public_input)
+                .arg("--private").arg(private_input)
                 .arg("--output").arg(output));
         },
         Verify => {
             run(zkwasm_command("prove")
-                .arg("--private").arg(input)
+                .arg("--public").arg(public_input)
+                .arg("--private").arg(private_input)
                 .arg("--output").arg(output.clone()));
 
             run(Command::new("zkwasm-cli")
diff --git a/zkvms/zkwasm/wrapper_macro/Cargo.toml b/zkvms/zkwasm/wrapper_macro/Cargo.toml
index 14348e5..6223ed8 100644
--- a/zkvms/zkwasm/wrapper_macro/Cargo.toml
+++ b/zkvms/zkwasm/wrapper_macro/Cargo.toml
@@ -5,3 +5,6 @@ edition = "2021"
 
 [lib]
 proc-macro = true
+
+[dependencies]
+toml = "0.8.19"
diff --git a/zkvms/zkwasm/wrapper_macro/src/lib.rs b/zkvms/zkwasm/wrapper_macro/src/lib.rs
index edb2fcd..8fe7fd0 100644
--- a/zkvms/zkwasm/wrapper_macro/src/lib.rs
+++ b/zkvms/zkwasm/wrapper_macro/src/lib.rs
@@ -2,15 +2,9 @@ use proc_macro::TokenStream;
 
 #[path = "../../../../guests_macro/src/parse_fn.rs"]
 mod parse_fn;
-use crate::parse_fn::{ split_fn, args_split, args_divide, group_streams };
-
-#[proc_macro]
-pub fn make_wrapper(item: TokenStream) -> TokenStream {
-    let (name, args, ret) = split_fn(&item);
-    let (patterns, types) = args_divide(&args);
-
-    let mut out = TokenStream::new();
+use crate::parse_fn::{ split_fn, args_divide_grouped, args_divide_public, group_streams };
 
+fn insert_reads(out: &mut TokenStream, patterns: &Vec<TokenStream>, types: &Vec<TokenStream>, readfn: &str) {
     for i in 0..patterns.len() {
         let type_note: String = format!("{}", types[i])
             .chars()
@@ -20,10 +14,26 @@ pub fn make_wrapper(item: TokenStream) -> TokenStream {
                  _  => c,
             })
             .collect();
-        out.extend(format!("let {} : {} = read!({});", patterns[i], types[i], type_note).parse::<TokenStream>());
+        out.extend(format!("let {} : {} = read!({} {});", patterns[i], types[i], type_note, readfn).parse::<TokenStream>());
     }
+}
+
+#[proc_macro]
+pub fn make_wrapper(item: TokenStream) -> TokenStream {
+    let (name, args, ret) = split_fn(&item);
+
+    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 mut out = TokenStream::new();
+
+    insert_reads(&mut out, &pub_pat, &pub_typ, "read_public");
+    insert_reads(&mut out, &prv_pat, &prv_typ, "read_private");
 
-    let ts_patterns = group_streams(&patterns);
+    let (ts_patterns, _) = args_divide_grouped(&args);
 
     out.extend(format!("let result = zkp::{}{}; assert(result); write(result as u64);", name, ts_patterns).parse::<TokenStream>());