1
0
mirror of https://github.com/bspeice/qadapt synced 2024-11-13 01:28:08 -05:00

Fix return handling

This commit is contained in:
Bradlee Speice 2018-11-15 23:08:02 -05:00
parent 5c9d09cba6
commit 8591f30ac3
5 changed files with 110 additions and 54 deletions

View File

@ -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" }

View File

@ -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"

View File

@ -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 => {

View File

@ -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,

View File

@ -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));
}