auto_enums/auto_enum/
context.rs

1// SPDX-License-Identifier: Apache-2.0 OR MIT
2
3use std::{
4    cell::RefCell, collections::hash_map::DefaultHasher, hash::Hasher as _, iter, mem, thread,
5};
6
7use proc_macro2::TokenStream;
8use quote::format_ident;
9#[cfg(feature = "type_analysis")]
10use syn::Type;
11use syn::{
12    parse::{Parse, ParseStream},
13    parse_quote, Attribute, Error, Expr, Ident, ItemEnum, Macro, Path, Result, Token,
14};
15
16use super::visitor::{Dummy, Visitor};
17use crate::utils::{expr_call, path, replace_expr, unit, Node};
18
19// -----------------------------------------------------------------------------
20// Context
21
22/// Config for related to `visitor::Visitor` type.
23#[derive(Clone, Copy, PartialEq)]
24pub(super) enum VisitMode {
25    Default,
26    Return(/* count */ usize),
27    Try,
28}
29
30/// Config for related to `expr::child_expr`.
31#[derive(Clone, Copy, PartialEq)]
32pub(super) enum VisitLastMode {
33    Default,
34    /*
35    local: `let .. = || {};`
36    or
37    expr: `|| {}`
38    not
39    item_fn: `fn _() -> Fn*() { || {} }`
40    */
41    /// `Stmt::Semi(..)` - never visit last expr
42    Never,
43}
44
45/// The default identifier of expression level marker.
46pub(super) const DEFAULT_MARKER: &str = "marker";
47
48pub(super) struct Context {
49    builder: Builder,
50
51    /// The identifier of the marker macro of the current scope.
52    pub(super) current_marker: String,
53    /// All marker macro identifiers that may have effects on the current scope.
54    markers: Vec<String>,
55
56    // TODO: we may be able to replace some fields based on depth.
57    // depth: isize,
58    /// Currently, this is basically the same as `self.markers.len() == 1`.
59    root: bool,
60    /// This is `true` if other `auto_enum` attribute exists in the current scope.
61    pub(super) has_child: bool,
62
63    pub(super) visit_mode: VisitMode,
64    pub(super) visit_last_mode: VisitLastMode,
65
66    /// Span passed to `syn::Error::new_spanned`.
67    pub(super) span: TokenStream,
68    // - `None`: during checking.
69    // - `Some(None)`: there are no errors.
70    // - `Some(Some)`: there are errors.
71    #[allow(clippy::option_option)]
72    error: RefCell<Option<Option<Error>>>,
73
74    pub(super) args: Vec<Path>,
75    // if "type_analysis" feature is disabled, this field is always empty.
76    traits: Vec<Path>,
77}
78
79impl Context {
80    fn new(
81        span: TokenStream,
82        args: TokenStream,
83        root: bool,
84        mut markers: Vec<String>,
85        diagnostic: Option<Error>,
86    ) -> Result<Self> {
87        let Args { args, marker } = syn::parse2(args)?;
88
89        let current_marker = if let Some(marker) = marker {
90            // Currently, there is no reason to preserve the span, so convert `Ident` to `String`.
91            // This should probably be more efficient than calling `to_string` for each comparison.
92            // https://github.com/dtolnay/proc-macro2/blob/1.0.86/src/wrapper.rs#L723
93            let marker_string = marker.to_string();
94            if markers.contains(&marker_string) {
95                bail!(
96                    marker,
97                    "a custom marker name is specified that duplicated the name already used in the parent scope",
98                );
99            }
100            marker_string
101        } else {
102            DEFAULT_MARKER.to_owned()
103        };
104
105        markers.push(current_marker.clone());
106
107        Ok(Self {
108            builder: Builder::new(&span),
109            current_marker,
110            markers,
111            root,
112            has_child: false,
113            visit_mode: VisitMode::Default,
114            visit_last_mode: VisitLastMode::Default,
115            span,
116            error: RefCell::new(Some(diagnostic)),
117            args,
118            traits: vec![],
119        })
120    }
121
122    /// Make a new `Context` as a root.
123    pub(super) fn root(span: TokenStream, args: TokenStream) -> Result<Self> {
124        Self::new(span, args, true, Vec::with_capacity(1), None)
125    }
126
127    /// Make a new `Context` as a child based on a parent context `self`.
128    pub(super) fn make_child(&mut self, span: TokenStream, args: TokenStream) -> Result<Self> {
129        debug_assert!(self.has_child);
130        Self::new(
131            span,
132            args,
133            false,
134            mem::take(&mut self.markers),
135            self.error.borrow_mut().as_mut().unwrap().take(),
136        )
137    }
138
139    /// Merge a child `Context` into a parent context `self`.
140    pub(super) fn join_child(&mut self, mut child: Self) {
141        debug_assert!(self.markers.is_empty());
142        child.markers.pop();
143        mem::swap(&mut self.markers, &mut child.markers);
144
145        if let Some(message) = child.error.borrow_mut().take().unwrap() {
146            self.error(message);
147        }
148    }
149
150    pub(super) fn error(&self, message: Error) {
151        match self.error.borrow_mut().as_mut().unwrap() {
152            Some(base) => base.combine(message),
153            error @ None => *error = Some(message),
154        }
155    }
156
157    pub(super) fn check(self) -> Result<()> {
158        match self.error.borrow_mut().take().unwrap() {
159            Some(e) => Err(e),
160            None => Ok(()),
161        }
162    }
163
164    /// Returns `true` if one or more errors occurred.
165    pub(super) fn has_error(&self) -> bool {
166        self.error.borrow().as_ref().unwrap().is_some()
167    }
168
169    pub(super) fn visit_last(&self) -> bool {
170        self.visit_last_mode != VisitLastMode::Never && self.visit_mode != VisitMode::Try
171    }
172
173    /// Even if this is `false`, there are cases where this `auto_enum` attribute is handled as a
174    /// dummy. e.g., If `self.has_child && self.builder.variants.is_empty()` is true, this
175    /// `auto_enum` attribute is handled as a dummy.
176    pub(super) fn is_dummy(&self) -> bool {
177        // `auto_enum` attribute with no argument is handled as a dummy.
178        // if "type_analysis" feature is disabled, `self.traits` field is always empty.
179        self.args.is_empty() && self.traits.is_empty()
180    }
181
182    #[cfg(feature = "type_analysis")]
183    pub(super) fn variant_is_empty(&self) -> bool {
184        self.builder.variants.is_empty()
185    }
186
187    /// Returns `true` if `expr` is the marker macro that may have effects on the current scope.
188    pub(super) fn is_marker_expr(&self, expr: &Expr) -> bool {
189        match expr {
190            Expr::Macro(expr) => self.is_marker_macro(&expr.mac),
191            _ => false,
192        }
193    }
194
195    /// Returns `true` if `mac` is the marker macro that may have effects on the current scope.
196    pub(super) fn is_marker_macro(&self, mac: &Macro) -> bool {
197        let exact = self.is_marker_macro_exact(mac);
198        if exact || self.root {
199            return exact;
200        }
201
202        self.markers.iter().any(|marker| mac.path.is_ident(marker))
203    }
204
205    /// Returns `true` if `mac` is the marker macro of the current scope.
206    pub(super) fn is_marker_macro_exact(&self, mac: &Macro) -> bool {
207        mac.path.is_ident(&self.current_marker)
208    }
209
210    /// from `<expr>` into `Enum::VariantN(<expr>)`
211    pub(super) fn next_expr(&mut self, expr: Expr) -> Expr {
212        self.next_expr_with_attrs(vec![], expr)
213    }
214
215    /// from `<expr>` into `<attrs> Enum::VariantN(<expr>)`
216    pub(super) fn next_expr_with_attrs(&mut self, attrs: Vec<Attribute>, expr: Expr) -> Expr {
217        self.builder.next_expr(attrs, expr)
218    }
219
220    pub(super) fn replace_boxed_expr(&mut self, expr: &mut Option<Box<Expr>>) {
221        replace_expr(expr.get_or_insert_with(|| Box::new(unit())), |expr| {
222            if self.is_marker_expr(&expr) {
223                // Skip if `<expr>` is a marker macro.
224                expr
225            } else {
226                self.next_expr(expr)
227            }
228        });
229    }
230
231    // visitors
232
233    pub(super) fn visitor(&mut self, node: &mut impl Node) {
234        node.visited(&mut Visitor::new(self));
235    }
236
237    pub(super) fn dummy(&mut self, node: &mut impl Node) {
238        debug_assert!(self.args.is_empty());
239
240        node.visited(&mut Dummy::new(self));
241    }
242
243    // build
244
245    pub(super) fn build(&self, f: impl FnOnce(ItemEnum)) {
246        // As we know that an error will occur, it does not matter if there are not enough variants.
247        if !self.has_error() {
248            match self.builder.variants.len() {
249                1 => {}
250                0 if !self.has_child => {}
251                _ => {
252                    if !self.builder.variants.is_empty() {
253                        f(self.builder.build(&self.args, &self.traits));
254                    }
255                    return;
256                }
257            }
258
259            let (msg1, msg2) = match self.visit_last_mode {
260                VisitLastMode::Default => {
261                    ("branches or marker macros in total", "branch or marker macro")
262                }
263                VisitLastMode::Never => ("marker macros", "marker macro"),
264            };
265            self.error(format_err!(
266                self.span,
267                "`#[auto_enum]` is required two or more {}, there is {} {} in this statement",
268                msg1,
269                if self.builder.variants.is_empty() { "no" } else { "only one" },
270                msg2
271            ));
272        }
273    }
274
275    // type_analysis feature
276
277    #[cfg(feature = "type_analysis")]
278    pub(super) fn collect_impl_trait(&mut self, ty: &mut Type) -> bool {
279        super::type_analysis::collect_impl_trait(&self.args, &mut self.traits, ty)
280    }
281}
282
283impl Drop for Context {
284    fn drop(&mut self) {
285        if !thread::panicking() && self.error.borrow().is_some() {
286            panic!("context need to be checked");
287        }
288    }
289}
290
291// -----------------------------------------------------------------------------
292// Args
293
294mod kw {
295    syn::custom_keyword!(marker);
296}
297
298struct Args {
299    args: Vec<Path>,
300    marker: Option<Ident>,
301}
302
303impl Parse for Args {
304    fn parse(input: ParseStream<'_>) -> Result<Self> {
305        let mut args = Vec::with_capacity((!input.is_empty()) as usize);
306        let mut marker = None;
307        while !input.is_empty() {
308            if input.peek(kw::marker) && input.peek2(Token![=]) {
309                let i: kw::marker = input.parse()?;
310                let _: Token![=] = input.parse()?;
311                let ident: Ident = input.parse()?;
312                if marker.replace(ident).is_some() {
313                    bail!(i, "duplicate `marker` argument");
314                }
315            } else {
316                args.push(input.parse()?);
317            }
318
319            if input.is_empty() {
320                break;
321            }
322            let _: Token![,] = input.parse()?;
323        }
324
325        Ok(Self { args, marker })
326    }
327}
328
329// -----------------------------------------------------------------------------
330// Enum builder
331
332struct Builder {
333    ident: Ident,
334    variants: Vec<Ident>,
335}
336
337impl Builder {
338    fn new(input: &TokenStream) -> Self {
339        Self { ident: format_ident!("__Enum{}", hash(input)), variants: Vec::with_capacity(2) }
340    }
341
342    fn next_expr(&mut self, attrs: Vec<Attribute>, expr: Expr) -> Expr {
343        let variant = format_ident!("__Variant{}", self.variants.len());
344
345        let path =
346            path(iter::once(self.ident.clone().into()).chain(iter::once(variant.clone().into())));
347
348        self.variants.push(variant);
349
350        expr_call(attrs, path, expr)
351    }
352
353    fn build(&self, args: &[Path], traits: &[Path]) -> ItemEnum {
354        let derive = args.iter().chain(traits);
355        let ident = &self.ident;
356        let ty_generics = &self.variants;
357        let variants = &self.variants;
358        let fields = &self.variants;
359
360        parse_quote! {
361            #[allow(non_camel_case_types)]
362            #[::auto_enums::enum_derive(#(#derive),*)]
363            enum #ident<#(#ty_generics),*> {
364                #(#variants(#fields),)*
365            }
366        }
367    }
368}
369
370/// Returns the hash value of the input AST.
371fn hash(input: &TokenStream) -> u64 {
372    let mut hasher = DefaultHasher::new();
373    hasher.write(input.to_string().as_bytes());
374    hasher.finish()
375}