Fix return handling

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "qadapt"
-version = "0.4.0"
+version = "0.5.0"
 authors = ["Bradlee Speice <bradlee@speice.io>"]
 description = "The Quick And Dirty Allocation Profiling Tool"
 license = "Apache-2.0"
@@ -22,4 +22,4 @@ libc = "0.2"
 spin = "0.4"
 thread-id = "3.3"
 
-qadapt-macro = { version = "0.4.0", path = "./qadapt-macro" }
+qadapt-macro = { version = "0.5.0", path = "./qadapt-macro" }

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) => {
-        {
+    ($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) => {
-        {
+    ($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) => {
-        {
+    ($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),
+        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 => {

tests/allocations.rs

@@ -12,7 +12,7 @@ static Q: QADAPT = QADAPT;
 fn test_copy() {
     enter_protected();
     let v = 0u8;
-    let v2 = v;
+    let _v2 = v;
     exit_protected();
 }
 
@@ -24,7 +24,7 @@ fn test_allocate() {
     exit_protected();
 }
 
-fn unit_result(b: bool) -> Result<(), ()> {
+fn return_unit_result(b: bool) -> Result<(), ()> {
     if b {
         Ok(())
     } else {
@@ -33,16 +33,16 @@ fn unit_result(b: bool) -> Result<(), ()> {
 }
 
 #[test]
-fn test_unit_result() {
+fn unit_result() {
     enter_protected();
-    unit_result(true).unwrap();
-    unit_result(false).unwrap_err();
+    return_unit_result(true).unwrap();
+    return_unit_result(false).unwrap_err();
     exit_protected();
 }
 
 #[test]
 #[should_panic]
-fn test_vec_push() {
+fn vec_push() {
     let mut v = Vec::new();
     enter_protected();
     v.push(0);
@@ -54,7 +54,7 @@ fn test_vec_push() {
 }
 
 #[test]
-fn test_vec_push_capacity() {
+fn vec_push_capacity() {
     let mut v = Vec::with_capacity(1);
     enter_protected();
     v.push(0);
@@ -64,14 +64,14 @@ fn test_vec_push_capacity() {
 }
 
 #[test]
-fn test_vec_with_zero() {
+fn vec_with_zero() {
     enter_protected();
     let _v: Vec<u8> = Vec::with_capacity(0);
     exit_protected();
 }
 
 #[test]
-fn test_vec_new() {
+fn vec_new() {
     enter_protected();
     let _v: Vec<u8> = Vec::new();
     exit_protected();
@@ -79,7 +79,7 @@ fn test_vec_new() {
 
 #[test]
 #[should_panic]
-fn test_vec_with_one() {
+fn vec_with_one() {
     enter_protected();
     let v: Vec<u8> = Vec::with_capacity(1);
     // We don't make it here in debug mode, but in release mode,

tests/macros.rs

@@ -87,3 +87,48 @@ fn return_result(r: Result<usize, io::Error>) -> Result<Result<usize, io::Error>
 fn macro_return_result() {
     return_result(Ok(16)).unwrap().unwrap();
 }
+
+#[allocate_panic]
+fn branching_return(a: bool, b: bool, c: bool) -> u8 {
+    if a {
+        if b {
+            if c {
+                return 1;
+            } else {
+                return 2;
+            }
+        } else {
+            if c {
+                return 3;
+            } else {
+                return 4;
+            }
+        }
+    } else {
+        if b {
+            if c {
+                return 5;
+            } else {
+                return 6;
+            }
+        } else {
+            if c {
+                return 7;
+            } else {
+                return 8;
+            }
+        }
+    }
+}
+
+#[test]
+fn macro_branching_return() {
+    assert_eq!(1, branching_return(true, true, true));
+    assert_eq!(2, branching_return(true, true, false));
+    assert_eq!(3, branching_return(true, false, true));
+    assert_eq!(4, branching_return(true, false, false));
+    assert_eq!(5, branching_return(false, true, true));
+    assert_eq!(6, branching_return(false, true, false));
+    assert_eq!(7, branching_return(false, false, true));
+    assert_eq!(8, branching_return(false, false, false));
+}