Fix return handling

qadapt-macro/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "qadapt-macro"
-version = "0.4.0"
+version = "0.5.0"
 authors = ["Bradlee Speice <bradlee@speice.io>"]
 description = "The Quick And Dirty Allocation Profiling Tool - Support Macros"
 license = "Apache-2.0"

qadapt-macro/src/lib.rs

@@ -15,36 +15,30 @@ use proc_macro::TokenTree;
 use std::iter::FromIterator;
 
 macro_rules! group {
-    ($delim:expr, $ts:expr) => {
-        {
-            let _tt: TokenTree = ::proc_macro::Group::new($delim, $ts).into();
-            _tt
-        }
-    };
+    ($delim:expr, $ts:expr) => {{
+        let _tt: TokenTree = ::proc_macro::Group::new($delim, $ts).into();
+        _tt
+    }};
     ($delim:expr) => {
         group!($delim, ::proc_macro::TokenStream::new())
     };
 }
 
 macro_rules! ident {
-    ($name:expr, $span:expr) => {
-        {
-            let _tt: TokenTree = ::proc_macro::Ident::new($name, $span).into();
-            _tt
-        }
-    };
+    ($name:expr, $span:expr) => {{
+        let _tt: TokenTree = ::proc_macro::Ident::new($name, $span).into();
+        _tt
+    }};
     ($name:expr) => {
         ident!($name, ::proc_macro::Span::call_site())
     };
 }
 
 macro_rules! punct {
-    ($ch:expr, $spacing:expr) => {
-        {
-            let _tt: TokenTree = ::proc_macro::Punct::new($ch, $spacing).into();
-            _tt
-        }
-    };
+    ($ch:expr, $spacing:expr) => {{
+        let _tt: TokenTree = ::proc_macro::Punct::new($ch, $spacing).into();
+        _tt
+    }};
 }
 
 macro_rules! token_stream {
@@ -64,6 +58,7 @@ macro_rules! token_stream {
     };
 }
 
+#[rustfmt::skip]
 fn release_guard(fn_name: &str) -> TokenStream {
     // #[cfg(any(debug, test))]
     // { ::qadapt::`fn_name`() }
@@ -92,8 +87,12 @@ fn release_guard(fn_name: &str) -> TokenStream {
     )
 }
 
+/// Generate the body of a function that is protected from making allocations.
+/// The code is conditionally compiled so that all QADAPT-related bits
+/// will be removed for release/bench builds, making the proc_macro safe
+/// to leave on in production.
+#[rustfmt::skip]
 fn protected_body(fn_body: Group) -> TokenTree {
-    // TODO: Don't wrap the release guards in another brace
     group!(Delimiter::Brace, token_stream!(
         group!(Delimiter::Brace, release_guard("enter_protected")),
         ident!("let"),
@@ -101,21 +100,44 @@ fn protected_body(fn_body: Group) -> TokenTree {
         punct!('=', Spacing::Alone),
         fn_body.into(),
         punct!(';', Spacing::Alone),
-        group!(Delimiter::Brace, release_guard("exit_protected")),
-        ident!("__ret__")
+        punct!('#', Spacing::Alone),
+        // When `return` statements are involved, this code can get marked as
+        // unreachable because of early exit
+        group!(Delimiter::Bracket, token_stream!(
+            ident!("allow"),
+            group!(Delimiter::Parenthesis, token_stream!(
+                ident!("unreachable_code")
+            ))
+        )),
+        group!(Delimiter::Brace, token_stream!(
+            group!(Delimiter::Brace, release_guard("exit_protected")),
+            ident!("__ret__")
+        ))
     ))
 }
 
-fn contains_return(ts: TokenStream) -> bool {
-    for tt in ts.into_iter() {
-        match tt {
-            TokenTree::Group(ref g) => if contains_return(g.stream()) { return true },
-            TokenTree::Ident(ref i) if i.to_string() == "return" => return true,
-            _ => (),
-        }
+/// Walk through a TokenStream (typically from a Group Brace) and prepend calls
+/// to `return` with an exit guard.
+fn escape_return(ts: TokenStream) -> TokenStream {
+    let mut protected: Vec<TokenTree> = Vec::new();
+
+    let mut tt_iter = ts.into_iter();
+    while let Some(tt) = tt_iter.next() {
+        let mut tokens = match tt {
+            TokenTree::Group(ref g) if g.delimiter() == Delimiter::Brace => {
+                vec![group!(Delimiter::Brace, escape_return(g.stream()))]
+            }
+            TokenTree::Ident(ref i) if i.to_string() == "return" => vec![
+                group!(Delimiter::Brace, release_guard("exit_protected")),
+                tt.clone(),
+            ],
+            t => vec![t],
+        };
+
+        protected.extend(tokens.into_iter());
     }
 
-    false
+    TokenStream::from_iter(protected.into_iter())
 }
 
 /// Set up the QADAPT allocator to trigger a panic if any allocations happen during
@@ -126,18 +148,7 @@ fn contains_return(ts: TokenStream) -> bool {
 /// separate thread, QADAPT will not trigger a panic.
 #[proc_macro_attribute]
 pub fn allocate_panic(_attr: TokenStream, item: TokenStream) -> TokenStream {
-    if contains_return(item.clone()) {
-        let mut iter = item.clone().into_iter();
-        iter.next();
-        let fn_name = iter.next().unwrap();
-        eprintln!("QADAPT does not currently support `return` \
-                   statements in functions. Function named `{}` \
-                   DOES NOT have allocation guards in place.", fn_name);
-        return item;
-    }
-
     let mut protected_fn: Vec<TokenTree> = Vec::new();
-
     let mut item_iter = item.into_iter();
 
     // First, get the function body we're replicating
@@ -145,7 +156,7 @@ pub fn allocate_panic(_attr: TokenStream, item: TokenStream) -> TokenStream {
     while let Some(tt) = item_iter.next() {
         match tt {
             TokenTree::Group(ref g) if g.delimiter() == Delimiter::Brace => {
-                fn_body = Some(g.clone());
+                fn_body = Some(Group::new(Delimiter::Brace, escape_return(g.stream())));
                 break;
             }
             tt => {