cvlr_hook/
lib.rs

1use proc_macro::TokenStream;
2use quote::{quote, ToTokens};
3use syn::{parse_macro_input, punctuated::Punctuated, ItemFn, Meta, Token};
4
5/**
6* This macro is used to insert a hook at the start of a function.
7* # Example
8* #[cvlr_hook_on_entry(hook())]
9  fn t1() {
10      // hook inserted here
11      println!("t1");
12  }
13
14  expands to
15
16  fn t1() {
17      hook();
18      println!("t1");
19  }
20*/
21#[allow(clippy::doc_overindented_list_items)]
22#[proc_macro_attribute]
23pub fn cvlr_hook_on_entry(attr: TokenStream, input: TokenStream) -> TokenStream {
24    // parse the attribute argument
25    let attr = parse_macro_input!(attr with Punctuated::<Meta, Token![,]>::parse_terminated);
26
27    if attr.len() != 1 {
28        return quote! {
29            compile_error!("Expected 1 argument");
30        }
31        .into();
32    }
33
34    let arg = attr.first().unwrap();
35
36    // parse the input tokens and make sure it is a function
37    let mut fn_item = parse_macro_input!(input as ItemFn);
38
39    // insert tokens_start into fn item statements at position 0
40    let tokens_start: syn::Stmt = syn::parse_quote! { #arg; };
41
42    fn_item.block.stmts.insert(0, tokens_start);
43
44    fn_item.into_token_stream().into()
45}
46
47/**
48* This macro is used to insert a hook at the end of a function.
49* If the function returns a value, the hook is inserted before the return statement.
50* # Example
51* #[cvlr_hook_on_exit(hook())]
52  fn t1() {
53      assert_eq!(1, 1);
54      assert_eq!(2, 2);
55      // hook inserted here
56  }
57
58  expands to
59
60  fn t1() {
61      assert_eq!(1, 1);
62      assert_eq!(2, 2);
63      hook()
64  }
65*/
66#[allow(clippy::doc_overindented_list_items)]
67#[proc_macro_attribute]
68pub fn cvlr_hook_on_exit(attr: TokenStream, input: TokenStream) -> TokenStream {
69    // parse the attribute argument
70
71    let attr = parse_macro_input!(attr with Punctuated::<Meta, Token![,]>::parse_terminated);
72    if attr.len() != 1 {
73        return quote! {
74            compile_error!("Expected 1 argument");
75        }
76        .into();
77    }
78
79    let arg = &attr.first().unwrap();
80
81    // parse the input tokens and make sure it is a function
82    let mut fn_item = parse_macro_input!(input as ItemFn);
83    let ret_type = &fn_item.sig.output;
84
85    // create tokens_end
86    let stmt_end: syn::Stmt = syn::parse_quote! { #arg; };
87
88    // len of fn item statements
89    let len = fn_item.block.stmts.len();
90
91    match ret_type {
92        syn::ReturnType::Default => {
93            // insert tokens_end into fn item statements at position len
94            fn_item.block.stmts.insert(len, stmt_end);
95        }
96        syn::ReturnType::Type(_, _) => {
97            // insert tokens_end into fn item statements at position len-1
98            fn_item.block.stmts.insert(len - 1, stmt_end);
99        }
100    }
101
102    fn_item.into_token_stream().into()
103}