VisitMut<T> and VisitMutwith<F> (#680)

This is a groundwork for dts generator.

common/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "swc_common"
-version = "0.5.3"
+version = "0.5.4"
 authors = ["강동윤 <kdy1997.dev@gmail.com>"]
 license = "Apache-2.0/MIT"
 repository = "https://github.com/swc-project/swc.git"

common/src/fold.rs

@@ -48,6 +48,27 @@ pub trait Visit<T: ?Sized> {
     }
 }
 
+/// Visitor based on a type system.
+///
+/// This trait requires `#![feature(specialization)]`.
+pub trait VisitMut<T: ?Sized> {
+    fn visit_mut(&mut self, node: &mut T);
+
+    /// Creates a folder which applies `folder` after `self`.
+    fn then<F>(self, visitor: F) -> AndThen<Self, F>
+    where
+        Self: Sized,
+        F: VisitMut<T>,
+    {
+        AndThen {
+            first: self,
+            second: visitor,
+        }
+    }
+}
+
+// ----- ----- impl for Box<F> ----- -----
+
 impl<T, F: ?Sized> Fold<T> for Box<F>
 where
     T: FoldWith<Self>,
@@ -58,6 +79,16 @@ where
     }
 }
 
+impl<T: ?Sized, F: ?Sized> VisitMut<T> for Box<F>
+where
+    T: VisitMutWith<Self>,
+    F: VisitMut<T>,
+{
+    fn visit_mut(&mut self, node: &mut T) {
+        (**self).visit_mut(node)
+    }
+}
+
 impl<T: ?Sized, F: ?Sized> Visit<T> for Box<F>
 where
     T: VisitWith<Self>,
@@ -68,6 +99,8 @@ where
     }
 }
 
+// ----- ----- impl for &'a mut F ----- -----
+
 impl<'a, T, F: ?Sized> Fold<T> for &'a mut F
 where
     T: FoldWith<Self>,
@@ -78,6 +111,16 @@ where
     }
 }
 
+impl<'a, T, F: ?Sized> VisitMut<T> for &'a mut F
+where
+    T: VisitMutWith<Self>,
+    F: VisitMut<T>,
+{
+    fn visit_mut(&mut self, node: &mut T) {
+        (**self).visit_mut(node)
+    }
+}
+
 impl<'a, T: ?Sized, F: ?Sized> Visit<T> for &'a mut F
 where
     T: VisitWith<Self>,
@@ -88,6 +131,8 @@ where
     }
 }
 
+// ----- ----- default impl for F ----- -----
+
 impl<T, F> Fold<T> for F
 where
     T: FoldWith<F>,
@@ -97,6 +142,15 @@ where
     }
 }
 
+impl<T: ?Sized, F> VisitMut<T> for F
+where
+    T: VisitMutWith<F>,
+{
+    default fn visit_mut(&mut self, t: &mut T) {
+        t.visit_mut_children(self)
+    }
+}
+
 impl<T: ?Sized, F> Visit<T> for F
 where
     T: VisitWith<F>,
@@ -109,7 +163,7 @@ where
 /// Trait implemented for types which know how to fold itself.
 ///
 ///
-///#Derive
+/// # Derive
 ///
 /// This trait can be derived with `#[derive(Fold)]`.
 ///
@@ -136,7 +190,34 @@ pub trait FoldWith<F>: Sized {
 /// Trait implemented for types which know how to visit itself.
 ///
 ///
-///#Derive
+/// # Derive
+///
+/// This trait can be derived with `#[derive(Fold)]`.
+///
+/// Note that derive ignores all fields with primitive type
+/// because it would encourage mistakes. Use new type instead.
+///
+/// `#[fold(ignore)]` can be used to ignore a field.
+pub trait VisitMutWith<F> {
+    /// This is used by default implementation of `Fold<Self>::fold`.
+    fn visit_mut_children(&mut self, f: &mut F);
+
+    /// Call `f.fold(self)`.
+    ///
+    /// This bypasses a type inference bug which is caused by specialization.
+
+    fn visit_mut_with(&mut self, f: &mut F)
+    where
+        F: VisitMut<Self>,
+    {
+        f.visit_mut(self)
+    }
+}
+
+/// Trait implemented for types which know how to visit itself.
+///
+///
+/// # Derive
 ///
 /// This trait can be derived with `#[derive(Fold)]`.
 ///
@@ -157,6 +238,8 @@ pub trait VisitWith<F> {
     }
 }
 
+// ----- ----- impl for &T ----- -----
+
 impl<'a, T: ?Sized, F> VisitWith<F> for &'a T
 where
     F: Visit<T>,
@@ -166,6 +249,19 @@ where
     }
 }
 
+// ----- ----- impl for &mut T ----- -----
+
+impl<'a, T: ?Sized, F> VisitMutWith<F> for &'a mut T
+where
+    F: VisitMut<T>,
+{
+    fn visit_mut_children(&mut self, f: &mut F) {
+        f.visit_mut(&mut **self)
+    }
+}
+
+// ----- ----- impl for Box<T> ----- -----
+
 impl<T, F> FoldWith<F> for Box<T>
 where
     F: Fold<T>,
@@ -175,6 +271,15 @@ where
     }
 }
 
+impl<T, F> VisitMutWith<F> for Box<T>
+where
+    F: VisitMut<T>,
+{
+    fn visit_mut_children(&mut self, f: &mut F) {
+        f.visit_mut(self)
+    }
+}
+
 impl<T: ?Sized, F> VisitWith<F> for Box<T>
 where
     F: Visit<T>,
@@ -184,6 +289,8 @@ where
     }
 }
 
+// ----- ----- impl for Vec<T> ----- -----
+
 impl<T, F> FoldWith<F> for Vec<T>
 where
     F: Fold<T>,
@@ -194,6 +301,15 @@ where
     }
 }
 
+impl<T, F> VisitMutWith<F> for Vec<T>
+where
+    F: VisitMut<T>,
+{
+    fn visit_mut_children(&mut self, f: &mut F) {
+        self.iter_mut().for_each(|node| f.visit_mut(node))
+    }
+}
+
 impl<T, F> VisitWith<F> for Vec<T>
 where
     F: Visit<T>,
@@ -203,6 +319,17 @@ where
     }
 }
 
+// ----- ----- impl for [T] ----- -----
+
+impl<T, F> VisitMutWith<F> for [T]
+where
+    F: VisitMut<T>,
+{
+    fn visit_mut_children(&mut self, f: &mut F) {
+        self.iter_mut().for_each(|node| f.visit_mut(node))
+    }
+}
+
 impl<T, F> VisitWith<F> for [T]
 where
     F: Visit<T>,
@@ -212,6 +339,8 @@ where
     }
 }
 
+// ----- ----- impl for Option<T> ----- -----
+
 impl<T, F> FoldWith<F> for Option<T>
 where
     F: Fold<T>,
@@ -221,6 +350,17 @@ where
     }
 }
 
+impl<T, F> VisitMutWith<F> for Option<T>
+where
+    F: VisitMut<T>,
+{
+    fn visit_mut_children(&mut self, f: &mut F) {
+        if let Some(ref mut node) = *self {
+            f.visit_mut(node)
+        }
+    }
+}
+
 impl<T, F> VisitWith<F> for Option<T>
 where
     F: Visit<T>,
@@ -232,34 +372,46 @@ where
     }
 }
 
+// ----- ----- impl for String ----- -----
+
+/// No op.
 impl<F> FoldWith<F> for String {
-    /// No op.
-
     fn fold_children(self, _: &mut F) -> Self {
         self
     }
 }
 
+/// No op.
+impl<F> VisitMutWith<F> for String {
+    fn visit_mut_children(&mut self, _: &mut F) {}
+}
+
+/// No op.
 impl<F> VisitWith<F> for String {
-    /// No op.
-
     fn visit_children(&self, _: &mut F) {}
 }
 
-impl<F, S: StaticAtomSet> FoldWith<F> for Atom<S> {
-    /// No op.
+// ----- ----- impl for string_cache::Atom ----- -----
 
+/// No op.
+impl<F, S: StaticAtomSet> FoldWith<F> for Atom<S> {
     fn fold_children(self, _: &mut F) -> Self {
         self
     }
 }
 
-impl<F, S: StaticAtomSet> VisitWith<F> for Atom<S> {
-    /// No op.
+/// No op.
+impl<F, S: StaticAtomSet> VisitMutWith<F> for Atom<S> {
+    fn visit_mut_children(&mut self, _: &mut F) {}
+}
 
+/// No op.
+impl<F, S: StaticAtomSet> VisitWith<F> for Atom<S> {
     fn visit_children(&self, _: &mut F) {}
 }
 
+// ----- ----- impl FoldWith for Either<A, B> ----- -----
+
 impl<A, B, F> FoldWith<F> for Either<A, B>
 where
     F: Fold<A> + Fold<B>,
@@ -272,6 +424,18 @@ where
     }
 }
 
+impl<A, B, F> VisitMutWith<F> for Either<A, B>
+where
+    F: VisitMut<A> + VisitMut<B>,
+{
+    fn visit_mut_children(&mut self, f: &mut F) {
+        match self {
+            Either::Left(l) => f.visit_mut(l),
+            Either::Right(r) => f.visit_mut(r),
+        }
+    }
+}
+
 impl<A, B, F> VisitWith<F> for Either<A, B>
 where
     F: Visit<A> + Visit<B>,
@@ -284,6 +448,8 @@ where
     }
 }
 
+// ----- ----- impl Fold for Either<A, B> ----- -----
+
 impl<A, B, T> Fold<T> for Either<A, B>
 where
     T: FoldWith<A> + FoldWith<B> + FoldWith<Self>,
@@ -296,6 +462,18 @@ where
     }
 }
 
+impl<A, B, T> VisitMut<T> for Either<A, B>
+where
+    T: VisitMutWith<A> + VisitMutWith<B> + VisitMutWith<Self>,
+{
+    fn visit_mut(&mut self, node: &mut T) {
+        match self {
+            Either::Left(l) => node.visit_mut_with(l),
+            Either::Right(r) => node.visit_mut_with(r),
+        }
+    }
+}
+
 impl<A, B, T> Visit<T> for Either<A, B>
 where
     T: VisitWith<A> + VisitWith<B> + VisitWith<Self>,
@@ -308,6 +486,8 @@ where
     }
 }
 
+// ----- ----- impl FoldWith for Arc<T> ----- -----
+
 impl<T, F> VisitWith<F> for Arc<T>
 where
     T: ?Sized,
@@ -318,6 +498,8 @@ where
     }
 }
 
+// ----- ----- impl FoldWith for Cow<T> ----- -----
+
 impl<'a, A, F> FoldWith<F> for Cow<'a, A>
 where
     A: Clone + FoldWith<F>,
@@ -329,6 +511,15 @@ where
     }
 }
 
+impl<'a, A, F> VisitMutWith<F> for Cow<'a, A>
+where
+    A: Clone + VisitMutWith<F>,
+{
+    fn visit_mut_children(&mut self, f: &mut F) {
+        self.to_mut().visit_mut_with(f)
+    }
+}
+
 impl<A, F> VisitWith<F> for Cow<'_, A>
 where
     A: Clone + VisitWith<F>,

common/src/lib.rs

@@ -1,11 +1,15 @@
 #![cfg_attr(feature = "fold", feature(specialization))]
 
 #[cfg(feature = "fold")]
-pub use self::fold::{Fold, FoldWith, Visit, VisitWith};
+pub use self::fold::{Fold, FoldWith, Visit, VisitMut, VisitMutWith, VisitWith};
 pub use self::{
     errors::{SourceMapper, SourceMapperDyn},
-    pos::*,
-    source_map::{FileLines, FileLoader, FileName, FilePathMapping, SourceMap, SpanSnippetError},
+    pos::{
+        hygiene, BytePos, CharPos, FileName, Globals, Loc, LocWithOpt, Mark, MultiSpan, SourceFile,
+        SourceFileAndBytePos, SourceFileAndLine, Span, SpanData, SpanLinesError, Spanned,
+        SyntaxContext, DUMMY_SP, GLOBALS, NO_EXPANSION,
+    },
+    source_map::{FileLines, FileLoader, FilePathMapping, SourceMap, SpanSnippetError},
 };
 pub use ast_node::{ast_node, DeserializeEnum, Fold, Spanned};
 pub use from_variant::FromVariant;

common/src/pos.rs

@@ -1,5 +1,5 @@
 #[cfg(feature = "fold")]
-use crate::fold::{FoldWith, VisitWith};
+use crate::fold::{FoldWith, VisitMutWith, VisitWith};
 pub use crate::syntax_pos::{
     hygiene, BytePos, CharPos, FileName, Globals, Loc, LocWithOpt, Mark, MultiSpan, SourceFile,
     SourceFileAndBytePos, SourceFileAndLine, Span, SpanData, SpanLinesError, SyntaxContext,
@@ -92,16 +92,22 @@ where
     }
 }
 
+/// No op as span does not have any child.
 #[cfg(feature = "fold")]
 impl<F> FoldWith<F> for Span {
-    /// No op as span does not have any child.
     fn fold_children(self, _: &mut F) -> Span {
         self
     }
 }
 
+/// No op as span does not have any child.
+#[cfg(feature = "fold")]
+impl<F> VisitMutWith<F> for Span {
+    fn visit_mut_children(&mut self, _: &mut F) {}
+}
+
+/// No op as span does not have any child.
 #[cfg(feature = "fold")]
 impl<F> VisitWith<F> for Span {
-    /// No op as span does not have any child.
     fn visit_children(&self, _: &mut F) {}
 }

common/src/syntax_pos.rs

@@ -19,6 +19,7 @@ mod analyze_source_file;
 pub mod hygiene;
 mod span_encoding;
 
+#[derive(Default)]
 pub struct Globals {
     span_interner: Lock<span_encoding::SpanInterner>,
     hygiene_data: Lock<hygiene::HygieneData>,

common/src/syntax_pos/hygiene.rs

@@ -130,7 +130,7 @@ impl Mark {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub(crate) struct HygieneData {
     marks: Vec<MarkData>,
     syntax_contexts: Vec<SyntaxContextData>,

common/tests/concurrent.rs

@@ -3,14 +3,15 @@ use std::{env, path::PathBuf, sync::Arc};
 use swc_common::{FilePathMapping, SourceFile, SourceMap};
 
 fn init() {
-    rayon::ThreadPoolBuilder::new()
+    let _ = rayon::ThreadPoolBuilder::new()
         .num_threads(100)
-        .build_global()
-        .unwrap();
+        .build_global();
 }
 
 #[test]
 fn no_overlap() {
+    init();
+
     let cm = Arc::new(SourceMap::new(FilePathMapping::empty()));
 
     let files: Vec<Arc<SourceFile>> = (0..100000)

macros/ast_node/src/fold.rs

@@ -5,7 +5,7 @@ use syn::*;
 
 #[derive(Debug, FromField)]
 #[darling(attributes(fold))]
-struct FieldAttrs {
+pub struct FieldAttrs {
     ///
     #[darling(default)]
     pub ignore: bool,
@@ -182,7 +182,7 @@ pub fn derive(input: DeriveInput) -> ItemImpl {
     derive_generics.append_to(item)
 }
 
-fn should_skip_field(field: &Field) -> bool {
+pub fn should_skip_field(field: &Field) -> bool {
     let attrs = FieldAttrs::from_field(field).expect("#[derive(Fold)]: failed to parse attribute");
     if attrs.ignore {
         return true;
@@ -198,7 +198,7 @@ fn should_skip_field(field: &Field) -> bool {
     false
 }
 
-fn normalize_type_for_bound(ty: Type) -> Type {
+pub fn normalize_type_for_bound(ty: Type) -> Type {
     use syn::fold::Fold;
 
     struct Norm;

macros/ast_node/src/lib.rs

@@ -12,6 +12,7 @@ mod enum_deserialize;
 mod fold;
 mod spanned;
 mod visit;
+mod visit_mut;
 
 /// Implements `FoldWith<F>` and `VisitWith<F>`.
 ///
@@ -28,17 +29,20 @@ pub fn derive_fold(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let name = input.ident.clone();
 
     let fold_item = self::fold::derive(input.clone());
-    let visit_item = self::visit::derive(input);
+    let visit_item = self::visit::derive(input.clone());
+    let visit_mut_item = self::visit_mut::derive(input);
     let item = Quote::new(def_site::<Span>()).quote_with(smart_quote!(
         Vars {
-            fold_item: fold_item,
-            visit_item: visit_item,
+            fold_item,
+            visit_item,
+            visit_mut_item,
             NAME: Ident::new(&format!("IMPL_FOLD_FOR_{}",name), Span::call_site()),
         },
         {
             const NAME: () = {
                 fold_item
                 visit_item
+                visit_mut_item
             };
         }
     ));

macros/ast_node/src/visit.rs

@@ -1,20 +1,8 @@
-use darling::FromField;
-use pmutil::{smart_quote, Quote, ToTokensExt};
+use crate::fold::{normalize_type_for_bound, should_skip_field};
+use pmutil::{smart_quote, Quote};
 use swc_macros_common::prelude::*;
 use syn::*;
 
-#[derive(Debug, FromField)]
-#[darling(attributes(fold))]
-struct FieldAttrs {
-    ///
-    #[darling(default)]
-    pub ignore: bool,
-
-    /// Should we add bound for the field's type?
-    #[darling(default)]
-    pub bound: bool,
-}
-
 pub fn derive(input: DeriveInput) -> ItemImpl {
     let mut derive_generics = Derive::new(&input);
 
@@ -149,49 +137,3 @@ pub fn derive(input: DeriveInput) -> ItemImpl {
         .parse();
     derive_generics.append_to(item)
 }
-
-fn should_skip_field(field: &Field) -> bool {
-    let attrs = FieldAttrs::from_field(field).expect("#[derive(Fold)]: failed to parse attribute");
-    if attrs.ignore {
-        return true;
-    }
-
-    let ty_str = field.ty.dump().to_string();
-    match &*ty_str {
-        "bool" | "usize" | "u128" | "u64" | "u32" | "u16" | "u8" | "isize" | "i128" | "i64"
-        | "i32" | "i16" | "i8" | "f64" | "f32" | "String" => return true,
-        _ => {}
-    }
-
-    false
-}
-
-fn normalize_type_for_bound(ty: Type) -> Type {
-    use syn::fold::Fold;
-
-    struct Norm;
-    impl Fold for Norm {
-        fn fold_path(&mut self, path: Path) -> Path {
-            if path.segments.len() == 1 {
-                let seg = &path.segments[0];
-                if seg.ident != "Box" && seg.ident != "Option" && seg.ident != "Vec" {
-                    return path.clone();
-                }
-
-                if let PathArguments::AngleBracketed(ref args) = seg.arguments {
-                    if args.args.len() == 1 {
-                        if let GenericArgument::Type(ref ty) = *args.args.last().unwrap() {
-                            if let Type::Path(TypePath { ref path, .. }) = *ty {
-                                return self.fold_path(path.clone());
-                            }
-                        }
-                    }
-                }
-            }
-
-            fold::fold_path(self, path)
-        }
-    }
-
-    Norm.fold_type(ty)
-}

macros/ast_node/src/visit_mut.rs

@@ -0,0 +1,142 @@
+use crate::fold::{normalize_type_for_bound, should_skip_field};
+use pmutil::{smart_quote, Quote};
+use swc_macros_common::prelude::*;
+use syn::*;
+
+pub fn derive(input: DeriveInput) -> ItemImpl {
+    let mut derive_generics = Derive::new(&input);
+
+    let preds = derive_generics
+        .all_generic_fields()
+        .into_iter()
+        .filter(|f| {
+            f.attrs.iter().any(|attr| {
+                is_attr_name(attr, "fold") && (attr.tokens.to_string().contains("bound"))
+            })
+        })
+        .map(|f| f.ty.clone())
+        .map(normalize_type_for_bound)
+        .map(|ty| {
+            Quote::new(def_site::<Span>())
+                .quote_with(smart_quote!(
+                    Vars { Type: &ty },
+                    (Type: swc_common::VisitMutWith<__V>)
+                ))
+                .parse()
+        });
+    derive_generics.add_where_predicates(preds);
+
+    let arms = Binder::new_from(&input)
+        .variants()
+        .into_iter()
+        .map(|v| {
+            let (pat, bindings) = v.bind("_", Some(def_site()), Some(def_site()));
+
+            let fields: Punctuated<Stmt, token::Semi> = bindings
+                .into_iter()
+                .filter_map(|binding| {
+                    // This closure will not be called for unit-like struct.
+
+                    let value = if should_skip_field(binding.field()) {
+                        None
+                    } else {
+                        Some(
+                            Quote::new(def_site::<Span>())
+                                .quote_with(smart_quote!(
+                                    Vars {
+                                        FieldType: &binding.field().ty,
+                                        binded_field: binding.name(),
+                                    },
+                                    {
+                                        swc_common::VisitMut::<FieldType>::visit_mut(
+                                            _v,
+                                            binded_field,
+                                        );
+                                    }
+                                ))
+                                .parse::<Stmt>(),
+                        )
+                    };
+
+                    let _attrs = binding
+                        .field()
+                        .attrs
+                        .iter()
+                        .filter(|attr| is_attr_name(attr, "cfg"))
+                        .cloned()
+                        .collect::<Vec<_>>();
+
+                    value
+                })
+                .map(|t| Element::Punctuated(t, def_site()))
+                .collect();
+
+            let body = match *v.data() {
+                // Handle unit-like structs separately
+                Fields::Unit => Box::new(
+                    Quote::new(def_site::<Span>())
+                        .quote_with(smart_quote!(Vars {}, {
+                            {
+                                // no-op
+                            }
+                        }))
+                        .parse(),
+                ),
+                _ => Box::new(
+                    Quote::new(def_site::<Span>())
+                        .quote_with(smart_quote!(Vars { fields }, {
+                            {
+                                fields
+                            }
+                        }))
+                        .parse(),
+                ),
+            };
+
+            Arm {
+                body,
+
+                attrs: v
+                    .attrs()
+                    .iter()
+                    .filter(|attr| is_attr_name(attr, "cfg"))
+                    .cloned()
+                    .collect(),
+                pat,
+                guard: None,
+                fat_arrow_token: def_site(),
+                comma: Some(def_site()),
+            }
+        })
+        .collect();
+
+    let body = Expr::Match(ExprMatch {
+        attrs: Default::default(),
+        match_token: def_site(),
+        brace_token: def_site(),
+        expr: Box::new(
+            Quote::new(def_site::<Span>())
+                .quote_with(smart_quote!(Vars {}, { *self }))
+                .parse(),
+        ),
+        arms,
+    });
+
+    let item = Quote::new(def_site::<Span>())
+        .quote_with(smart_quote!(
+            Vars {
+                Type: &input.ident,
+                body,
+            },
+            {
+                impl<__V> swc_common::VisitMutWith<__V> for Type {
+                    #[inline]
+                    fn visit_mut_children(&mut self, _v: &mut __V) {
+                        body
+                    }
+                }
+            }
+        ))
+        .parse();
+    derive_generics.append_to(item)
+}