diff options
| -rw-r--r-- | zkvms/zkm/wrapper_macro/src/lib.rs | 26 |
1 files changed, 5 insertions, 21 deletions
diff --git a/zkvms/zkm/wrapper_macro/src/lib.rs b/zkvms/zkm/wrapper_macro/src/lib.rs index cb8060a..4e00cc5 100644 --- a/zkvms/zkm/wrapper_macro/src/lib.rs +++ b/zkvms/zkm/wrapper_macro/src/lib.rs @@ -2,9 +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, group_streams, split_fn, -}; +use crate::parse_fn::FunctionDefinition; /// Create a body, which reads all public and private inputs, stores them in /// variables, then executes the guest entrypoint function with those arguments @@ -31,28 +29,14 @@ use crate::parse_fn::{ /// ``` #[proc_macro] pub fn make_wrapper(item: TokenStream) -> TokenStream { - let (name, args, ret) = split_fn(&item); - - 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 fd = FunctionDefinition::new(&item); let mut out = TokenStream::new(); // NOTE: The first read returns public data, the second returns private - 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!("let {} : {} = read();", fd.grouped_public_patterns(), fd.grouped_public_types()).parse::<TokenStream>()); + out.extend(format!("let {} : {} = read();", fd.grouped_private_patterns(), fd.grouped_private_types()).parse::<TokenStream>()); - out.extend(format!("commit::<{}>(&zkp::{}{});", ret, name, ts_patterns).parse::<TokenStream>()); + out.extend(format!("commit::<{}>(&zkp::{}({}));", fd.return_type, fd.name, fd.grouped_patterns()).parse::<TokenStream>()); let mut block = TokenStream::new(); block.extend(format!("{{ {} }}", out).parse::<TokenStream>()); |