function_wrapper/lib.rs
1//! Rust attribute macro library that makes it easy to wrap functions in code that runs before and / or after a function executes.
2//!
3//! Repository: <https://github.com/ChandlerJayCalkins/function-wrapper>
4//!
5//! This function
6//!
7//! ```rust
8//! #[wrap]
9//! fn hello() -> bool
10//! {
11//! println!("Hello there!");
12//! println!("This is some code.");
13//! true
14//! }
15//! ```
16//!
17//! which is being wrapped by this attribute
18//!
19//! ```rust
20//! use function_wrapper::WrappedFn;
21//! extern crate proc_macro;
22//! extern crate proc_macro2;
23//! use syn::parse_macro_input;
24//! use quote::quote;
25//!
26//! // Adds print statements before and after a function executes.
27//! #[proc_macro_attribute]
28//! pub fn wrap(_: proc_macro::TokenStream, item: proc_macro::TokenStream) -> proc_macro::TokenStream
29//! {
30//! // Parse input as a WrappedFn object from the function-wrapper crate.
31//! let mut function = parse_macro_input!(item as WrappedFn);
32//! // Put a print statement in the code that gets run before the function.
33//! function.set_pre_code(quote!{ println!("Hi at the start :)"); });
34//! // Put a print statement in the code that gets run after the function.
35//! function.set_post_code(quote!{ println!("Hi at the end :)"); });
36//! // Convert the function into a TokenStream and return it.
37//! proc_macro2::TokenStream::from(function).into()
38//! }
39//! ```
40//!
41//! will turn into this after being compiled.
42//!
43//! ```rust
44//! fn hello() -> bool
45//! {
46//! println!("Hi at the start :)");
47//! let mut wrapper = ||
48//! {
49//! println!("Hello there!");
50//! println!("This is some code.");
51//! true
52//! };
53//! let result = wrapper();
54//! println!("Hi at the end :)");
55//! result
56//! }
57//! ```
58//!
59//! If only pre-code is added, a wrapper closure and extra return expression won't be added since they are unecessary in this case.
60//! If only post-code is added, the wrapper closure and return expression will still be added out of necessity.
61//!
62
63use proc_macro2::{TokenStream, Span};
64use syn::{ItemFn, Block, Ident, /* Type, ReturnType */};
65use syn::parse::{Parse, ParseStream};
66use syn;
67use quote::{quote, ToTokens};
68use core::iter::Extend;
69
70/// String table of error messages
71const ERROR_STRS: [&str; 1] =
72[
73 // Error message for when no tokens are given to parse in the `syn::parse()` method.
74 "expected function"
75];
76
77// /// Contains the type variants that wrapped function can return.
78// #[derive(Clone)]
79// pub enum WrappedFnOutput
80// {
81// /// No return type was given. Usually represented as `()`.
82// Default,
83// /// All other explicitly written types. Contains a `syn::Type` as the internal value.
84// Type(Type)
85// }
86
87/// Function that can have code inserted before and after the rest of the function executes.
88/// Can be constructed with `syn::parse()` and other variations of parsing from the `syn` crate.
89///
90/// Example:
91///
92/// ```rust
93/// let mut function = parse_macro_input!(token_stream as WrappedFn);
94/// ```
95///
96/// Code that runs before the function can be set using the `set_pre_code()` method.
97///
98/// Example:
99///
100/// ```rust
101/// function.set_pre_code(quote!{ println!("Hi at the start :)"); });
102/// ```
103///
104/// Code that runs after the function can be set using the `set_post_code()` method.
105///
106/// Example:
107///
108/// ```rust
109/// function.set_post_code(quote!{ println!("Hi at the end :)"); });
110/// ```
111#[derive(Clone, Debug)]
112pub struct WrappedFn
113{
114 /// `syn::ItemFn` that contains all of the data of the original function, including the code inside, the function signature, any attributes, etc.
115 pub function: ItemFn,
116 /// Contains code that gets run before the rest of the function.
117 pub pre_code: Option<TokenStream>,
118 /// Contains code that gets run after the rest of the function.
119 pub post_code: Option<TokenStream>,
120 // /// The arguments to the function.
121 // pub args: Vec<FnArgData>,
122 // /// Return type.
123 // pub output: WrappedFnOutput,
124 /// Identifier token for the closure that wraps all of the original code from the wrapped function. `wrapper` by default.
125 pub wrapper_ident: Ident,
126 /// Identifier token for the variable that holds the return value of the wrapped function. `result` by default.
127 pub result_ident: Ident
128}
129
130impl WrappedFn
131{
132 /// Sets the code that gets run before the rest of the function executes.
133 pub fn set_pre_code(&mut self, pre_code: TokenStream)
134 {
135 self.pre_code = Some(pre_code);
136 }
137
138 /// Sets the code that gets run after the rest of the function executes.
139 pub fn set_post_code(&mut self, post_code: TokenStream)
140 {
141 self.post_code = Some(post_code);
142 }
143
144 /// Removes any code that was going to be added before the rest of the function.
145 pub fn remove_pre_code(&mut self)
146 {
147 self.pre_code = None;
148 }
149
150 /// Removes any code that was going to be added after the rest of the function.
151 pub fn remove_post_code(&mut self)
152 {
153 self.post_code = None;
154 }
155
156 /// Changes the identifier for the closure that wraps the code of the original function (`wrapper` by default).
157 pub fn set_wrapper_ident(&mut self, ident: &str)
158 {
159 self.wrapper_ident = Ident::new(ident, Span::call_site());
160 }
161
162 /// Changes the identifier for the variable that holds the value that the function returns (`result` by default).
163 pub fn set_result_ident(&mut self, ident: &str)
164 {
165 self.result_ident = Ident::new(ident, Span::call_site());
166 }
167
168 /// Inserts the unwrapped original code from a function into a function block.
169 ///
170 /// Inputs:
171 ///
172 /// `function_block`: The block of code that goes inside the function where the original code is re-added.
173 ///
174 /// `og_code`: The original code from the function.
175 fn add_unwrapped_code(function_block: &mut TokenStream, og_code: &Block)
176 {
177 // Convert the function's old code block into a TokenStream and add it after the pre code
178 function_block.extend(quote!{ #og_code });
179 }
180
181 /// Wraps the original code of a function in a closure and inserts code after it inside a function block.
182 ///
183 /// Inputs:
184 ///
185 /// `og_code`: The original code of the function that comes before the post code.
186 ///
187 /// `wrapper_ident`: Identifier token for the closure that wraps all of the original code from the wrapped function.
188 ///
189 /// `result_ident`: Identifier token for the variable that holds the return value of the wrapped function.
190 ///
191 /// `function_block`: The block of code that goes inside the function where the wrapper code and post code is added.
192 ///
193 /// `post_code`: The code to be inserted that runs at the end of the function.
194 fn add_wrapped_post_code(og_code: &Block, wrapper_ident: &Ident, result_ident: &Ident, function_block: &mut TokenStream, post_code: &TokenStream)
195 {
196 // Wrap the code in a closure, get the result of running that closure, and turn all of it into a TokenStream
197 let wrapper_code = quote!
198 {
199 let mut #wrapper_ident = || #og_code ;
200 let #result_ident = #wrapper_ident ();
201 };
202 // Get a TokenStream of the return line
203 let return_line = quote!{ #result_ident };
204 // Add the wrapped code that came with the function
205 function_block.extend(wrapper_code);
206 // Add the code that runs after the rest of the function
207 function_block.extend(post_code.clone());
208 // Add the line that returns the return value
209 function_block.extend(return_line);
210 }
211
212 /// Gets a `syn::ItemFn` of a function that just had code inserted into it.
213 ///
214 /// Inputs:
215 ///
216 /// `function`: The original function that is being wrapped.
217 ///
218 /// `function_block`: The new block of code that is replacing the old one inside the function that is being wrapped.
219 ///
220 /// Outputs: A `syn::ItemFn` of the newly wrapped function.
221 fn get_wrapped_function(function: &ItemFn, function_block: &TokenStream) -> ItemFn
222 {
223 // Wrap all of this code inside curly braces
224 let function_block = quote!{ { #function_block } };
225 // Creates a clone of the function
226 let mut new_function = function.clone();
227 // Put the new code block inside the new function
228 new_function.block = syn::parse(function_block.into()).unwrap();
229 // Return the function with the new code
230 new_function
231 }
232}
233
234/// Main way to construct a `WrappedFn`.
235/// Can be constructed using `syn::parse_macro_input` like this:
236///
237/// ```rust
238/// let mut function = parse_macro_input!(token_stream as WrappedFn);
239/// ```
240impl Parse for WrappedFn
241{
242 /// Constructs a WrappedFn from a `syn::ParseStream`.
243 fn parse(input: ParseStream) -> syn::Result<Self>
244 {
245 // If no tokens were given to parse, throw an error
246 if input.is_empty()
247 {
248 return Err(syn::Error::new(input.span(), ERROR_STRS[0]))
249 }
250 // Attempt to parse the input tokens as a function
251 let function: ItemFn = input.parse()?;
252 // Get the return type
253 // let output = match function.sig.output.clone()
254 // {
255 // // If a return type was explicitly given, extract it
256 // ReturnType::Type(_, o) => WrappedFnOutput::Type(*o),
257 // // If no return type was given, use the default return type variant (usually represented as ())
258 // ReturnType::Default => WrappedFnOutput::Default
259 // };
260 // Construct a WrappedFn to return
261 let wrapped_function = Self
262 {
263 function: function,
264 pre_code: None,
265 post_code: None,
266 // output: output,
267 wrapper_ident: Ident::new("wrapper", Span::call_site()),
268 result_ident: Ident::new("result", Span::call_site())
269 };
270 Ok(wrapped_function)
271 }
272}
273
274/// Allows `WrappedFn`s to be converted to `syn::ItemFn`s for easier use in procedural macros.
275impl From<&WrappedFn> for ItemFn
276{
277 /// Converts a `WrappedFn` into a `syn::ItemFn`.
278 fn from(function: &WrappedFn) -> Self
279 {
280 // Determine whether the function has code that gets run before / after the rest of the function or not
281 match (&function.pre_code, &function.post_code)
282 {
283 // If the function has some code to get run both before and after the function
284 (Some(pre_code), Some(post_code)) =>
285 {
286 // Create a new block of code that will replace the old one in the function and start it with the pre code
287 let mut function_block = pre_code.clone();
288 // Wrap and add the original function code and add the post code to the new function block
289 WrappedFn::add_wrapped_post_code(&function.function.block, &function.wrapper_ident, &function.result_ident, &mut function_block, post_code);
290 // Replaces the function's code block with the new one and returns the function
291 WrappedFn::get_wrapped_function(&function.function, &function_block)
292 },
293 // If the function has some code to get run before the function but not after
294 (Some(pre_code), None) =>
295 {
296 // Create a new block of code that will replace the old one in the function and start it with the pre code
297 let mut function_block = pre_code.clone();
298 // Add the function's original code after the pre code
299 WrappedFn::add_unwrapped_code(&mut function_block, &function.function.block);
300 // Replaces the function's code block with the new one and returns the function
301 WrappedFn::get_wrapped_function(&function.function, &function_block)
302 },
303 // If the function has some code to get run after the function but not before
304 (None, Some(post_code)) =>
305 {
306 // Create a new block of code that will replace the old one in the function
307 let mut function_block = TokenStream::new();
308 // Wrap and add the original function code and add the post code to the new function block
309 WrappedFn::add_wrapped_post_code(&function.function.block, &function.wrapper_ident, &function.result_ident, &mut function_block, post_code);
310 // Replaces the function's code block with the new one and returns the function
311 WrappedFn::get_wrapped_function(&function.function, &function_block)
312 },
313 // If the function has no code to insert anywhere
314 (None, None) =>
315 {
316 // Just return the function the way it is
317 function.function.clone()
318 }
319 }
320 }
321}
322
323/// Allows `WrappedFn`s to be converted to `syn::ItemFn`s for easier use in procedural macros.
324impl From<WrappedFn> for ItemFn
325{
326 /// Converts a `WrappedFn` into a `syn::ItemFn`.
327 fn from(function: WrappedFn) -> Self
328 {
329 // Use the `From<&WrappedFn> for ItemFn` implementation
330 Self::from(&function)
331 }
332}
333
334/// Allows `WrappedFn`s to be converted into `proc_macro2::TokenStream`s for easier use in procedural attribute macros.
335impl From<&WrappedFn> for TokenStream
336{
337 /// Converts a `WrappedFn` into a `proc_macro2::TokenStream`.
338 fn from(function: &WrappedFn) -> Self
339 {
340 // Convert the function to a `syn::ItemFn` and then convert that into a `TokenStream`
341 let function = ItemFn::from(function);
342 quote!{ #function }
343 }
344}
345
346/// Allows `WrappedFn`s to be converted into `proc_macro2::TokenStream`s for easier use in procedural attribute macros.
347impl From<WrappedFn> for TokenStream
348{
349 /// Converts a `WrappedFn` into a `proc_macro2::TokenStream`.
350 fn from(function: WrappedFn) -> Self
351 {
352 // Convert the function to a `syn::ItemFn` and then convert that into a TokenStream
353 let function = ItemFn::from(function);
354 quote!{ #function }
355 }
356}
357
358/// Allows `WrappedFn`s to be easily appended to and converted to `proc_macro2::TokenStream`s.
359impl ToTokens for WrappedFn
360{
361 /// Adds a `WrappedFn` to the end of a `proc_macro2::TokenStream`.
362 fn to_tokens(&self, tokens: &mut TokenStream)
363 {
364 // Convert the function into a TokenStream
365 let function: TokenStream = self.into();
366 // Add the function to the end of the TokenStream
367 tokens.extend(function);
368 }
369}