qadapt_macro/
lib.rs

1//! Helper macros to use with the QADAPT allocator system
2//!
3//! **Please note**: This crate has been deprecated in favor of [alloc-counter](https://crates.io/crates/alloc_counter)
4//!
5//! This crate is intended for managing the QADAPT allocator,
6//! and is unusable on its own.
7//!
8// TODO: This causes issues, but I can't track down why
9// #![deny(missing_docs)]
10extern crate proc_macro;
11
12use proc_macro::Delimiter;
13use proc_macro::Group;
14use proc_macro::Spacing;
15use proc_macro::TokenStream;
16use proc_macro::TokenTree;
17use std::iter::FromIterator;
18
19macro_rules! group {
20    ($delim:expr, $ts:expr) => {{
21        let _tt: TokenTree = ::proc_macro::Group::new($delim, $ts).into();
22        _tt
23    }};
24    ($delim:expr) => {
25        group!($delim, ::proc_macro::TokenStream::new())
26    };
27}
28
29macro_rules! ident {
30    ($name:expr, $span:expr) => {{
31        let _tt: TokenTree = ::proc_macro::Ident::new($name, $span).into();
32        _tt
33    }};
34    ($name:expr) => {
35        ident!($name, ::proc_macro::Span::call_site())
36    };
37}
38
39macro_rules! punct {
40    ($ch:expr, $spacing:expr) => {{
41        let _tt: TokenTree = ::proc_macro::Punct::new($ch, $spacing).into();
42        _tt
43    }};
44}
45
46macro_rules! token_stream {
47    ($($tt:expr,)*) => {
48        {
49            let _v: Vec<::proc_macro::TokenTree> = vec![$($tt),*];
50            let _ts: TokenStream = ::proc_macro::TokenStream::from_iter(_v.into_iter());
51            _ts
52        }
53    };
54    ($($tt:expr),*) => {
55        {
56            let _v: Vec<::proc_macro::TokenTree> = vec![$($tt),*];
57            let _ts: TokenStream = ::proc_macro::TokenStream::from_iter(_v.into_iter());
58            _ts
59        }
60    };
61}
62
63/// Generate the body of a function that is protected from making allocations.
64/// The code is conditionally compiled so that all QADAPT-related bits
65/// will be removed for release/bench builds, making the proc_macro safe
66/// to leave on in production.
67#[rustfmt::skip]
68fn protected_body(fn_body: Group) -> TokenTree {
69    group!(Delimiter::Brace, token_stream!(
70        group!(Delimiter::Brace, token_stream!(
71            punct!(':', Spacing::Joint),
72            punct!(':', Spacing::Alone),
73            ident!("qadapt"),
74            punct!(':', Spacing::Joint),
75            punct!(':', Spacing::Alone),
76            ident!("enter_protected"),
77            group!(Delimiter::Parenthesis)
78        )),
79        ident!("let"),
80        ident!("__ret__"),
81        punct!('=', Spacing::Alone),
82        fn_body.into(),
83        punct!(';', Spacing::Alone),
84        punct!('#', Spacing::Alone),
85        // When `return` statements are involved, this code can get marked as
86        // unreachable because of early exit
87        group!(Delimiter::Bracket, token_stream!(
88            ident!("allow"),
89            group!(Delimiter::Parenthesis, token_stream!(
90                ident!("unreachable_code")
91            ))
92        )),
93        group!(Delimiter::Brace, token_stream!(
94            group!(Delimiter::Brace, token_stream!(
95                punct!(':', Spacing::Joint),
96                punct!(':', Spacing::Alone),
97                ident!("qadapt"),
98                punct!(':', Spacing::Joint),
99                punct!(':', Spacing::Alone),
100                ident!("exit_protected"),
101                group!(Delimiter::Parenthesis)
102            )),
103            ident!("__ret__")
104        ))
105    ))
106}
107
108/// Walk through a TokenStream (typically from a Group Brace) and prepend calls
109/// to `return` with an exit guard.
110fn escape_return(ts: TokenStream) -> TokenStream {
111    let mut protected: Vec<TokenTree> = Vec::new();
112    let mut in_closure: bool = false;
113
114    let mut tt_iter = ts.into_iter();
115    while let Some(tt) = tt_iter.next() {
116        let tokens = match tt {
117            TokenTree::Group(ref g) if g.delimiter() == Delimiter::Brace && !in_closure => {
118                vec![group!(Delimiter::Brace, escape_return(g.stream()))]
119            }
120            TokenTree::Ident(ref i) if i.to_string() == "return" && !in_closure => vec![
121                group!(
122                    Delimiter::Brace,
123                    token_stream!(
124                        punct!(':', Spacing::Joint),
125                        punct!(':', Spacing::Alone),
126                        ident!("qadapt"),
127                        punct!(':', Spacing::Joint),
128                        punct!(':', Spacing::Alone),
129                        ident!("exit_protected"),
130                        group!(Delimiter::Parenthesis)
131                    )
132                ),
133                tt.clone(),
134            ],
135            TokenTree::Punct(ref p) if p.as_char() == '|' => {
136                in_closure = true;
137                vec![tt.clone()]
138            }
139            TokenTree::Punct(ref p) if p.as_char() == ';' => {
140                in_closure = false;
141                vec![tt.clone()]
142            }
143            t => vec![t],
144        };
145
146        protected.extend(tokens.into_iter());
147    }
148
149    TokenStream::from_iter(protected.into_iter())
150}
151
152/// Set up the QADAPT allocator to trigger a panic if any allocations happen during
153/// calls to this function.
154///
155/// QADAPT will only track allocations in the current function call;
156/// if (for example) this function receives the results of an allocation in a
157/// separate thread, or defers allocations via closure/Future, those results
158/// will not trigger an error.
159#[proc_macro_attribute]
160#[deprecated(
161    since = "1.0.3",
162    note = "Please use the `alloc_counter` crate instead."
163)]
164pub fn no_alloc(_attr: TokenStream, item: TokenStream) -> TokenStream {
165    let mut protected_fn: Vec<TokenTree> = Vec::new();
166    let mut item_iter = item.into_iter();
167
168    // First, get the function body we're replicating
169    let mut fn_body = None;
170    while let Some(tt) = item_iter.next() {
171        match tt {
172            TokenTree::Group(ref g) if g.delimiter() == Delimiter::Brace => {
173                fn_body = Some(Group::new(Delimiter::Brace, escape_return(g.stream())));
174                break;
175            }
176            tt => {
177                protected_fn.push(tt.clone());
178            }
179        }
180    }
181
182    protected_fn.push(protected_body(fn_body.as_ref().unwrap().clone()));
183
184    while let Some(tt) = item_iter.next() {
185        protected_fn.push(tt)
186    }
187
188    TokenStream::from_iter(protected_fn.into_iter())
189}
qadapt_macro/lib.rs

qadapt_macro/
lib.rs
