From 84739c4dd103a22ac8f7865a5d4263dd478daae3 Mon Sep 17 00:00:00 2001 From: Bradlee Speice Date: Mon, 12 Nov 2018 22:30:09 -0500 Subject: [PATCH] Handle actual functions --- qadapt-macro/src/lib.rs | 75 +++++++++++++++++++++++++++++++++++------ src/lib.rs | 6 ++-- tests/macros.rs | 57 +++++++++++++++++++++++++++++-- 3 files changed, 124 insertions(+), 14 deletions(-) diff --git a/qadapt-macro/src/lib.rs b/qadapt-macro/src/lib.rs index 6470ff6..fbcc611 100644 --- a/qadapt-macro/src/lib.rs +++ b/qadapt-macro/src/lib.rs @@ -20,7 +20,27 @@ type G = proc_macro::Group; type I = proc_macro::Ident; type P = proc_macro::Punct; -fn protected_group(fn_body: G) -> TokenTree { +fn protected_body(fn_name: &str, args: G) -> TokenTree { + let mut args_filtered = Vec::new(); + let mut waiting_for_comma = false; + let mut in_type = 0; + for tt in args.stream().into_iter() { + match tt { + TokenTree::Ident(ref _i) if !waiting_for_comma && in_type == 0 => { + args_filtered.push(tt.clone()); + waiting_for_comma = true; + } + TokenTree::Punct(ref p) if p.as_char() == '<' => in_type += 1, + TokenTree::Punct(ref p) if p.as_char() == '>' => in_type -= 1, + TokenTree::Punct(ref p) if p.as_char() == ',' && in_type == 0 => { + waiting_for_comma = false; + args_filtered.push(tt.clone()) + } + _ => () + } + } + let args_group = G::new(Delimiter::Parenthesis, TS::from_iter(args_filtered)); + let tt: Vec = vec![ P::new(':', Spacing::Joint).into(), P::new(':', Spacing::Alone).into(), @@ -30,7 +50,12 @@ fn protected_group(fn_body: G) -> TokenTree { I::new("enter_protected", Span::call_site()).into(), G::new(Delimiter::Parenthesis, TokenStream::new()).into(), P::new(';', Spacing::Alone).into(), - fn_body.into(), + I::new("let", Span::call_site()).into(), + I::new("__ret__", Span::call_site()).into(), + P::new('=', Spacing::Alone).into(), + I::new(fn_name, Span::call_site()).into(), + args_group.into(), + P::new(';', Spacing::Alone).into(), P::new(':', Spacing::Joint).into(), P::new(':', Spacing::Alone).into(), I::new("qadapt", Span::call_site()).into(), @@ -39,6 +64,7 @@ fn protected_group(fn_body: G) -> TokenTree { I::new("exit_protected", Span::call_site()).into(), G::new(Delimiter::Parenthesis, TokenStream::new()).into(), P::new(';', Spacing::Alone).into(), + I::new("__ret__", Span::call_site()).into(), ]; G::new(Delimiter::Brace, TS::from_iter(tt)).into() @@ -52,19 +78,48 @@ fn protected_group(fn_body: G) -> TokenTree { /// separate thread, QADAPT will not trigger a panic. #[proc_macro_attribute] pub fn allocate_panic(_attr: TokenStream, item: TokenStream) -> TokenStream { - let mut ret: Vec = Vec::new(); + let mut original_fn: Vec = Vec::new(); + let mut protected_fn: Vec = Vec::new(); - let mut fn_body = None; let mut item_iter = item.into_iter(); + + // First, get the function name we're replacing + let mut fn_name = None; + let mut fn_args = None; while let Some(tt) = item_iter.next() { match tt { - TokenTree::Group(ref g) if g.delimiter() == Delimiter::Brace => { - fn_body = Some(g.clone()); - break; + TokenTree::Ident(ref i) if i.to_string() == "fn" => { + original_fn.push(tt.clone()); + protected_fn.push(tt.clone()); + } + TokenTree::Ident(ref i) if fn_args.is_none() => { + let changed_name = format!("__{}__", i.to_owned()); + original_fn.push(TokenTree::Ident(I::new(&changed_name, Span::call_site()))); + protected_fn.push(tt.clone()); + fn_name = Some(changed_name); + } + TokenTree::Group(ref g) if g.delimiter() == Delimiter::Parenthesis && fn_args.is_none() => { + original_fn.push(tt.clone()); + protected_fn.push(tt.clone()); + fn_args = Some(g.clone()); + } + TokenTree::Group(ref g) if g.delimiter() == Delimiter::Brace => { + original_fn.push(tt.clone()); + protected_fn.push(protected_body( + &fn_name.take().unwrap(), + fn_args.take().unwrap(), + )); + } + tt => { + original_fn.push(tt.clone()); + protected_fn.push(tt.clone()); } - tt => ret.push(tt), } } - ret.push(protected_group(fn_body.unwrap())); - TokenStream::from_iter(ret.into_iter()) + + let mut full = Vec::new(); + full.push(TS::from_iter(original_fn)); + full.push(TS::from_iter(protected_fn)); + let ts = TS::from_iter(full.into_iter()); + ts } diff --git a/src/lib.rs b/src/lib.rs index 219e0c6..1f1fe2f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,7 +11,7 @@ //! so the attribute is safe to leave everywhere. //! //! Currently this crate is Nightly-only, but will work once `const fn` is in Stable. -//! +//! //! Please also take a look at [qadapt-macro](https://github.com/bspeice/qadapt/tree/master/qadapt-macro) //! for some helper macros to make working with QADAPT a bit easier. #![deny(missing_docs)] @@ -135,7 +135,7 @@ unsafe impl GlobalAlloc for QADAPT { } } - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) { if alloc_immediate() { return free(ptr as *mut c_void); } @@ -152,11 +152,13 @@ unsafe impl GlobalAlloc for QADAPT { // Tripped a bad dealloc, but make sure further memory access during unwind // doesn't have issues PROTECTION_LEVEL.with(|v| *v.write() = 0); + /* panic!( "Unexpected deallocation for size {}, protection level: {}", layout.size(), v ) + */ } _ => (), } diff --git a/tests/macros.rs b/tests/macros.rs index 3d19ff1..3cd214a 100644 --- a/tests/macros.rs +++ b/tests/macros.rs @@ -1,4 +1,5 @@ extern crate qadapt; +use std::io; use qadapt::allocate_panic; use qadapt::QADAPT; @@ -20,13 +21,65 @@ fn no_allocate() { let _v: Vec<()> = Vec::with_capacity(0); } +#[allocate_panic] +fn no_allocate_ret() -> bool { + return true; +} + +#[allocate_panic] +fn no_allocate_implicit_ret() -> bool { + true +} + +#[allocate_panic] +fn no_allocate_arg(b: bool) -> bool { + b +} + +#[allocate_panic] +fn no_allocate_args(_b: bool, _u: usize, i: i64) -> i64 { + i +} + +#[allocate_panic] +fn return_result(r: Result) -> Result, ()> { + Ok(r) +} + #[test] -fn test_no_allocate() { +fn macro_no_allocate() { no_allocate(); } #[test] #[should_panic] -fn test_allocates() { +fn macro_allocates() { allocates(); } + +#[test] +fn macro_return() { + assert!(no_allocate_ret()); +} + +#[test] +fn macro_implicit_return() { + assert!(no_allocate_ret()); +} + +#[test] +fn macro_allocate_arg() { + no_allocate_arg(true); + no_allocate_arg(false); +} + +#[test] +fn macro_allocate_args() { + no_allocate_args(true, 0, -1); + no_allocate_args(false, 4, -90); +} + +#[test] +fn macro_return_result() { + return_result(Ok(16)).unwrap().unwrap(); +} \ No newline at end of file