quither_proc_macros/
lib.rs

1// Copyright 2021 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use ::proc_macro::TokenStream;
16use ::proc_macro2::TokenStream as TokenStream2;
17use ::quote::quote;
18use ::std::collections::HashSet;
19use ::syn::spanned::Spanned;
20use ::syn::visit::{Visit, visit_ident, visit_path, visit_path_arguments};
21use ::syn::visit_mut::{
22    VisitMut, visit_expr_match_mut, visit_expr_mut, visit_item_impl_mut, visit_item_mut,
23    visit_item_struct_mut, visit_item_type_mut, visit_path_mut,
24};
25use ::syn::{
26    BinOp, Block, Expr, ExprBinary, ExprBlock, ExprLit, ExprParen, ExprPath, ExprUnary,
27    GenericArgument, GenericParam, Generics, Ident, ImplItem, Item, ItemImpl, ItemStruct, Lit,
28    LitBool, Path, PathArguments, PathSegment, Stmt, Type, TypePath, UnOp, WhereClause,
29    WherePredicate, parse, parse_macro_input,
30};
31
32/// A proc macro to generate code for `quither` crate.
33///
34/// **This macro is for internal use by the `quither` crate. Do not use it directly in your own projects.**
35///
36/// This macro can be used for an attribute macro for `Item`s, like an `impl` block.
37/// Without passing any arguments, this macro copies the annotated item 7 times, with:
38///  - Replacing the path segment `Xither` (without generic parameters) or `Xither<X, Y>` (with 2 generic parameters)
39///    with `Either`, `Neither`, `Both`, `EitherOrNeither`, `EitherOrBoth` or `NeitherOrBoth`.
40///    MAKE SURE TO NOT EXPOSE THE NAME `Xither` TO THE `quither` CRATE'S PUBLIC API OR DOCUMENTATION.
41///  - Replacing the path segment like `Xither<X, Y, e, n, b>`, where `e`, `n`, and `b` are boolean
42///    constants, with corresponding variant types like `Either<X, Y>`, `EitherOrBoth<X, Y>`, etc.
43///    `e`, `n`, and `b` indicate whether the corresponding variant type has `Either`, `Neither`,
44///    or `Both` enum variants.
45///  - Replacing the `has_either`, `has_neither`, and `has_both` expressions with the corresponding
46///    boolean constants (i.e. `true` or `false`), depend on the current copied item's state.
47///  - A syntax subtree inside this macro can also have `quither(...)` attribute to enable / disable
48///    that syntax subtree. The attribute accepts a boolean expression consist of:
49///    - `true` or `false` literals
50///    - `has_either`, `has_neither`, and `has_both` expressions
51///    - `&&` and `||` operators
52///    - `!` operator.
53///  - Attributes `#[either]` / `#[neither]` / `#[both]` works as a shortcut for `#[quither(has_either)]`,
54///    `#[quither(has_neither)]`, and `#[quither(has_both)]` respectively.
55///
56/// The outmost `#[quither]` attribute can take arguments as same as the inner one,
57/// which controls which variant types are generated. For example, if you specify
58/// `#[quither(has_either && !has_neither)]`, then the macro will generate `Either` and `EitherOrBoth`
59/// types only.
60#[proc_macro_attribute]
61pub fn quither(args: TokenStream, input: TokenStream) -> TokenStream {
62    let args_expr_opt: Option<Expr> = (!args.is_empty()).then(|| parse(args).unwrap());
63
64    let ast = parse_macro_input!(input as Item);
65    let mut results = Vec::<TokenStream2>::new();
66    for (has_either, has_neither, has_both) in [
67        (true, false, false),
68        (false, true, false),
69        (false, false, true),
70        (true, true, false),
71        (true, false, true),
72        (false, true, true),
73        (true, true, true),
74    ] {
75        let mut ast = ast.clone();
76        let mut processor = CodeProcessor {
77            has_either,
78            has_neither,
79            has_both,
80        };
81        if let Some(false) = args_expr_opt
82            .as_ref()
83            .and_then(|args_expr| processor.check_quither_condition(&args_expr))
84        {
85            continue;
86        }
87        processor.visit_item_mut(&mut ast);
88
89        let generated_item = quote! { #ast };
90        results.push(generated_item);
91    }
92    quote! {
93        #(#results)*
94    }
95    .into()
96}
97
98struct CodeProcessor {
99    has_either: bool,
100    has_neither: bool,
101    has_both: bool,
102}
103
104impl VisitMut for CodeProcessor {
105    fn visit_path_mut(&mut self, node: &mut Path) {
106        // Type `Xither<L, R>` or `Xither<L, R, has_either, has_neither, has_both>` will be
107        // replaced with something like `EitherOrBoth<L, R>` depend on the bool arguments.
108        for segment in node.segments.iter_mut() {
109            self.replace_quither_path_segment(segment, |ident, new_part| {
110                ident.to_string().replace("Xither", new_part)
111            });
112        }
113        visit_path_mut(self, node);
114    }
115
116    fn visit_expr_mut(&mut self, expr: &mut Expr) {
117        self.replace_has_quither_expr(expr);
118        visit_expr_mut(self, expr);
119    }
120
121    fn visit_item_mut(&mut self, item: &mut Item) {
122        visit_item_mut(self, item);
123    }
124
125    fn visit_item_struct_mut(&mut self, item_struct: &mut ItemStruct) {
126        visit_item_struct_mut(self, item_struct);
127        self.replace_quither_type_definition(&mut item_struct.ident, &mut item_struct.generics);
128    }
129
130    fn visit_item_type_mut(&mut self, item_type: &mut syn::ItemType) {
131        visit_item_type_mut(self, item_type);
132        self.replace_quither_type_definition(&mut item_type.ident, &mut item_type.generics);
133    }
134
135    fn visit_item_impl_mut(&mut self, item_impl: &mut ItemImpl) {
136        visit_item_impl_mut(self, item_impl);
137
138        item_impl.items.retain_mut(|item| {
139            let attr_vec = match item {
140                ImplItem::Fn(item_fn) => &mut item_fn.attrs,
141                ImplItem::Const(item_const) => &mut item_const.attrs,
142                ImplItem::Type(item_type) => &mut item_type.attrs,
143                ImplItem::Macro(item_macro) => &mut item_macro.attrs,
144                _ => return true,
145            };
146            let quither_attr_result =
147                find_first_and_remove_vec_mut(attr_vec, |attr| self.check_attr_is_true(attr));
148            match quither_attr_result {
149                Some(true) => true,
150                Some(false) => false, // Remove the item if the attribute is false.
151                None => true,
152            }
153        });
154
155        // If the `<L, R>` arguments are not found in the impl, we need to remove them.
156        // This can happen when only the `Neither` type is used.
157        let unused_type_params = self
158            .remove_unused_type_params(item_impl)
159            .collect::<Vec<_>>();
160        if let Some(where_clause) = &mut item_impl.generics.where_clause {
161            self.remove_where_clause_for_type_params(where_clause, unused_type_params.iter());
162        }
163    }
164
165    fn visit_expr_match_mut(&mut self, ma: &mut syn::ExprMatch) {
166        visit_expr_match_mut(self, ma);
167        ma.arms.retain_mut(|arm| {
168            let attr_vec = &mut arm.attrs;
169            find_first_and_remove_vec_mut(attr_vec, |attr| self.check_attr_is_true(attr))
170                .unwrap_or(true)
171        });
172    }
173}
174
175impl CodeProcessor {
176    fn replace_quither_path_segment<F>(&self, segment: &mut PathSegment, new_name_gen: F)
177    where
178        F: FnOnce(&str, &str) -> String,
179    {
180        let ident_string = segment.ident.to_string();
181        if !ident_string.contains("Xither") {
182            return;
183        }
184        let Some(bool_args) = self.implicit_345th_bool_arguments_for_path_segment(segment) else {
185            return;
186        };
187        let new_name_part = Self::quither_name_gen(bool_args);
188        segment.ident = Ident::new(
189            &new_name_gen(&ident_string, new_name_part),
190            segment.ident.span(),
191        );
192        if bool_args == (false, true, false) {
193            // For `Neither` type, we need to remove the `<L, R>` arguments.
194            segment.arguments = PathArguments::None
195        } else if let PathArguments::AngleBracketed(syn_args) = &mut segment.arguments {
196            // For other types, we need to keep the `<L, R>` arguments, but remove the
197            // `<has_either, has_neither, has_both>` arguments if available.
198            while syn_args.args.len() > 2 {
199                syn_args.args.pop();
200            }
201            if syn_args.args.is_empty() {
202                segment.arguments = PathArguments::None
203            }
204        }
205    }
206
207    fn replace_quither_type_definition(&self, ident: &mut Ident, params: &mut Generics) {
208        if !self.has_either && !self.has_both {
209            // For `Neither` type, we need to remove the `<L, R>` params after `impl`.
210            params.params.clear();
211            params.where_clause = None;
212        }
213
214        let ident_str = ident.to_string();
215        if let Some(_) = ident_str.find("Xither") {
216            let new_ident_str = ident_str.replace(
217                "Xither",
218                Self::quither_name_gen((self.has_either, self.has_neither, self.has_both)),
219            );
220            *ident = Ident::new(&new_ident_str, ident.span());
221        }
222    }
223
224    fn replace_has_quither_expr(&self, expr: &mut Expr) {
225        let Expr::Path(ExprPath { path, .. }) = expr else {
226            return;
227        };
228        let Some(ident) = path.get_ident() else {
229            return;
230        };
231        let found_value = if ident == "has_either" {
232            Some(self.has_either)
233        } else if ident == "has_neither" {
234            Some(self.has_neither)
235        } else if ident == "has_both" {
236            Some(self.has_both)
237        } else {
238            None
239        };
240        if let Some(found_value) = found_value {
241            *expr = Expr::Lit(ExprLit {
242                lit: Lit::Bool(LitBool {
243                    value: found_value,
244                    span: expr.span(),
245                }),
246                attrs: Vec::new(),
247            });
248        }
249    }
250
251    fn implicit_345th_bool_arguments_for_path_segment(
252        &self,
253        segment: &syn::PathSegment,
254    ) -> Option<(bool, bool, bool)> {
255        if let PathArguments::AngleBracketed(syn_args) = &segment.arguments {
256            let args = syn_args.args.clone().into_pairs().collect::<Vec<_>>();
257            if args.len() == 5 {
258                let has_either = self.generic_argument_as_a_bool(&args[2].value())?;
259                let has_neither = self.generic_argument_as_a_bool(&args[3].value())?;
260                let has_both = self.generic_argument_as_a_bool(&args[4].value())?;
261                return Some((has_either, has_neither, has_both));
262            }
263        }
264        Some((self.has_either, self.has_neither, self.has_both))
265    }
266
267    fn generic_argument_as_a_bool(&self, arg: &GenericArgument) -> Option<bool> {
268        if let GenericArgument::Const(arg_expr) = arg {
269            self.check_quither_condition(&arg_expr)
270        } else if let GenericArgument::Type(Type::Path(TypePath { path, .. })) = arg {
271            self.check_quither_condition_for_path(&path)
272        } else {
273            None
274        }
275    }
276
277    fn remove_unused_type_params(&self, item_impl: &mut ItemImpl) -> impl Iterator<Item = Ident> {
278        let mut type_param_finder = TypeParamFinder::default();
279        type_param_finder.visit_type(&item_impl.self_ty);
280        if let Some((_, trait_path, _)) = &item_impl.trait_ {
281            type_param_finder.visit_path(trait_path);
282        }
283        let (used, unused) = item_impl
284            .generics
285            .params
286            .iter()
287            .cloned()
288            .partition::<Vec<_>, _>(|param| {
289                let GenericParam::Type(tp) = param else {
290                    return false;
291                };
292                type_param_finder.does_appear(&tp.ident)
293            });
294        item_impl.generics.params = used.into_iter().collect();
295        item_impl.generics.params.pop_punct();
296        unused.into_iter().filter_map(|param| {
297            let GenericParam::Type(tp) = param else {
298                return None;
299            };
300            Some(tp.ident)
301        })
302    }
303
304    fn remove_where_clause_for_type_params<'a>(
305        &self,
306        where_clause: &mut WhereClause,
307        unused_type_params: impl Iterator<Item = &'a Ident> + Clone,
308    ) {
309        where_clause.predicates = where_clause
310            .predicates
311            .iter()
312            .cloned()
313            .filter_map(|pred| {
314                let WherePredicate::Type(tp) = &pred else {
315                    return Some(pred);
316                };
317                let mut finder = TypeParamFinder::default();
318                finder.visit_type(&tp.bounded_ty);
319                if unused_type_params
320                    .clone()
321                    .any(|param| finder.does_appear(param))
322                {
323                    // Any of the given unused type params is found in the bounded type,
324                    // so we need to remove this predicate.
325                    None
326                } else {
327                    Some(pred)
328                }
329            })
330            .collect();
331        where_clause.predicates.pop_punct();
332    }
333
334    fn check_attr_is_true(&self, attr: &syn::Attribute) -> Option<bool> {
335        let attr_path = attr.meta.path();
336        if attr_path.is_ident("either") {
337            return Some(self.has_either);
338        } else if attr_path.is_ident("neither") {
339            return Some(self.has_neither);
340        } else if attr_path.is_ident("both") {
341            return Some(self.has_both);
342        } else if attr_path.is_ident("quither") {
343            return self.check_quither_condition(&attr.parse_args().ok()?);
344        } else {
345            return None;
346        }
347    }
348
349    fn check_quither_condition(&self, args: &Expr) -> Option<bool> {
350        match args {
351            Expr::Binary(ExprBinary {
352                left, right, op, ..
353            }) => {
354                let left = self.check_quither_condition(left)?;
355                let right = self.check_quither_condition(right)?;
356                match op {
357                    BinOp::And(_) => Some(left && right),
358                    BinOp::Or(_) => Some(left || right),
359                    _ => None,
360                }
361            }
362            Expr::Unary(ExprUnary {
363                expr,
364                op: UnOp::Not(_),
365                ..
366            }) => self.check_quither_condition(expr).map(|b| !b),
367            Expr::Paren(ExprParen { expr, .. }) => self.check_quither_condition(expr),
368            Expr::Block(ExprBlock {
369                block: Block { stmts, .. },
370                ..
371            }) => {
372                if stmts.len() != 1 {
373                    return None;
374                }
375                let Some(Stmt::Expr(expr, _)) = stmts.first() else {
376                    return None;
377                };
378                self.check_quither_condition(expr)
379            }
380            Expr::Path(ExprPath { path, .. }) => self.check_quither_condition_for_path(path),
381            Expr::Lit(ExprLit {
382                lit: Lit::Bool(LitBool { value, .. }),
383                ..
384            }) => Some(*value),
385            _ => None,
386        }
387    }
388
389    fn check_quither_condition_for_path(&self, path: &Path) -> Option<bool> {
390        if path.is_ident("has_either") {
391            Some(self.has_either)
392        } else if path.is_ident("has_neither") {
393            Some(self.has_neither)
394        } else if path.is_ident("has_both") {
395            Some(self.has_both)
396        } else {
397            None
398        }
399    }
400
401    fn quither_name_gen(bool_args: (bool, bool, bool)) -> &'static str {
402        match bool_args {
403            (true, true, true) => "Quither",
404            (true, true, false) => "EitherOrNeither",
405            (true, false, true) => "EitherOrBoth",
406            (true, false, false) => "Either",
407            (false, true, true) => "NeitherOrBoth",
408            (false, true, false) => "Neither",
409            (false, false, true) => "Both",
410            (false, false, false) => "Unreachable",
411        }
412    }
413}
414
415#[derive(Default, Debug)]
416struct TypeParamFinder {
417    idents: HashSet<Ident>,
418}
419
420impl TypeParamFinder {
421    fn does_appear(&self, ident: &Ident) -> bool {
422        self.idents.contains(ident)
423    }
424}
425
426impl<'ast> Visit<'ast> for TypeParamFinder {
427    fn visit_path(&mut self, path: &'ast Path) {
428        visit_path(self, path);
429        if let Some(ident) = path.get_ident() {
430            self.idents.insert(ident.clone());
431        }
432    }
433
434    fn visit_ident(&mut self, ident: &'ast Ident) {
435        visit_ident(self, ident);
436        self.idents.insert(ident.clone());
437    }
438
439    fn visit_path_segment(&mut self, segment: &'ast PathSegment) {
440        // Do recurse into the type params, but not for the segment's ident
441        // because it's already visited in `visit_path`.
442        // We only need the standalone ident, not the path including `::`.
443        visit_path_arguments(self, &segment.arguments);
444    }
445}
446
447fn find_first_and_remove_vec_mut<T, U, F>(vec: &mut Vec<T>, f: F) -> Option<U>
448where
449    F: Fn(&mut T) -> Option<U>,
450{
451    for (i, item) in vec.iter_mut().enumerate() {
452        if let Some(u) = f(item) {
453            vec.remove(i);
454            return Some(u);
455        }
456    }
457    None
458}
459
460#[test]
461fn test_quither_condition_value() {
462    use ::syn::parse_quote;
463
464    let cp = CodeProcessor {
465        has_either: true,
466        has_neither: false,
467        has_both: true,
468    };
469    assert_eq!(
470        Some(true),
471        cp.check_quither_condition(&parse_quote! { true })
472    );
473    assert_eq!(
474        Some(true),
475        cp.check_quither_condition(&parse_quote! { has_either })
476    );
477    assert_eq!(
478        Some(true),
479        cp.check_quither_condition(&parse_quote! { { true } })
480    );
481    assert_eq!(
482        Some(false),
483        cp.check_quither_condition(&parse_quote! { { has_either && has_neither } })
484    );
485}
486
487#[test]
488fn test_visit_path_mut() {
489    use ::proc_macro2::Span;
490    use ::syn::parse_quote_spanned;
491
492    let mut cp = CodeProcessor {
493        has_either: true,
494        has_neither: false,
495        has_both: true,
496    };
497    let span = Span::call_site();
498
499    let mut path = parse_quote_spanned! { span => Xither<L, R> };
500    cp.visit_path_mut(&mut path);
501    assert_eq!(path, parse_quote_spanned! { span => EitherOrBoth<L, R> });
502
503    let mut path = parse_quote_spanned! { span => Xither<L, R, false, false, true> };
504    cp.visit_path_mut(&mut path);
505    assert_eq!(path, parse_quote_spanned! { span => Both<L, R, > });
506
507    let mut path = parse_quote_spanned! { span => Xither<L, R, has_both, has_both, has_neither> };
508    cp.visit_path_mut(&mut path);
509    assert_eq!(
510        path,
511        parse_quote_spanned! { span => EitherOrNeither<L, R, > }
512    );
513}
quither_proc_macros/lib.rs

quither_proc_macros/
lib.rs
