diff options
| -rw-r--r-- | zkvms/nexus/guest/src/main.rs | 2 | ||||
| -rw-r--r-- | zkvms/nexus/wrapper_macro/src/lib.rs | 19 |
2 files changed, 7 insertions, 14 deletions
diff --git a/zkvms/nexus/guest/src/main.rs b/zkvms/nexus/guest/src/main.rs index 75b4d7d..680cf81 100644 --- a/zkvms/nexus/guest/src/main.rs +++ b/zkvms/nexus/guest/src/main.rs @@ -1,6 +1,6 @@ #![cfg_attr(target_arch = "riscv32", no_std, no_main, allow(unused_imports))] -use nexus_rt::{postcard, println, read_private_input, write_public_output}; +use nexus_rt::{postcard, println, read_private_input, read_public_input, write_public_output}; extern crate alloc; use alloc::{collections::*, vec::*}; diff --git a/zkvms/nexus/wrapper_macro/src/lib.rs b/zkvms/nexus/wrapper_macro/src/lib.rs index 061a33f..218359b 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::{args_divide_grouped, args_divide_public, args_split, split_fn}; +use crate::parse_fn::FunctionDefinition; /// Creates a body, which reads all inputs, stores them in variables, then /// writes the ones, defined as public in `default_public_input.toml` to the @@ -32,23 +32,16 @@ use crate::parse_fn::{args_divide_grouped, args_divide_public, args_split, split /// ``` #[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 fd = FunctionDefinition::new(&item); let mut out = TokenStream::new(); - out.extend( - format!( - "let {} = read_private_input::<{}>().unwrap();", - ts_patterns, ts_types - ) - .parse::<TokenStream>(), - ); + out.extend(format!("let {} = read_public_input::<{}>().unwrap();", fd.grouped_public_patterns(), fd.grouped_public_types()).parse::<TokenStream>()); + out.extend(format!("let {} = read_private_input::<{}>().unwrap();", fd.grouped_private_patterns(), fd.grouped_private_types()).parse::<TokenStream>()); out.extend( format!( - "write_public_output::<{}>(&zkp::{}{});", - ret, name, ts_patterns + "write_public_output::<{}>(&zkp::{}({}));", + fd.return_type, fd.name, fd.grouped_patterns() ) .parse::<TokenStream>(), ); |