contructor_derive/
lib.rs

1//! Registers a function to be called before/after main (if an executable)
2//! or when loaded/unloaded (if a dynamic library).
3//!
4//! **Notes**
5//!
6//! Use this library is unsafe unless you want to interop directly with a FFI library.
7//!
8//! Please consider to use the `lazy-static` crate instead of it.
9//!
10//! Usage
11//! =====
12//!
13//! Add the following dependency to your Cargo manifest...
14//!
15//! ```toml
16//! [dependencies]
17//! contructor_derive = "0.1.0"
18//! ```
19//!
20//! Example
21//! =======
22//!
23//! ```
24//! #[macro_use]
25//! extern crate contructor_derive;
26//!
27//! pub static mut RAN: bool = false;
28//!
29//! #[constructor]
30//! extern "C" fn set_ran() {
31//!     unsafe { RAN = true }
32//! }
33//!
34//! #[destructor]
35//! extern "C" fn reset_ran() {
36//!     unsafe { RAN = false }
37//! }
38//!
39//! fn main() {
40//!     assert!(unsafe { RAN });
41//! }
42//! ```
43
44extern crate proc_macro;
45extern crate proc_macro2;
46#[macro_use]
47extern crate quote;
48extern crate syn;
49
50use proc_macro::TokenStream;
51use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
52use syn::{Expr, Item, ItemFn, Lit};
53
54/// Registers a function to be called before main (if an executable) or when loaded (if a dynamic library).
55#[proc_macro_attribute]
56pub fn constructor(args: TokenStream, input: TokenStream) -> TokenStream {
57    let item: Item = syn::parse(input).unwrap();
58
59    if let Item::Fn(ref func) = item {
60        let priority = parse_priority(args);
61
62        gen_ctor(func, priority).into()
63    } else {
64        panic!("constructor!{} is only defined for function!");
65    }
66}
67
68/// Registers a function to be called after main (if an executable) or when unloaded (if a dynamic library).
69#[proc_macro_attribute]
70pub fn destructor(args: TokenStream, input: TokenStream) -> TokenStream {
71    let item: Item = syn::parse(input).unwrap();
72
73    if let Item::Fn(ref func) = item {
74        let priority = parse_priority(args);
75
76        gen_dtor(func, priority).into()
77    } else {
78        panic!("destructor!{} is only defined for function!");
79    }
80}
81
82fn parse_priority(args: TokenStream) -> Option<u64> {
83    if !args.is_empty() {
84        let expr: Expr = syn::parse(args).unwrap();
85
86        if let Expr::Lit(lit) = expr {
87            if let Lit::Int(n) = lit.lit {
88                return Some(n.value());
89            }
90        }
91    }
92
93    None
94}
95
96fn gen_ctor(func: &ItemFn, _priority: Option<u64>) -> TokenStream2 {
97    let mod_name = Ident::new(&format!("{}_ctor", func.ident), Span::call_site());
98    let func_name = &func.ident;
99
100    let ctor = if cfg!(target_os = "linux") {
101        quote! {
102            #[link_section = ".ctors"]
103            #[no_mangle]
104            pub static #func_name: extern fn() = super::#func_name;
105        }
106    } else if cfg!(target_os = "macos") {
107        quote! {
108            #[link_section = "__DATA,__mod_init_func"]
109            #[no_mangle]
110            pub static #func_name: extern fn() = super::#func_name;
111        }
112    } else if cfg!(target_os = "windows") {
113        quote! {
114            #[link_section = ".CRT$XCU"]
115            #[no_mangle]
116            pub static #func_name: extern fn() = super::#func_name;
117        }
118    } else {
119        unimplemented!()
120    };
121
122    quote!{
123        #func
124
125        #[doc(hidden)]
126        pub mod #mod_name {
127            #ctor
128        }
129    }
130}
131
132fn gen_dtor(func: &ItemFn, _priority: Option<u64>) -> TokenStream2 {
133    let mod_name = Ident::new(&format!("{}_dtor", func.ident), Span::call_site());
134    let func_name = &func.ident;
135    let ctor = if cfg!(target_os = "linux") {
136        quote! {
137            #[link_section = ".dtors"]
138            #[no_mangle]
139            pub static #func_name: extern fn() = super::#func_name;
140        }
141    } else if cfg!(target_os = "macos") {
142        quote! {
143            #[link_section = "__DATA,__mod_term_func"]
144            #[no_mangle]
145            pub static #func_name: extern fn() = super::#func_name;
146        }
147    } else if cfg!(target_os = "windows") {
148        quote! {
149            #[link_section = ".CRT$XPU"]
150            #[no_mangle]
151            pub static #func_name: extern fn() = super::#func_name;
152        }
153    } else {
154        unimplemented!()
155    };
156
157    quote!{
158        #func
159
160        #[doc(hidden)]
161        pub mod #mod_name {
162            #ctor
163        }
164    }
165}
contructor_derive/lib.rs

contructor_derive/
lib.rs
