auto_enums/
utils.rs

1// SPDX-License-Identifier: Apache-2.0 OR MIT
2
3use std::{iter, mem};
4
5use proc_macro2::TokenStream;
6use syn::{
7    punctuated::Punctuated, token, visit_mut::VisitMut, Arm, Attribute, Block, Expr, ExprBlock,
8    ExprCall, ExprPath, ExprTuple, ItemFn, Local, Meta, Path, PathSegment, Stmt, StmtMacro,
9};
10
11pub(crate) fn path(segments: impl IntoIterator<Item = PathSegment>) -> Path {
12    Path { leading_colon: None, segments: segments.into_iter().collect() }
13}
14
15pub(crate) fn block(stmts: Vec<Stmt>) -> Block {
16    Block { brace_token: token::Brace::default(), stmts }
17}
18
19pub(crate) fn expr_block(block: Block) -> Expr {
20    Expr::Block(ExprBlock { attrs: vec![], label: None, block })
21}
22
23pub(crate) fn expr_call(attrs: Vec<Attribute>, path: Path, arg: Expr) -> Expr {
24    Expr::Call(ExprCall {
25        attrs,
26        func: Box::new(Expr::Path(ExprPath { attrs: vec![], qself: None, path })),
27        paren_token: token::Paren::default(),
28        args: iter::once(arg).collect(),
29    })
30}
31
32pub(crate) fn unit() -> Expr {
33    Expr::Tuple(ExprTuple {
34        attrs: vec![],
35        paren_token: token::Paren::default(),
36        elems: Punctuated::new(),
37    })
38}
39
40pub(crate) fn replace_expr(this: &mut Expr, f: impl FnOnce(Expr) -> Expr) {
41    *this = f(mem::replace(this, Expr::Verbatim(TokenStream::new())));
42}
43
44pub(crate) fn replace_block(this: &mut Block, f: impl FnOnce(Block) -> Expr) {
45    // `brace_token` of the block that passed to `f` should have `call_site` span.
46    // If `f` generates unused braces containing the span of `this.brace_token`,
47    // this will cause confusing warnings: https://github.com/rust-lang/rust/issues/71080
48    let stmts = mem::take(&mut this.stmts);
49    this.stmts = vec![Stmt::Expr(f(block(stmts)), None)];
50}
51
52pub(crate) fn path_eq(path: &syn::Path, expected_crates: &[&str], expected_path: &[&str]) -> bool {
53    if path.segments.len() == 1 && path.segments[0].ident == expected_path.last().unwrap() {
54        return true;
55    }
56    if path.segments.len() == expected_path.len() + 1 {
57        if !expected_crates.iter().any(|&c| path.segments[0].ident == c) {
58            return false;
59        }
60        for i in 1..path.segments.len() {
61            if path.segments[i].ident != expected_path[i - 1] {
62                return false;
63            }
64        }
65        return true;
66    }
67    false
68}
69
70// -----------------------------------------------------------------------------
71// extension traits
72
73pub(crate) trait VecExt<T> {
74    fn find_remove(&mut self, predicate: impl FnMut(&T) -> bool) -> Option<T>;
75}
76
77impl<T> VecExt<T> for Vec<T> {
78    fn find_remove(&mut self, predicate: impl FnMut(&T) -> bool) -> Option<T> {
79        self.iter().position(predicate).map(|i| self.remove(i))
80    }
81}
82
83// -----------------------------------------------------------------------------
84// node
85
86pub(crate) trait Node {
87    fn visited(&mut self, visitor: &mut impl VisitMut);
88}
89
90impl Node for Stmt {
91    fn visited(&mut self, visitor: &mut impl VisitMut) {
92        visitor.visit_stmt_mut(self);
93    }
94}
95
96impl Node for Local {
97    fn visited(&mut self, visitor: &mut impl VisitMut) {
98        visitor.visit_local_mut(self);
99    }
100}
101
102impl Node for Expr {
103    fn visited(&mut self, visitor: &mut impl VisitMut) {
104        visitor.visit_expr_mut(self);
105    }
106}
107
108impl Node for Arm {
109    fn visited(&mut self, visitor: &mut impl VisitMut) {
110        visitor.visit_arm_mut(self);
111    }
112}
113
114impl Node for Block {
115    fn visited(&mut self, visitor: &mut impl VisitMut) {
116        visitor.visit_block_mut(self);
117    }
118}
119
120impl Node for ItemFn {
121    fn visited(&mut self, visitor: &mut impl VisitMut) {
122        visitor.visit_item_fn_mut(self);
123    }
124}
125
126// -----------------------------------------------------------------------------
127// helper for handling attributes
128
129pub(crate) trait Attrs {
130    fn attrs(&self) -> &[Attribute];
131
132    fn any_attr(&self, ident: &str) -> bool {
133        self.attrs().iter().any(|attr| attr.path().is_ident(ident))
134    }
135
136    fn any_empty_attr(&self, ident: &str) -> bool {
137        self.attrs().iter().any(|attr| matches!(&attr.meta, Meta::Path(p) if p.is_ident(ident)))
138    }
139
140    fn attrs_mut(&mut self) -> Option<&mut Vec<Attribute>>;
141
142    fn find_remove_attr(&mut self, ident: &str) -> Option<Attribute> {
143        self.attrs_mut()?.find_remove(|attr| attr.path().is_ident(ident))
144    }
145}
146
147impl Attrs for Arm {
148    fn attrs(&self) -> &[Attribute] {
149        &self.attrs
150    }
151
152    fn attrs_mut(&mut self) -> Option<&mut Vec<Attribute>> {
153        Some(&mut self.attrs)
154    }
155}
156
157impl Attrs for Local {
158    fn attrs(&self) -> &[Attribute] {
159        &self.attrs
160    }
161
162    fn attrs_mut(&mut self) -> Option<&mut Vec<Attribute>> {
163        Some(&mut self.attrs)
164    }
165}
166
167impl Attrs for StmtMacro {
168    fn attrs(&self) -> &[Attribute] {
169        &self.attrs
170    }
171
172    fn attrs_mut(&mut self) -> Option<&mut Vec<Attribute>> {
173        Some(&mut self.attrs)
174    }
175}
176
177impl Attrs for Stmt {
178    fn attrs(&self) -> &[Attribute] {
179        match self {
180            Stmt::Expr(expr, _) => expr.attrs(),
181            Stmt::Local(local) => local.attrs(),
182            Stmt::Macro(mac) => mac.attrs(),
183            // Ignore nested items.
184            Stmt::Item(_) => &[],
185        }
186    }
187
188    fn attrs_mut(&mut self) -> Option<&mut Vec<Attribute>> {
189        match self {
190            Stmt::Expr(expr, _) => expr.attrs_mut(),
191            Stmt::Local(local) => local.attrs_mut(),
192            Stmt::Macro(mac) => mac.attrs_mut(),
193            // Ignore nested items.
194            Stmt::Item(_) => None,
195        }
196    }
197}
198
199macro_rules! attrs_impl {
200    ($($Expr:ident($Struct:ident),)*) => {
201        impl Attrs for Expr {
202            fn attrs(&self) -> &[Attribute] {
203                // #[cfg_attr(test, deny(non_exhaustive_omitted_patterns))]
204                match self {
205                    $(Expr::$Expr(syn::$Struct { attrs, .. }))|* => &attrs,
206                    _ => &[],
207                }
208            }
209
210            fn attrs_mut(&mut self) -> Option<&mut Vec<Attribute>> {
211                // #[cfg_attr(test, deny(non_exhaustive_omitted_patterns))]
212                match self {
213                    $(Expr::$Expr(syn::$Struct { attrs, .. }))|* => Some(attrs),
214                    _ => None,
215                }
216            }
217        }
218    };
219}
220
221attrs_impl! {
222    Array(ExprArray),
223    Assign(ExprAssign),
224    Async(ExprAsync),
225    Await(ExprAwait),
226    Binary(ExprBinary),
227    Block(ExprBlock),
228    Break(ExprBreak),
229    Call(ExprCall),
230    Cast(ExprCast),
231    Closure(ExprClosure),
232    Const(ExprConst),
233    Continue(ExprContinue),
234    Field(ExprField),
235    ForLoop(ExprForLoop),
236    Group(ExprGroup),
237    If(ExprIf),
238    Index(ExprIndex),
239    Infer(ExprInfer),
240    Let(ExprLet),
241    Lit(ExprLit),
242    Loop(ExprLoop),
243    Macro(ExprMacro),
244    Match(ExprMatch),
245    MethodCall(ExprMethodCall),
246    Paren(ExprParen),
247    Path(ExprPath),
248    Range(ExprRange),
249    Reference(ExprReference),
250    Repeat(ExprRepeat),
251    Return(ExprReturn),
252    Struct(ExprStruct),
253    Try(ExprTry),
254    TryBlock(ExprTryBlock),
255    Tuple(ExprTuple),
256    Unary(ExprUnary),
257    Unsafe(ExprUnsafe),
258    While(ExprWhile),
259    Yield(ExprYield),
260}