From 8bc768e40b04480644cabe2e262ce0d30acad58b Mon Sep 17 00:00:00 2001
From: Bradlee Speice <bradlee@speice.io>
Date: Sun, 11 Nov 2018 22:22:24 -0500
Subject: [PATCH] Get the proc macro working. It's time to release.

---
 qadapt-macro/src/lib.rs | 56 +++++++++++++++++++++++++++++++++++------
 src/lib.rs              |  5 ++++
 tests/macros.rs         | 11 +++++---
 3 files changed, 61 insertions(+), 11 deletions(-)

diff --git a/qadapt-macro/src/lib.rs b/qadapt-macro/src/lib.rs
index 7494913..300c485 100644
--- a/qadapt-macro/src/lib.rs
+++ b/qadapt-macro/src/lib.rs
@@ -7,8 +7,42 @@
 // #![deny(missing_docs)]
 extern crate proc_macro;
 
-use proc_macro::TokenTree;
 use proc_macro::TokenStream;
+use proc_macro::TokenTree;
+use proc_macro::Spacing;
+use proc_macro::Span;
+use proc_macro::Delimiter;
+use std::iter::FromIterator;
+
+type TT = proc_macro::TokenTree;
+type TS = proc_macro::TokenStream;
+type G = proc_macro::Group;
+type I = proc_macro::Ident;
+type P = proc_macro::Punct;
+
+fn protected_group(fn_body: G) -> TokenTree {
+    let tt: Vec<TT> = vec![
+        P::new(':', Spacing::Joint).into(),
+        P::new(':', Spacing::Alone).into(),
+        I::new("qadapt", Span::call_site()).into(),
+        P::new(':', Spacing::Joint).into(),
+        P::new(':', Spacing::Alone).into(),
+        I::new("enter_protected", Span::call_site()).into(),
+        G::new(Delimiter::Parenthesis, TokenStream::new()).into(),
+        P::new(';', Spacing::Alone).into(),
+        fn_body.into(),
+        P::new(':', Spacing::Joint).into(),
+        P::new(':', Spacing::Alone).into(),
+        I::new("qadapt", Span::call_site()).into(),
+        P::new(':', Spacing::Joint).into(),
+        P::new(':', Spacing::Alone).into(),
+        I::new("exit_protected", Span::call_site()).into(),
+        G::new(Delimiter::Parenthesis, TokenStream::new()).into(),
+        P::new(';', Spacing::Alone).into(),
+    ];
+
+    G::new(Delimiter::Brace, TS::from_iter(tt)).into()
+}
 
 /// Set up the QADAPT allocator to trigger a panic if any allocations happen during
 /// calls to this function.
@@ -18,13 +52,19 @@ use proc_macro::TokenStream;
 /// separate thread, QADAPT will not trigger a panic.
 #[proc_macro_attribute]
 pub fn allocate_panic(_attr: TokenStream, item: TokenStream) -> TokenStream {
-    let ret = item.clone();
+    let mut ret: Vec<TokenTree> = Vec::new();
 
-    let mut token_iter = item.into_iter();
-    match token_iter.next() {
-        Some(TokenTree::Ident(ref i)) if i.to_string() == "fn" => (),
-        _ => panic!("#[allocate_panic] macro can only be applied to functions")
+    let mut fn_body = None;
+    let mut item_iter = item.into_iter();
+    while let Some(tt) = item_iter.next() {
+        match tt {
+            TokenTree::Group(ref g) if g.delimiter() == Delimiter::Brace => {
+                fn_body = Some(g.clone());
+                break;
+            },
+            tt => ret.push(tt)
+        }
     }
-
-    ret
+    ret.push(protected_group(fn_body.unwrap()));
+    TokenStream::from_iter(ret.into_iter())
 }
diff --git a/src/lib.rs b/src/lib.rs
index 44b2338..e617732 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -73,6 +73,11 @@ pub fn exit_protected() {
 
 static INTERNAL_ALLOCATION: RwLock<usize> = RwLock::new(usize::max_value());
 
+/// Get the current "protection level" in QADAPT: calls to enter_protected() - exit_protected()
+pub fn protection_level() -> usize {
+    PROTECTION_LEVEL.try_with(|v| *v.read() ).unwrap_or(0)
+}
+
 fn claim_internal_alloc() {
     loop {
         match INTERNAL_ALLOCATION.write() {
diff --git a/tests/macros.rs b/tests/macros.rs
index 25a9cea..edb6d54 100644
--- a/tests/macros.rs
+++ b/tests/macros.rs
@@ -1,14 +1,21 @@
 extern crate qadapt;
 
+use qadapt::QADAPT;
 use qadapt::allocate_panic;
 
+#[global_allocator]
+static Q: QADAPT = QADAPT;
+
 #[allocate_panic]
 fn allocates() {
-    let _v: Vec<()> = Vec::with_capacity(1);
+    assert_eq!(::qadapt::protection_level(), 1);
+    let mut v = Vec::new();
+    v.push(1);
 }
 
 #[allocate_panic]
 fn no_allocate() {
+    assert_eq!(::qadapt::protection_level(), 1);
     let _v: Vec<()> = Vec::with_capacity(0);
 }
 
@@ -17,10 +24,8 @@ fn test_no_allocate() {
     no_allocate();
 }
 
-/*
 #[test]
 #[should_panic]
 fn test_allocates() {
     allocates();
 }
-*/