1 files changed, 32 insertions, 26 deletions

diff --git a/src/ast.rs b/src/ast.rs
index b156493..88a2b96 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -1,16 +1,12 @@
-use proc_macro::{Delimiter, Literal, Span, TokenStream, TokenTree};
+use proc_macro::{Delimiter, Ident, Literal, Span, TokenStream, TokenTree};
 
-use crate::{
-    iter::TokenIter,
-    to_tokens::ToTokens,
-    utils::{parse_as_empty, tt_span},
-    Result,
-};
+use crate::{iter::TokenIter, to_tokens::ToTokens, utils::tt_span, Result};
 
 pub(crate) struct Func {
-    pub(crate) attrs: Vec<Attribute>,
-    pub(crate) sig: Vec<TokenTree>,
-    pub(crate) body: TokenTree,
+    attrs: Vec<Attribute>,
+    // [const] [async] [unsafe] [extern [<abi>]] fn
+    sig: Vec<TokenTree>,
+    body: TokenStream,
     pub(crate) print_const: bool,
 }
 
@@ -19,10 +15,9 @@ pub(crate) fn parse_input(input: TokenStream) -> Result<Func> {
 
     let attrs = parse_attrs(&mut input)?;
     let sig = parse_signature(&mut input);
-    let body = input.next();
-    parse_as_empty(input)?;
+    let body: TokenStream = input.collect();
 
-    if body.is_none()
+    if body.is_empty()
         || !sig
             .iter()
             .any(|tt| if let TokenTree::Ident(i) = tt { i.to_string() == "fn" } else { false })
@@ -32,19 +27,8 @@ pub(crate) fn parse_input(input: TokenStream) -> Result<Func> {
             "#[const_fn] attribute may only be used on functions"
         ));
     }
-    if !sig
-        .iter()
-        .any(|tt| if let TokenTree::Ident(i) = tt { i.to_string() == "const" } else { false })
-    {
-        let span = sig
-            .iter()
-            .position(|tt| if let TokenTree::Ident(i) = tt { i.to_string() == "fn" } else { false })
-            .map(|i| sig[i].span())
-            .unwrap();
-        return Err(error!(span, "#[const_fn] attribute may only be used on const functions"));
-    }
 
-    Ok(Func { attrs, sig, body: body.unwrap(), print_const: true })
+    Ok(Func { attrs, sig, body, print_const: true })
 }
 
 impl ToTokens for Func {
@@ -66,10 +50,32 @@ impl ToTokens for Func {
 
 fn parse_signature(input: &mut TokenIter) -> Vec<TokenTree> {
     let mut sig = Vec::new();
+    let mut has_const = false;
     loop {
         match input.peek() {
-            Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Brace => break,
             None => break,
+            Some(TokenTree::Ident(i)) if !has_const => {
+                match &*i.to_string() {
+                    "fn" => {
+                        sig.push(TokenTree::Ident(Ident::new("const", i.span())));
+                        sig.push(input.next().unwrap());
+                        break;
+                    }
+                    "const" => {
+                        has_const = true;
+                    }
+                    "async" | "unsafe" | "extern" => {
+                        has_const = true;
+                        sig.push(TokenTree::Ident(Ident::new("const", i.span())));
+                    }
+                    _ => {}
+                }
+                sig.push(input.next().unwrap());
+            }
+            Some(TokenTree::Ident(i)) if i.to_string() == "fn" => {
+                sig.push(input.next().unwrap());
+                break;
+            }
             Some(_) => sig.push(input.next().unwrap()),
         }
     }