From 8591f30ac3d9c32916c63e437a4929dff8f51032 Mon Sep 17 00:00:00 2001 From: Bradlee Speice Date: Thu, 15 Nov 2018 23:08:02 -0500 Subject: [PATCH] Fix `return` handling --- Cargo.toml | 4 +- qadapt-macro/Cargo.toml | 2 +- qadapt-macro/src/lib.rs | 93 +++++++++++++++++++++++------------------ tests/allocations.rs | 20 ++++----- tests/macros.rs | 45 ++++++++++++++++++++ 5 files changed, 110 insertions(+), 54 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index f4f6a25..0a5f117 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "qadapt" -version = "0.4.0" +version = "0.5.0" authors = ["Bradlee Speice "] 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" } \ No newline at end of file +qadapt-macro = { version = "0.5.0", path = "./qadapt-macro" } \ No newline at end of file diff --git a/qadapt-macro/Cargo.toml b/qadapt-macro/Cargo.toml index 39630be..53c5adf 100644 --- a/qadapt-macro/Cargo.toml +++ b/qadapt-macro/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "qadapt-macro" -version = "0.4.0" +version = "0.5.0" authors = ["Bradlee Speice "] description = "The Quick And Dirty Allocation Profiling Tool - Support Macros" license = "Apache-2.0" diff --git a/qadapt-macro/src/lib.rs b/qadapt-macro/src/lib.rs index b2a7951..04d2b9a 100644 --- a/qadapt-macro/src/lib.rs +++ b/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 = 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 = 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 => { diff --git a/tests/allocations.rs b/tests/allocations.rs index 5f0388d..13616a8 100644 --- a/tests/allocations.rs +++ b/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 = Vec::with_capacity(0); exit_protected(); } #[test] -fn test_vec_new() { +fn vec_new() { enter_protected(); let _v: Vec = 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 = Vec::with_capacity(1); // We don't make it here in debug mode, but in release mode, diff --git a/tests/macros.rs b/tests/macros.rs index cd5d252..f537142 100644 --- a/tests/macros.rs +++ b/tests/macros.rs @@ -87,3 +87,48 @@ fn return_result(r: Result) -> Result 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)); +}