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}