1 files changed, 17 insertions, 20 deletions
diff --git a/zkvms/risc0/wrapper_macro/src/lib.rs b/zkvms/risc0/wrapper_macro/src/lib.rs
index 7fa328b..c43f5c1 100644
--- a/zkvms/risc0/wrapper_macro/src/lib.rs
+++ b/zkvms/risc0/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;
 
 /// Create a body, which reads all inputs, stores them in variables, then
 /// commits the ones, defined as public in `default_public_input.toml` to the
@@ -32,29 +32,26 @@ 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 args_split = args_split(&args);
+    let fd = FunctionDefinition::new(&item);
 
-    let mut out = TokenStream::new();
-    for arg in args_split {
-        out.extend(format!("let {} = read();", arg).parse::<TokenStream>());
-    }
+    let args = fd
+        .arguments()
+        .into_iter()
+        .map(|x| format!("let {x} = read();"))
+        .collect::<String>();
 
-    let public_inputs = toml::from_str::<toml::Table>(include_str!(concat!(
-        env!("INPUTS_DIR"),
-        "/default_public_input.toml"
-    )))
-    .unwrap();
-    let public_patterns = args_divide_public(&args, &public_inputs.keys().collect())
-        .0
-         .0;
-    for pattern in public_patterns.iter() {
-        out.extend(format!("commit(&{});", pattern).parse::<TokenStream>());
-    }
+    let mut out = TokenStream::new();
+    out.extend(args.parse::<TokenStream>());
 
-    let (ts_patterns, _) = args_divide_grouped(&args);
+    let commits = fd
+        .public_patterns()
+        .clone()
+        .into_iter()
+        .map(|x| format!("commit(&{x});"))
+        .collect::<String>();
+    out.extend(commits.parse::<TokenStream>());
 
-    out.extend(format!("commit(&zkp::{}{});", name, ts_patterns).parse::<TokenStream>());
+    out.extend(format!("commit(&zkp::{}({}));", fd.name, fd.grouped_patterns()).parse::<TokenStream>());
 
     let mut block = TokenStream::new();
     block.extend(format!("{{ {} }}", out).parse::<TokenStream>());