Make proc_macro suck a lot less

.gitignore

@@ -1,3 +1,4 @@
 /target
 **/*.rs.bk
 Cargo.lock
+*.swp

qadapt-macro/src/lib.rs

@@ -8,77 +8,102 @@
 extern crate proc_macro;
 
 use proc_macro::Delimiter;
+use proc_macro::Group;
 use proc_macro::Spacing;
-use proc_macro::Span;
 use proc_macro::TokenStream;
 use proc_macro::TokenTree;
 use std::iter::FromIterator;
 
-type TT = proc_macro::TokenTree;
+macro_rules! group {
-type TS = proc_macro::TokenStream;
+    ($delim:expr, $ts:expr) => {
-type G = proc_macro::Group;
+        {
-type I = proc_macro::Ident;
+            let _tt: TokenTree = ::proc_macro::Group::new($delim, $ts).into();
-type P = proc_macro::Punct;
+            _tt
-
+        }
-fn release_guard(fn_name: &str) -> TokenStream {
+    };
-    let rel = I::new("release", Span::call_site());
+    ($delim:expr) => {
-    let not_rel: Vec<TokenTree> = vec![
+        group!($delim, ::proc_macro::TokenStream::new())
-        I::new("not", Span::call_site()).into(),
+    };
-        G::new(Delimiter::Parenthesis, TokenTree::Ident(rel).into()).into()
-    ];
-    let cfg_not_rel: Vec<TokenTree> = vec![
-        I::new("cfg", Span::call_site()).into(),
-        G::new(Delimiter::Parenthesis, TS::from_iter(not_rel.into_iter())).into()
-    ];
-    let guarded: Vec<TokenTree> = vec![
-        P::new('#', Spacing::Alone).into(),
-        G::new(Delimiter::Bracket, TS::from_iter(cfg_not_rel.into_iter())).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(fn_name, Span::call_site()).into(),
-        G::new(Delimiter::Parenthesis, TokenStream::new()).into(),
-    ];
-
-    TS::from_iter(guarded.into_iter())
 }
 
-fn protected_body(fn_name: &str, args: G) -> TokenTree {
+macro_rules! ident {
-    let mut args_filtered = Vec::new();
+    ($name:expr, $span:expr) => {
-    let mut waiting_for_comma = false;
+        {
-    let mut in_type = 0;
+            let _tt: TokenTree = ::proc_macro::Ident::new($name, $span).into();
-    for tt in args.stream().into_iter() {
+            _tt
-        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));
+    ($name:expr) => {
+        ident!($name, ::proc_macro::Span::call_site())
+    };
+}
 
-    let tt: Vec<TT> = vec![
+macro_rules! punct {
-        G::new(Delimiter::Brace, release_guard("enter_protected")).into(),
+    ($ch:expr, $spacing:expr) => {
-        I::new("let", Span::call_site()).into(),
+        {
-        I::new("__ret__", Span::call_site()).into(),
+            let _tt: TokenTree = ::proc_macro::Punct::new($ch, $spacing).into();
-        P::new('=', Spacing::Alone).into(),
+            _tt
-        I::new(fn_name, Span::call_site()).into(),
+        }
-        args_group.into(),
+    };
-        P::new(';', Spacing::Alone).into(),
+}
-        G::new(Delimiter::Brace, release_guard("exit_protected")).into(),
-        I::new("__ret__", Span::call_site()).into(),
-    ];
 
-    G::new(Delimiter::Brace, TS::from_iter(tt)).into()
+macro_rules! token_stream {
+    ($($tt:expr,)*) => {
+        {
+            let _v: Vec<::proc_macro::TokenTree> = vec![$($tt),*];
+            let _ts: TokenStream = ::proc_macro::TokenStream::from_iter(_v.into_iter());
+            _ts
+        }
+    };
+    ($($tt:expr),*) => {
+        {
+            let _v: Vec<::proc_macro::TokenTree> = vec![$($tt),*];
+            let _ts: TokenStream = ::proc_macro::TokenStream::from_iter(_v.into_iter());
+            _ts
+        }
+    };
+}
+
+fn release_guard(fn_name: &str) -> TokenStream {
+    // #[cfg(any(debug, test))]
+    // { ::qadapt::`fn_name`() }
+    token_stream!(
+        punct!('#', Spacing::Alone),
+        group!(Delimiter::Bracket, token_stream!(
+            ident!("cfg"),
+            group!(Delimiter::Parenthesis, token_stream!(
+                ident!("any"),
+                group!(Delimiter::Parenthesis, token_stream!(
+                    ident!("debug"),
+                    punct!(',', Spacing::Alone),
+                    ident!("test")
+                )),
+            )),
+        )),
+        group!(Delimiter::Brace, token_stream!(
+            punct!(':', Spacing::Joint),
+            punct!(':', Spacing::Alone),
+            ident!("qadapt"),
+            punct!(':', Spacing::Joint),
+            punct!(':', Spacing::Alone),
+            ident!(fn_name),
+            group!(Delimiter::Parenthesis)
+        ))
+    )
+}
+
+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"),
+        ident!("__ret__"),
+        punct!('=', Spacing::Alone),
+        fn_body.into(),
+        punct!(';', Spacing::Alone),
+        group!(Delimiter::Brace, release_guard("exit_protected")),
+        ident!("__ret__")
+    ))
 }
 
 /// Set up the QADAPT allocator to trigger a panic if any allocations happen during
@@ -89,48 +114,29 @@ fn protected_body(fn_name: &str, args: G) -> TokenTree {
 /// separate thread, QADAPT will not trigger a panic.
 #[proc_macro_attribute]
 pub fn allocate_panic(_attr: TokenStream, item: TokenStream) -> TokenStream {
-    let mut original_fn: Vec<TokenTree> = Vec::new();
     let mut protected_fn: Vec<TokenTree> = Vec::new();
 
     let mut item_iter = item.into_iter();
 
-    // First, get the function name we're replacing
+    // First, get the function body we're replicating
-    let mut fn_name = None;
+    let mut fn_body = None;
-    let mut fn_args = None;
     while let Some(tt) = item_iter.next() {
         match tt {
-            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());
+                fn_body = Some(g.clone());
-                protected_fn.push(protected_body(
+                break;
-                    &fn_name.take().unwrap(),
-                    fn_args.take().unwrap(),
-                ));
             }
             tt => {
-                original_fn.push(tt.clone());
                 protected_fn.push(tt.clone());
             }
         }
     }
 
-    let mut full = Vec::new();
+    protected_fn.push(protected_body(fn_body.as_ref().unwrap().clone()));
-    full.push(TS::from_iter(original_fn));
+
-    full.push(TS::from_iter(protected_fn));
+    while let Some(tt) = item_iter.next() {
-    let ts = TS::from_iter(full.into_iter());
+        protected_fn.push(tt)
-    ts
+    }
+
+    TokenStream::from_iter(protected_fn.into_iter())
 }