profi_attributes/
lib.rs

1use proc_macro::{TokenStream, TokenTree};
2
3#[proc_macro_attribute]
4pub fn profile(_: TokenStream, items: TokenStream) -> TokenStream {
5    let mut items = items.into_iter().collect::<Vec<_>>();
6
7    let Some(TokenTree::Group(body)) = items.pop() else {
8        return r#"compile_error!("Expected function body");"#.parse().unwrap();
9    };
10
11    let cr = match proc_macro_crate::crate_name("profi").unwrap() {
12        proc_macro_crate::FoundCrate::Itself => std::borrow::Cow::Borrowed("profi"),
13        proc_macro_crate::FoundCrate::Name(n) => std::borrow::Cow::Owned(n),
14    };
15    let profile = {
16        use proc_macro::{Delimiter as D, Group, Ident, Punct, Spacing as S, Span};
17
18        [
19            TokenTree::Punct(Punct::new(':', S::Joint)),
20            TokenTree::Punct(Punct::new(':', S::Alone)),
21            TokenTree::Ident(Ident::new(&cr, Span::call_site())),
22            TokenTree::Punct(Punct::new(':', S::Joint)),
23            TokenTree::Punct(Punct::new(':', S::Alone)),
24            TokenTree::Ident(Ident::new("prof", Span::call_site())),
25            TokenTree::Punct(Punct::new('!', S::Alone)),
26            TokenTree::Group(Group::new(D::Parenthesis, TokenStream::new())),
27            TokenTree::Punct(Punct::new(';', S::Alone)),
28            TokenTree::Group(body),
29        ]
30    };
31    let tree = TokenTree::from(proc_macro::Group::new(
32        proc_macro::Delimiter::Brace,
33        TokenStream::from_iter(profile),
34    ));
35    items.push(tree);
36
37    TokenStream::from_iter(items)
38}
39
40#[proc_macro_attribute]
41pub fn main(_: TokenStream, items: TokenStream) -> TokenStream {
42    let mut items = items.into_iter().collect::<Vec<_>>();
43
44    let Some(TokenTree::Group(body)) = items.pop() else {
45        return r#"compile_error!("Expected function body");"#.parse().unwrap();
46    };
47
48    let cr = match proc_macro_crate::crate_name("profi").unwrap() {
49        proc_macro_crate::FoundCrate::Itself => std::borrow::Cow::Borrowed("profi"),
50        proc_macro_crate::FoundCrate::Name(n) => std::borrow::Cow::Owned(n),
51    };
52    let profile = {
53        use proc_macro::{Delimiter as D, Group, Ident, Punct, Spacing as S, Span};
54
55        [
56            TokenTree::Punct(Punct::new(':', S::Joint)),
57            TokenTree::Punct(Punct::new(':', S::Alone)),
58            TokenTree::Ident(Ident::new(&cr, Span::call_site())),
59            TokenTree::Punct(Punct::new(':', S::Joint)),
60            TokenTree::Punct(Punct::new(':', S::Alone)),
61            TokenTree::Ident(Ident::new("print_on_exit", Span::call_site())),
62            TokenTree::Punct(Punct::new('!', S::Alone)),
63            TokenTree::Group(Group::new(D::Parenthesis, TokenStream::new())),
64            TokenTree::Punct(Punct::new(';', S::Alone)),
65            TokenTree::Group(body),
66        ]
67    };
68    let tree = TokenTree::from(proc_macro::Group::new(
69        proc_macro::Delimiter::Brace,
70        TokenStream::from_iter(profile),
71    ));
72    items.push(tree);
73
74    TokenStream::from_iter(items)
75}