From ef33084183298d15c7a957ed959246fbc9c810fe Mon Sep 17 00:00:00 2001 From: Kamen Mladenov Date: Mon, 7 Apr 2025 16:55:22 +0300 Subject: feat(zkvms/zkm): Update to use FunctionDefinition --- zkvms/zkm/wrapper_macro/src/lib.rs | 26 +++++--------------------- 1 file changed, 5 insertions(+), 21 deletions(-) (limited to 'zkvms') 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::(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::()); - out.extend(format!("let {} : {} = read();", prv_pat, prv_typ).parse::()); + out.extend(format!("let {} : {} = read();", fd.grouped_public_patterns(), fd.grouped_public_types()).parse::()); + out.extend(format!("let {} : {} = read();", fd.grouped_private_patterns(), fd.grouped_private_types()).parse::()); - out.extend(format!("commit::<{}>(&zkp::{}{});", ret, name, ts_patterns).parse::()); + out.extend(format!("commit::<{}>(&zkp::{}({}));", fd.return_type, fd.name, fd.grouped_patterns()).parse::()); let mut block = TokenStream::new(); block.extend(format!("{{ {} }}", out).parse::()); -- cgit v1.2.3