crossroads/
lib.rs

1/*
2 * Copyright (c) 2022 Janosch Reppnow <janoschre+rust@gmail.com>.
3 *
4 * Permission is hereby granted, free of charge, to any person obtaining a copy
5 * of this software and associated documentation files (the "Software"), to deal
6 * in the Software without restriction, including without limitation the rights
7 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
8 * copies of the Software, and to permit persons to whom the Software is
9 * furnished to do so, subject to the following conditions:
10 *
11 * The above copyright notice and this permission notice shall be included in all
12 * copies or substantial portions of the Software.
13 *
14 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
15 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
16 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
17 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
18 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
19 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
20 * SOFTWARE.
21 */
22
23//! A proc macro that turns one function into many - along user-defined forks points.
24//!
25//! # Motivation and Usage
26//! This crate allows you to define multiple functions sharing significant part of their logic.
27//! To understand why this is useful, consider the following example for a set of unit tests
28//! (the ```#[test]``` attribute is only commented out to have these be picked up by doctest):
29//!
30//! ```rust
31//! use std::collections::HashMap;
32//!
33//! // #[test]
34//! fn empty_be_default() {
35//!     let map: HashMap<String, usize> = Default::default();
36//!     
37//!     assert!(map.is_empty());
38//! }
39//!
40//! // #[test]
41//! fn empty_after_clear() {
42//!     let mut map: HashMap<String, usize> = Default::default();
43//!
44//!     map.insert("test".to_string(), 1);
45//!     map.clear();
46//!
47//!     assert!(map.is_empty());
48//! }
49//!
50//! // #[test]
51//! fn empty_after_remove() {
52//!     let mut map: HashMap<String, usize> = Default::default();
53//!
54//!     map.insert("test".to_string(), 1);
55//!     map.remove("test");
56//!
57//!     assert!(map.is_empty());
58//! }
59//! ```
60//!
61//! With this crate, you can write the following instead:
62//!
63//! ```rust
64//! use std::collections::HashMap;
65//! use crossroads::crossroads;
66//!
67//! #[crossroads]
68//! // #[test]
69//! fn empty() {
70//!     let mut map: HashMap<String, usize> = Default::default();
71//!
72//!     match fork!() {
73//!         by_default => {}
74//!         after_add => {
75//!             map.insert("Key".to_owned(), 1337);
76//!             match fork!() {
77//!                 and_remove => map.remove("Key"),
78//!                 and_clear => map.clear(),
79//!             };
80//!         }
81//!     }
82//!
83//!     assert!(map.is_empty());
84//! }
85//! ```
86//!
87//! The ```#[crossroads]``` macro will replace the function with as many functions as there are distinct paths through your fork points.
88//! In this case, it will generate:
89//! ```rust
90//! // #[test]
91//! fn empty_by_default() { /* ... */ }
92//! // #[test]
93//! fn empty_after_add_and_remove() { /* ... */ }
94//! // #[test]
95//! fn empty_after_add_and_clear() { /* ... */ }
96//! ```
97//!
98//! The contents of the methods are the result of replacing the ```match``` expressions with a
99//! block containing the expression specified in the corresponding ```match``` arms.
100//!
101//! You can find the above example in the ```examples``` folder and confirm that it will indeed produce the following output when run as a test:
102//! ```text
103//! running 3 tests
104//! test empty_by_default ... ok
105//! test empty_after_add_and_clear ... ok
106//! test empty_after_add_and_remove ... ok
107//! ```
108//!
109//! # Questions and Answers
110//!
111//! 1. Why did you decide to use the ```match```-based syntax and not implement a new one?
112//! The main reason for using the ```match``` syntax in the way this crate does is to make it as
113//! compatible as possible with code formattting tools such as ```rustfmt```.
114//! See the ```select!``` macros used in the async context for an example of issues a new syntax can cause.
115
116use proc_macro::TokenStream;
117use std::collections::VecDeque;
118
119use syn::__private::{ToTokens, TokenStream2};
120use syn::spanned::Spanned;
121use syn::visit::Visit;
122use syn::visit_mut::VisitMut;
123use syn::{visit, visit_mut, Block, Expr, ExprBlock, Ident, ItemFn, Pat, Stmt};
124
125type Paths<T> = Vec<Vec<T>>;
126
127struct PathFinder {
128    paths: Paths<String>,
129}
130
131impl PathFinder {
132    fn new(paths: Paths<String>) -> Self {
133        Self { paths }
134    }
135
136    fn into_inner(self) -> Paths<String> {
137        self.paths
138    }
139}
140
141impl<'ast> Visit<'ast> for PathFinder {
142    fn visit_expr(&mut self, expr: &'ast Expr) {
143        if match expr {
144            Expr::Match(mtch) => {
145                match mtch.expr.as_ref() {
146                    Expr::Macro(mac) => {
147                        match mac.mac.path.segments.first() {
148                            Some(segment) if segment.ident == "fork" => {
149                                let mut new_paths = Paths::default();
150                                assert!(!mtch.arms.is_empty(), "Must have at least one branch in match branches with fork!()! {:?}", mtch.span());
151                                for arm in &mtch.arms {
152                                    if let Pat::Ident(ident) = &arm.pat {
153                                        let mut this_paths = self.paths.clone();
154                                        for path in &mut this_paths {
155                                            path.push(ident.ident.to_string());
156                                        }
157
158                                        let mut this_pathfinder = PathFinder::new(this_paths);
159                                        this_pathfinder.visit_expr(arm.body.as_ref());
160
161                                        new_paths.append(&mut this_pathfinder.into_inner());
162                                    } else {
163                                        panic!(
164                                            "Must use only idents with a fork!() match! {:?}",
165                                            arm.span()
166                                        );
167                                    }
168                                }
169
170                                self.paths = new_paths;
171                                false
172                            }
173                            _ => true,
174                        }
175                        // TODO: Proper handling of namespace..
176                        // match mac.mac.path.segments.iter().map(|segment| segment.ident.to_string()).collect::<Vec<String>>().as_ref::<[&str]>() {
177                        //     ["fork"] | ["crossroads", "fork"] => {}
178                        //     _ => {}
179                        // }
180                    }
181                    _ => true,
182                }
183            }
184            _ => true,
185        } {
186            visit::visit_expr(self, expr);
187        }
188    }
189}
190
191struct Rewriter {
192    along_path: VecDeque<String>,
193}
194
195impl Rewriter {
196    fn new(path: impl Into<VecDeque<String>>) -> Self {
197        Self {
198            along_path: path.into(),
199        }
200    }
201}
202
203impl VisitMut for Rewriter {
204    fn visit_expr_mut(&mut self, expr: &mut Expr) {
205        if let Some(mut replacement) = if let Expr::Match(mtch) = &expr {
206            match mtch.expr.as_ref() {
207                Expr::Macro(mac) => {
208                    match mac.mac.path.segments.first() {
209                        Some(segment) if segment.ident == "fork" => {
210                            let current = self
211                                .along_path
212                                .pop_front()
213                                .expect("There should always be enough identifiers in this list.");
214                            assert!(!mtch.arms.is_empty(), "Must have at least one branch in match branches with fork!()! {:?}", mtch.span());
215
216                            let mut ret = None;
217
218                            for arm in &mtch.arms {
219                                if let Pat::Ident(ident) = &arm.pat {
220                                    if ident.ident == current {
221                                        ret = Some(Expr::Block(ExprBlock {
222                                            attrs: mtch.attrs.clone(),
223                                            label: None,
224                                            block: Block {
225                                                brace_token: Default::default(),
226                                                stmts: vec![Stmt::Expr(Expr::clone(
227                                                    arm.body.as_ref(),
228                                                ))],
229                                            },
230                                        }));
231                                        break;
232                                    }
233                                } else {
234                                    panic!(
235                                        "Must use only idents with a fork!() match! {:?}",
236                                        arm.span()
237                                    );
238                                }
239                            }
240
241                            if let Some(ret) = ret {
242                                Some(ret)
243                            } else {
244                                panic!("Did not find identifier {} in corresponding match statement. This is almost certainly a bug, please feel free to report it. {:?}", current, mtch.span());
245                            }
246                        }
247                        _ => None,
248                    }
249                    // TODO: Proper handling of namespace..
250                    // match mac.mac.path.segments.iter().map(|segment| segment.ident.to_string()).collect::<Vec<String>>().as_ref::<[&str]>() {
251                    //     ["fork"] | ["crossroads", "fork"] => {}
252                    //     _ => {}
253                    // }
254                }
255                _ => None,
256            }
257        } else {
258            None
259        } {
260            std::mem::swap(expr, &mut replacement);
261            // This is kind of mean: If the expression that we are putting in place of the match is itself another match,
262            // it gets skipped here (as the recursive method assumes you have already visited the node that you give).
263            // As such, we need to manually recurse in this specific case.
264            self.visit_expr_mut(expr);
265        } else {
266            visit_mut::visit_expr_mut(self, expr);
267        }
268    }
269}
270
271/// An attribute macro that can be placed above ```FnItem```s, i.e. freestanding functions everywhere.
272/// It will replace the function with a set of functions induced by the different paths through the
273/// function along the ```match fork!() { a => { ... }, ... }``` points, where the name of the function is induced by the
274/// sequence of the ```identifier``` specified in the patterns of the ```match``` branches used with the for that specific function instance.
275///
276/// See the crate-level documentation for a concrete example.
277#[proc_macro_attribute]
278pub fn crossroads(_: TokenStream, input: TokenStream) -> TokenStream {
279    let function = syn::parse_macro_input!(input as ItemFn);
280
281    let name = function.sig.ident.to_string();
282
283    let mut paths = PathFinder::new(vec![vec![name]]);
284    paths.visit_block(&function.block);
285
286    let paths = paths.into_inner();
287
288    let mut new_functions: Vec<ItemFn> = Vec::with_capacity(paths.len());
289
290    for path in paths {
291        let mut path = VecDeque::from(path);
292        path.pop_front();
293        let mut function = function.clone();
294
295        let mut new_name = function.sig.ident.to_string();
296        for fork in &path {
297            new_name.push('_');
298            new_name.push_str(fork)
299        }
300
301        function.sig.ident = Ident::new(&new_name, function.sig.ident.span());
302
303        let mut rewriter = Rewriter::new(path);
304        rewriter.visit_block_mut(&mut function.block);
305        new_functions.push(function);
306    }
307
308    let mut tokens = TokenStream2::new();
309    for function in new_functions {
310        function.to_tokens(&mut tokens);
311    }
312    tokens.into()
313}