diff options
| -rw-r--r-- | guests_macro/src/lib.rs | 2 | ||||
| -rw-r--r-- | guests_macro/src/parse_fn.rs | 6 |
2 files changed, 6 insertions, 2 deletions
diff --git a/guests_macro/src/lib.rs b/guests_macro/src/lib.rs index 84f2365..a01a84b 100644 --- a/guests_macro/src/lib.rs +++ b/guests_macro/src/lib.rs @@ -4,6 +4,6 @@ mod parse_fn; #[proc_macro_attribute] pub fn proving_entrypoint(_: TokenStream, mut item: TokenStream) -> TokenStream { let (name, args, ret) = parse_fn::split_fn(&item); - item.extend(format!("#[macro_export] macro_rules! entrypoint_expr {{ () => {{ make_wrapper!({}{} {}) }}; }}", name, args, ret).parse::<TokenStream>()); + item.extend(format!("#[macro_export] macro_rules! entrypoint_expr {{ () => {{ make_wrapper!({}{} -> {}) }}; }}", name, args, ret).parse::<TokenStream>()); item } diff --git a/guests_macro/src/parse_fn.rs b/guests_macro/src/parse_fn.rs index e78fb1b..1bcc3ba 100644 --- a/guests_macro/src/parse_fn.rs +++ b/guests_macro/src/parse_fn.rs @@ -1,7 +1,7 @@ use proc_macro::{ TokenStream, TokenTree, Delimiter, Spacing, Group }; /// Input: "fn name(...) -> ... { ... }" -/// Output: "name", "(...)", "-> ..." +/// Output: "name", "(...)", "..." pub fn split_fn(item: &TokenStream) -> (TokenStream, TokenStream, TokenStream) { let item = item.clone().into_iter(); @@ -22,6 +22,10 @@ pub fn split_fn(item: &TokenStream) -> (TokenStream, TokenStream, TokenStream) { TokenTree::Punct(ref punct) => { if punct.as_char() == '-' { out = &mut ret; + continue; + } + if punct.as_char() == '>' && out.is_empty() { + continue; } }, TokenTree::Group(ref group) => { |