feat(zkvms/nexus): Make guest commit public input

2 files changed, 18 insertions, 1 deletions

diff --git a/zkvms/nexus/wrapper_macro/Cargo.toml b/zkvms/nexus/wrapper_macro/Cargo.toml
index 14348e5..6223ed8 100644
--- a/zkvms/nexus/wrapper_macro/Cargo.toml
+++ b/zkvms/nexus/wrapper_macro/Cargo.toml
@@ -5,3 +5,6 @@ edition = "2021"
 
 [lib]
 proc-macro = true
+
+[dependencies]
+toml = "0.8.19"
diff --git a/zkvms/nexus/wrapper_macro/src/lib.rs b/zkvms/nexus/wrapper_macro/src/lib.rs
index 54ea2e9..0814e2a 100644
--- a/zkvms/nexus/wrapper_macro/src/lib.rs
+++ b/zkvms/nexus/wrapper_macro/src/lib.rs
@@ -2,7 +2,7 @@ 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 };
 
 #[proc_macro]
 pub fn make_wrapper(item: TokenStream) -> TokenStream {
@@ -12,6 +12,20 @@ pub fn make_wrapper(item: TokenStream) -> TokenStream {
 
     let mut out = TokenStream::new();
     out.extend(format!("let {} = read_private_input::<{}>().unwrap();", ts_patterns, ts_types).parse::<TokenStream>());
+
+    let public_inputs = toml::from_str::<toml::Table>(
+            include_str!(concat!(env!("INPUTS_DIR"), "/default_public_input.toml"))
+        )
+        .unwrap();
+    let (public_patterns, public_types) = args_divide_public(&args, &public_inputs.keys().collect());
+    let public_patterns: Vec<(TokenStream, TokenStream)> = public_patterns
+        .into_iter()
+        .zip(public_types.into_iter())
+        .collect();
+    for (pattern, ptype) in public_patterns {
+        out.extend(format!("write_output::<{}>(&{});", ptype, pattern).parse::<TokenStream>());
+    }
+
     out.extend(format!("write_output::<{}>(&zkp::{}{});", ret, name, ts_patterns).parse::<TokenStream>());
 
     let mut block = TokenStream::new();