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

cvlr_hook/
lib.rs
