use std::collections::HashSet;
use inflector::Inflector;
use proc_macro2::{Span, TokenStream};
use quote::{quote, ToTokens};
use syn::{
parse_quote, Arm, Attribute, Expr, Field, Fields, File, GenericArgument, Ident, Item, Lit,
LitInt, Path, PathArguments, Stmt, TraitItem, Type,
};
pub fn generate(crate_name: &Ident, node_types: &[&Item]) -> File {
let mut output = File {
shebang: None,
attrs: Vec::new(),
items: Vec::new(),
};
let mut all_types = all_field_types(node_types).into_iter().collect::<Vec<_>>();
all_types.sort_by_cached_key(|v| v.method_name());
let mut typedefs = HashSet::new();
for node_type in node_types {
match node_type {
Item::Enum(data) => {
typedefs.insert(FieldType::Normal(data.ident.to_string()));
}
Item::Struct(data) => {
typedefs.insert(FieldType::Normal(data.ident.to_string()));
}
_ => {}
}
}
let field_only_types = {
let mut all = all_types.clone();
all.retain(|ty| !typedefs.contains(ty));
all
};
output.attrs.push(parse_quote!(
));
output.attrs.push(parse_quote!(
#![allow(unused_variables)]
));
output.attrs.push(parse_quote!(
#![allow(clippy::all)]
));
output.items.push(parse_quote!(
use #crate_name::*;
));
output.items.push(parse_quote!(
pub use ::swc_visit::All;
));
for &kind in [TraitKind::Visit, TraitKind::VisitMut, TraitKind::Fold].iter() {
for &variant in [Variant::Normal, Variant::AstPath].iter() {
let g = Generator { kind, variant };
output.items.extend(g.declare_visit_trait(&all_types));
output.items.extend(g.declare_visit_with_trait());
output
.items
.extend(g.implement_visit_with_for_node_types(node_types));
output
.items
.extend(g.implement_visit_with_for_non_node_types(&field_only_types));
output
.items
.extend(g.implement_visit_with_for_generic_types());
}
}
output.items.push(parse_quote!(
#[cfg(any(docsrs, feature = "path"))]
pub type AstKindPath = swc_visit::AstKindPath<AstParentKind>;
));
output.items.push(parse_quote!(
#[cfg(any(docsrs, feature = "path"))]
pub type AstNodePath<'ast> = swc_visit::AstNodePath<AstParentNodeRef<'ast>>;
));
output.items.extend(define_fields(crate_name, node_types));
output
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum FieldType {
Normal(String),
Generic(String, Box<FieldType>),
}
impl ToTokens for FieldType {
fn to_tokens(&self, tokens: &mut TokenStream) {
match self {
FieldType::Normal(name) => {
let parsed: Path = syn::parse_str(name).expect("failed to parse path");
parsed.to_tokens(tokens);
}
FieldType::Generic(name, ty) => {
let name = Ident::new(name, Span::call_site());
let ty = &**ty;
quote!(#name<#ty>).to_tokens(tokens);
}
}
}
}
impl FieldType {
pub fn method_name(&self) -> String {
match self {
FieldType::Normal(name) => name.split("::").last().unwrap().to_snake_case(),
FieldType::Generic(name, ty) => match &**name {
"Option" => format!("opt_{}", ty.method_name()),
"Vec" => {
match &**ty {
FieldType::Generic(name, ty) if name == "Option" => {
return format!("opt_vec_{}s", ty.method_name())
}
_ => {}
}
format!("{}s", ty.method_name())
}
"Box" => ty.method_name(),
_ => todo!("method_name for generic type: {}", name),
},
}
}
}
fn all_field_types(node_types: &[&Item]) -> HashSet<FieldType> {
let mut all_types = HashSet::new();
for ty in node_types {
let type_name = match ty {
Item::Enum(data) => data.ident.to_string(),
Item::Struct(data) => data.ident.to_string(),
_ => continue,
};
all_types.insert(FieldType::Normal(type_name));
match ty {
Item::Enum(data) => {
for variant in &data.variants {
for field in &variant.fields {
let ty = &field.ty;
all_types.extend(all_types_in_ty(ty));
}
}
}
Item::Struct(data) => {
for field in &data.fields {
let ty = &field.ty;
all_types.extend(all_types_in_ty(ty));
}
}
_ => continue,
}
}
all_types
}
fn to_field_ty(ty: &Type) -> Option<FieldType> {
if let Some(ty) = extract_vec(ty) {
return to_field_ty(ty).map(|ty| FieldType::Generic("Vec".into(), Box::new(ty)));
}
if let Some(ty) = extract_generic("Box", ty) {
return to_field_ty(ty).map(|ty| FieldType::Generic("Box".into(), Box::new(ty)));
}
if let Some(ty) = extract_generic("Option", ty) {
return to_field_ty(ty).map(|ty| FieldType::Generic("Option".into(), Box::new(ty)));
}
match ty {
Type::Path(p) => {
let last = p.path.segments.last().unwrap();
if last.arguments.is_empty() {
let i = &last.ident;
if i == "bool"
|| i == "char"
|| i == "f32"
|| i == "f64"
|| i == "i8"
|| i == "i16"
|| i == "i32"
|| i == "i64"
|| i == "i128"
|| i == "isize"
|| i == "str"
|| i == "u8"
|| i == "u16"
|| i == "u32"
|| i == "u64"
|| i == "u128"
|| i == "usize"
{
return None;
}
return Some(FieldType::Normal(quote!(#p).to_string()));
}
todo!("to_field_ty: {:?}", ty)
}
_ => todo!("to_field_ty"),
}
}
fn all_types_in_ty(ty: &Type) -> Vec<FieldType> {
if let Some(ty) = extract_vec(ty) {
let mut types = all_types_in_ty(ty);
types.extend(to_field_ty(ty).map(|ty| FieldType::Generic("Vec".into(), Box::new(ty))));
return types;
}
if let Some(ty) = extract_generic("Box", ty) {
let mut types = all_types_in_ty(ty);
types.extend(to_field_ty(ty).map(|ty| FieldType::Generic("Box".into(), Box::new(ty))));
return types;
}
if let Some(ty) = extract_generic("Option", ty) {
let mut types = all_types_in_ty(ty);
types.extend(to_field_ty(ty).map(|ty| FieldType::Generic("Option".into(), Box::new(ty))));
return types;
}
to_field_ty(ty).into_iter().collect()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TraitKind {
Visit,
VisitMut,
Fold,
}
impl TraitKind {
pub fn method_prefix(self) -> &'static str {
match self {
TraitKind::Visit => "visit",
TraitKind::VisitMut => "visit_mut",
TraitKind::Fold => "fold",
}
}
pub fn trait_prefix(self) -> &'static str {
match self {
TraitKind::Visit => "Visit",
TraitKind::VisitMut => "VisitMut",
TraitKind::Fold => "Fold",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Variant {
Normal,
AstPath,
}
impl Variant {
pub fn method_suffix(self, is_visitor_method: bool) -> &'static str {
if self == Variant::Normal || is_visitor_method {
""
} else {
"_ast_path"
}
}
}
struct Generator {
kind: TraitKind,
variant: Variant,
}
impl Generator {
fn should_skip(&self, ty: &Type) -> bool {
if let Some(ty) = extract_generic("Box", ty) {
return self.should_skip(ty);
}
if let Some(ty) = extract_generic("Vec", ty) {
return self.should_skip(ty);
}
if let Some(ty) = extract_generic("Option", ty) {
return self.should_skip(ty);
}
let ty = to_field_ty(ty);
match ty {
Some(..) => {}
None => return true,
}
false
}
fn method_lifetime(&self) -> TokenStream {
match self.kind {
TraitKind::Visit => match self.variant {
Variant::Normal => quote!(),
Variant::AstPath => quote!(<'ast: 'r, 'r>),
},
TraitKind::VisitMut => quote!(),
TraitKind::Fold => quote!(),
}
}
fn parameter_type_token(&self, ty: TokenStream) -> TokenStream {
match self.kind {
TraitKind::Visit => match self.variant {
Variant::Normal => quote!(&#ty),
Variant::AstPath => quote!(&'ast #ty),
},
TraitKind::VisitMut => quote!(&mut #ty),
TraitKind::Fold => ty,
}
}
fn return_type_token(&self, ty: TokenStream) -> TokenStream {
match self.kind {
TraitKind::Visit => quote!(),
TraitKind::VisitMut => quote!(),
TraitKind::Fold => quote!(-> #ty),
}
}
fn arg_extra_token(&self) -> TokenStream {
match self.variant {
Variant::Normal => quote!(),
Variant::AstPath => quote!(, __ast_path),
}
}
fn param_extra_token(&self) -> TokenStream {
match self.variant {
Variant::Normal => quote!(),
Variant::AstPath => match self.kind {
TraitKind::Visit => {
quote!(, __ast_path: &mut AstNodePath<'r>)
}
TraitKind::VisitMut | TraitKind::Fold => quote!(, __ast_path: &mut AstKindPath),
},
}
}
fn trait_name(&self, with: bool) -> Ident {
let name = self.kind.trait_prefix();
let name = if with {
format!("{}With", name)
} else {
name.to_string()
};
match self.variant {
Variant::Normal => Ident::new(&name, Span::call_site()),
Variant::AstPath => Ident::new(&format!("{}AstPath", name), Span::call_site()),
}
}
fn base_trait_attrs(&self) -> Vec<Attribute> {
let mut attrs = Vec::new();
if self.variant == Variant::AstPath {
attrs.push(parse_quote!(#[cfg(any(docsrs, feature = "path"))]));
attrs.push(parse_quote!(#[cfg_attr(docsrs, doc(cfg(feature = "path")))]));
}
attrs
}
fn declare_visit_trait(&self, all_types: &[FieldType]) -> Vec<Item> {
let mut items = Vec::<Item>::new();
let lifetime = self.method_lifetime();
let ast_path_arg = self.arg_extra_token();
let ast_path_params = self.param_extra_token();
let with_trait_name = self.trait_name(true);
let trait_name = self.trait_name(false);
let attrs = self.base_trait_attrs();
let mut trait_methods = Vec::<TraitItem>::new();
let mut either_impl_methods = Vec::<TraitItem>::new();
let mut optional_impl_methods = Vec::<TraitItem>::new();
let mut ptr_impl_methods = Vec::<TraitItem>::new();
for ty in all_types {
if let FieldType::Generic(name, ..) = &ty {
if name == "Box" {
continue;
}
}
let type_name = quote!(#ty);
let return_type = self.return_type_token(quote!(#type_name));
let node_type = self.node_type_for_visitor_method(ty);
let type_param = self.parameter_type_token(quote!(#node_type));
let visit_method_name = Ident::new(
&format!(
"{}_{}{}",
self.kind.method_prefix(),
ty.method_name(),
self.variant.method_suffix(true)
),
Span::call_site(),
);
let visit_with_children_name = Ident::new(
&format!(
"{}_children_with{}",
self.kind.method_prefix(),
self.variant.method_suffix(false)
),
Span::call_site(),
);
let recurse_doc = "If you want to recurse, you need to call it manually.";
let method_doc = doc(&format!(
"Visit a node of type `{}`.\n\nBy default, this method calls \
[`{type_name}::{visit_with_children_name}`]. {recurse_doc}",
type_name
));
trait_methods.push(parse_quote!(
#method_doc
#[inline]
fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
<#node_type as #with_trait_name<Self>>::#visit_with_children_name(node, self #ast_path_arg)
}
));
either_impl_methods.push(parse_quote!(
#[inline]
fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
match self {
swc_visit::Either::Left(visitor) => {
#trait_name::#visit_method_name(visitor, node #ast_path_arg)
}
swc_visit::Either::Right(visitor) => {
#trait_name::#visit_method_name(visitor, node #ast_path_arg)
}
}
}
));
let else_block = if self.kind == TraitKind::Fold {
quote!(node)
} else {
quote!()
};
optional_impl_methods.push(parse_quote!(
#[inline]
fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
if self.enabled {
<V as #trait_name>::#visit_method_name(&mut self.visitor, node #ast_path_arg)
} else {
#else_block
}
}
));
ptr_impl_methods.push(parse_quote!(
#[inline]
fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
<V as #trait_name>::#visit_method_name(&mut **self, node #ast_path_arg)
}
));
}
items.push(parse_quote! {
#(#attrs)*
pub trait #trait_name {
#(#trait_methods)*
}
});
items.push(parse_quote! {
#(#attrs)*
impl<V> #trait_name for &mut V where V: ?Sized + #trait_name {
#(#ptr_impl_methods)*
}
});
items.push(parse_quote! {
#(#attrs)*
impl<V> #trait_name for Box<V> where V: ?Sized + #trait_name {
#(#ptr_impl_methods)*
}
});
items.push(parse_quote! {
#(#attrs)*
impl<A, B> #trait_name for ::swc_visit::Either<A, B>
where
A: #trait_name,
B: #trait_name,
{
#(#either_impl_methods)*
}
});
items.push(parse_quote! {
#(#attrs)*
impl<V> #trait_name for ::swc_visit::Optional<V>
where
V: #trait_name,
{
#(#optional_impl_methods)*
}
});
items
}
fn declare_visit_with_trait(&self) -> Vec<Item> {
let visitor_trait_name = self.trait_name(false);
let trait_name = self.trait_name(true);
let attrs = self.base_trait_attrs();
let mut visit_with_trait_methods: Vec<TraitItem> = Vec::new();
{
let lifetime = self.method_lifetime();
let ast_path_extra = self.param_extra_token();
let return_type = self.return_type_token(quote!(Self));
let receiver = self.parameter_type_token(quote!(self));
let visit_with_name = Ident::new(
&format!(
"{}_with{}",
self.kind.method_prefix(),
self.variant.method_suffix(false)
),
Span::call_site(),
);
let visit_with_children_name = Ident::new(
&format!(
"{}_children_with{}",
self.kind.method_prefix(),
self.variant.method_suffix(false)
),
Span::call_site(),
);
visit_with_trait_methods.push(parse_quote!(
fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_extra) #return_type;
));
visit_with_trait_methods.push(parse_quote!(
fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_extra) #return_type;
));
}
let mut items: Vec<Item> = Vec::new();
items.push(parse_quote!(
#(#attrs)*
pub trait #trait_name<V: ?Sized + #visitor_trait_name> {
#(#visit_with_trait_methods)*
}
));
items
}
fn implement_visit_with_for_node_types(&self, node_types: &[&Item]) -> Vec<Item> {
let visitor_trait_name = self.trait_name(false);
let trait_name = self.trait_name(true);
let attrs = self.base_trait_attrs();
let mut items: Vec<Item> = Vec::new();
for node_type in node_types {
let type_name = match node_type {
Item::Enum(data) => data.ident.clone(),
Item::Struct(data) => data.ident.clone(),
_ => continue,
};
let lifetime = self.method_lifetime();
let ast_path_arg = self.arg_extra_token();
let ast_path_param = self.param_extra_token();
let return_type = self.return_type_token(quote!(Self));
let receiver = self.parameter_type_token(quote!(self));
let visit_with_name = Ident::new(
&format!(
"{}_with{}",
self.kind.method_prefix(),
self.variant.method_suffix(false)
),
Span::call_site(),
);
let visit_with_children_name = Ident::new(
&format!(
"{}_children_with{}",
self.kind.method_prefix(),
self.variant.method_suffix(false)
),
Span::call_site(),
);
let visit_method_name = Ident::new(
&format!(
"{}_{}{}",
self.kind.method_prefix(),
type_name.to_string().to_snake_case(),
self.variant.method_suffix(true)
),
Span::call_site(),
);
let visit_with_doc = doc(&format!(
"Calls [{visitor_trait_name}`::{}`] with `self`.",
visit_method_name
));
let default_body: Expr = match node_type {
Item::Enum(data) => {
let name = &data.ident;
let mut match_arms = Vec::new();
for v in &data.variants {
let variant_name = &v.ident;
match_arms.push(self.default_visit_body(
quote!(#name::#variant_name),
name,
Some(variant_name),
&v.fields,
));
}
parse_quote!(match self { #(#match_arms)* })
}
Item::Struct(data) => {
let name = &data.ident;
let arm = self.default_visit_body(quote!(#name), name, None, &data.fields);
parse_quote!(match self { #arm })
}
_ => continue,
};
items.push(parse_quote!(
#(#attrs)*
impl<V: ?Sized + #visitor_trait_name> #trait_name<V> for #type_name {
#visit_with_doc
fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
<V as #visitor_trait_name>::#visit_method_name(visitor, self #ast_path_arg)
}
fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
#default_body
}
}
));
}
items
}
fn default_visit_body(
&self,
path: TokenStream,
type_name: &Ident,
enum_variant_name: Option<&Ident>,
fields: &Fields,
) -> Arm {
let ast_path_arg = match self.variant {
Variant::Normal => quote!(),
Variant::AstPath => quote!(, &mut *__ast_path),
};
let with_visitor_trait_name = self.trait_name(true);
let visit_with_name = Ident::new(
&format!(
"{}_with{}",
self.kind.method_prefix(),
self.variant.method_suffix(false)
),
Span::call_site(),
);
let fields_enum_name = Ident::new(&format!("{type_name}Field"), Span::call_site());
let enum_ast_path = match enum_variant_name {
Some(variant_name) if self.variant == Variant::AstPath => {
let field_variant = Ident::new(
&variant_name.to_string().to_pascal_case(),
Span::call_site(),
);
match self.kind {
TraitKind::Visit => Some(quote!(
let mut __ast_path = __ast_path
.with_guard(
AstParentNodeRef::#type_name(self, self::fields::#fields_enum_name::#field_variant),
);
)),
_ => Some(quote!(
let mut __ast_path = __ast_path
.with_guard(
AstParentKind::#type_name(self::fields::#fields_enum_name::#field_variant),
);
)),
}
}
_ => None,
};
match fields {
Fields::Named(n) => {
let mut stmts: Vec<Stmt> = Vec::new();
let mut bindings = Vec::new();
let mut reconstruct = match self.kind {
TraitKind::Visit | TraitKind::VisitMut => None,
TraitKind::Fold => Some(Vec::<TokenStream>::new()),
};
for field in &n.named {
let field_name = field.ident.as_ref().unwrap();
let ty = &field.ty;
bindings.push(field_name.clone());
let field_variant =
Ident::new(&field_name.to_string().to_pascal_case(), Span::call_site());
let mut ast_path_guard_expr: Option<Stmt> = None;
if self.variant == Variant::AstPath && !self.should_skip(ty) {
let mut kind = quote!(self::fields::#fields_enum_name::#field_variant);
if extract_vec(extract_generic("Option", ty).unwrap_or(ty)).is_some() {
kind = quote!(#kind(usize::MAX));
}
match self.kind {
TraitKind::Visit => {
ast_path_guard_expr = Some(parse_quote!(
let mut __ast_path = __ast_path
.with_guard(
AstParentNodeRef::#type_name(self, #kind),
);
));
}
_ => {
ast_path_guard_expr = Some(parse_quote!(
let mut __ast_path = __ast_path
.with_guard(
AstParentKind::#type_name(#kind),
);
));
}
}
}
if let Some(reconstructor) = &mut reconstruct {
if !self.should_skip(ty) {
stmts.push(parse_quote!(
let #field_name = {
#ast_path_guard_expr
<#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg)
};
));
}
reconstructor.push(parse_quote!(#field_name));
} else if !self.should_skip(ty) {
stmts.push(parse_quote!(
{
#ast_path_guard_expr
<#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg)
};
));
}
}
match self.kind {
TraitKind::Visit | TraitKind::VisitMut => {
parse_quote!(#path { #(#bindings),* } => {
#enum_ast_path;
#(#stmts)*
})
}
TraitKind::Fold => {
let reconstruct = reconstruct.unwrap();
parse_quote!(#path { #(#bindings),* } => {
#enum_ast_path;
#(#stmts)*
#path {
#(#reconstruct),*
}
})
}
}
}
Fields::Unnamed(u) => {
let mut stmts: Vec<Stmt> = Vec::new();
let mut bindings = Vec::<TokenStream>::new();
let mut reconstruct = match self.kind {
TraitKind::Visit | TraitKind::VisitMut => None,
TraitKind::Fold => Some(Vec::<TokenStream>::new()),
};
for (idx, field) in u.unnamed.iter().enumerate() {
let field_name = Ident::new(&format!("_field_{}", idx), Span::call_site());
let ty = &field.ty;
let binding_idx = Lit::Int(LitInt::new(&idx.to_string(), Span::call_site()));
bindings.push(parse_quote!(#binding_idx: #field_name));
if let Some(reconstructor) = &mut reconstruct {
if !self.should_skip(ty) {
stmts.push(parse_quote!(
let #field_name = <#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg);
));
}
reconstructor.push(parse_quote!(#binding_idx: #field_name));
} else if !self.should_skip(ty) {
stmts.push(parse_quote!(
<#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg);
));
}
}
match self.kind {
TraitKind::Visit | TraitKind::VisitMut => {
parse_quote!(#path { #(#bindings),* }=> {
#enum_ast_path;
#(#stmts)*
})
}
TraitKind::Fold => {
let reconstruct = reconstruct.unwrap();
parse_quote!(#path { #(#bindings),* } => {
#enum_ast_path;
#(#stmts)*
#path{#(#reconstruct),*}
})
}
}
}
Fields::Unit => match self.kind {
TraitKind::Visit | TraitKind::VisitMut => {
parse_quote!(#path => {})
}
TraitKind::Fold => parse_quote!(#path => #path,),
},
}
}
fn implement_visit_with_for_non_node_types(&self, non_leaf_types: &[FieldType]) -> Vec<Item> {
let visitor_trait_name = self.trait_name(false);
let visit_with_trait_name = self.trait_name(true);
let attrs = self.base_trait_attrs();
let lifetime = self.method_lifetime();
let ast_path_arg = self.arg_extra_token();
let ast_path_param = self.param_extra_token();
let return_type = self.return_type_token(quote!(Self));
let receiver = self.parameter_type_token(quote!(self));
let mut items: Vec<Item> = Vec::new();
for node_type in non_leaf_types {
let visit_with_name = Ident::new(
&format!(
"{}_with{}",
self.kind.method_prefix(),
self.variant.method_suffix(false)
),
Span::call_site(),
);
let visit_with_children_name = Ident::new(
&format!(
"{}_children_with{}",
self.kind.method_prefix(),
self.variant.method_suffix(false)
),
Span::call_site(),
);
let visit_method_name = Ident::new(
&format!(
"{}_{}{}",
self.kind.method_prefix(),
node_type.method_name(),
self.variant.method_suffix(true)
),
Span::call_site(),
);
let visit_with_doc = doc(&format!(
"Calls [{visitor_trait_name}`::{}`] with `self`. (Extra impl)",
visit_method_name
));
let default_body: Expr = match node_type {
FieldType::Normal(..) => match self.kind {
TraitKind::Visit => {
parse_quote!({})
}
TraitKind::VisitMut => {
parse_quote!({})
}
TraitKind::Fold => {
parse_quote!(self)
}
},
FieldType::Generic(name, inner) => match &**name {
"Vec" => {
let inner = inner.as_ref();
let inner_ty = quote!(#inner);
match (self.kind, self.variant) {
(TraitKind::Visit, Variant::Normal) => {
parse_quote!(self.iter().for_each(|item| {
<#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor #ast_path_arg)
}))
}
(TraitKind::Visit, Variant::AstPath) => {
parse_quote!(self.iter().enumerate().for_each(|(__idx, item)| {
let mut __ast_path = __ast_path.with_index_guard(__idx);
<#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor, &mut *__ast_path)
}))
}
(TraitKind::VisitMut, Variant::Normal) => {
parse_quote!(
self.iter_mut().for_each(|item| {
<#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor #ast_path_arg)
})
)
}
(TraitKind::VisitMut, Variant::AstPath) => {
parse_quote!(
self.iter_mut().enumerate().for_each(|(__idx, item)| {
let mut __ast_path = __ast_path.with_index_guard(__idx);
<#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor, &mut *__ast_path)
})
)
}
(TraitKind::Fold, Variant::Normal) => {
parse_quote!(
swc_visit::util::move_map::MoveMap::move_map(self, |item| {
<#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor #ast_path_arg)
})
)
}
(TraitKind::Fold, Variant::AstPath) => {
parse_quote!(
self.into_iter().enumerate().map(|(__idx, item)| {
let mut __ast_path = __ast_path.with_index_guard(__idx);
<#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor, &mut *__ast_path)
}).collect()
)
}
}
}
"Option" => {
let inner = inner.as_ref();
let inner_ty = quote!(#inner);
match self.kind {
TraitKind::Visit => {
parse_quote!(
match self {
Some(inner) => {
<#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
}
None => {}
}
)
}
TraitKind::VisitMut => {
parse_quote!(
match self {
Some(inner) => {
<#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
}
None => {}
}
)
}
TraitKind::Fold => {
parse_quote!(
self.map(|inner| {
<#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
})
)
}
}
}
"Box" => continue,
_ => unreachable!("unexpected generic type: {}", name),
},
};
let target_type = self.node_type_for_visitor_method(node_type);
items.push(parse_quote!(
#(#attrs)*
impl<V: ?Sized + #visitor_trait_name> #visit_with_trait_name<V> for #target_type {
#visit_with_doc
#[inline]
fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
<V as #visitor_trait_name>::#visit_method_name(visitor, self #ast_path_arg)
}
#[inline]
fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
#default_body
}
}
));
}
items
}
fn implement_visit_with_for_generic_types(&self) -> Vec<Item> {
let visit_trait_name = self.trait_name(false);
let visit_with_trait_name = self.trait_name(true);
let lifetime = self.method_lifetime();
let ast_path_arg = self.arg_extra_token();
let ast_path_param = self.param_extra_token();
let return_type = self.return_type_token(quote!(Self));
let attrs = self.base_trait_attrs();
let receiver = self.parameter_type_token(quote!(self));
let visit_with_name = Ident::new(
&format!(
"{}_with{}",
self.kind.method_prefix(),
self.variant.method_suffix(false)
),
Span::call_site(),
);
let visit_with_children_name = Ident::new(
&format!(
"{}_children_with{}",
self.kind.method_prefix(),
self.variant.method_suffix(false)
),
Span::call_site(),
);
let mut items = Vec::<Item>::new();
{
match self.kind {
TraitKind::Fold => {
items.push(parse_quote!(
#(#attrs)*
impl<V, T> #visit_with_trait_name<V> for std::boxed::Box<T>
where V: ?Sized + #visit_trait_name,
T: #visit_with_trait_name<V> {
#[inline]
fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
swc_visit::util::map::Map::map(self, |inner| {
<T as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
})
}
#[inline]
fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
swc_visit::util::map::Map::map(self, |inner| {
<T as #visit_with_trait_name<V>>::#visit_with_children_name(inner, visitor #ast_path_arg)
})
}
}
));
}
_ => {
let deref_expr = match self.kind {
TraitKind::Visit => {
quote!(&**self)
}
TraitKind::VisitMut => {
quote!(&mut **self)
}
TraitKind::Fold => {
unreachable!()
}
};
let restore_expr = match self.kind {
TraitKind::Visit => {
quote!()
}
TraitKind::VisitMut => {
quote!()
}
TraitKind::Fold => {
unreachable!()
}
};
items.push(parse_quote!(
#(#attrs)*
impl<V, T> #visit_with_trait_name<V> for std::boxed::Box<T>
where V: ?Sized + #visit_trait_name,
T: #visit_with_trait_name<V> {
#[inline]
fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
let v = <T as #visit_with_trait_name<V>>::#visit_with_name(#deref_expr, visitor #ast_path_arg);
#restore_expr
v
}
#[inline]
fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
let v = <T as #visit_with_trait_name<V>>::#visit_with_children_name(#deref_expr, visitor #ast_path_arg);
#restore_expr
v
}
}
));
}
}
}
if self.kind == TraitKind::Visit {
items.push(parse_quote!(
#(#attrs)*
impl<V, T> #visit_with_trait_name<V> for std::vec::Vec<T>
where V: ?Sized + #visit_trait_name,
[T]: #visit_with_trait_name<V> {
#[inline]
fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
let v = <[T] as #visit_with_trait_name<V>>::#visit_with_name(self, visitor #ast_path_arg);
v
}
#[inline]
fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
let v = <[T] as #visit_with_trait_name<V>>::#visit_with_children_name(self, visitor #ast_path_arg);
v
}
}
));
}
items
}
fn node_type_for_visitor_method(&self, node_type: &FieldType) -> TokenStream {
match self.kind {
TraitKind::Visit => match node_type {
FieldType::Generic(name, inner) if name == "Vec" => {
let inner_ty = quote!(#inner);
quote!([#inner_ty])
}
_ => quote!(#node_type),
},
_ => quote!(#node_type),
}
}
}
fn doc(s: &str) -> Attribute {
parse_quote!(#[doc = #s])
}
fn field_variant(type_name: &Ident, field: &Field) -> Option<(TokenStream, Option<Arm>)> {
if let Some(field_name) = &field.ident {
let variant_name = Ident::new(&field_name.to_string().to_pascal_case(), Span::call_site());
let variant_doc = doc(&format!("Represents [`{type_name}::{field_name}`]"));
if extract_vec(extract_generic("Option", &field.ty).unwrap_or(&field.ty)).is_some() {
let v = quote!(
#variant_doc
#variant_name(usize)
);
let arg = parse_quote!(
Self::#variant_name(idx) => {
assert_initial_index(*idx, index);
*idx = index;
},
);
return Some((v, Some(arg)));
}
return Some((
quote!(
#variant_doc
#variant_name
),
None,
));
}
None
}
fn extract_vec(ty: &Type) -> Option<&Type> {
extract_generic("Vec", ty)
}
fn extract_generic<'a>(name: &str, ty: &'a Type) -> Option<&'a Type> {
if let Type::Path(p) = ty {
let last = p.path.segments.last().unwrap();
if !last.arguments.is_empty() && last.ident == name {
match &last.arguments {
PathArguments::AngleBracketed(tps) => {
let arg = tps.args.first().unwrap();
match arg {
GenericArgument::Type(arg) => return Some(arg),
_ => unimplemented!("generic parameter other than type"),
}
}
_ => unimplemented!("Box() -> T or Box without a type parameter"),
}
}
}
if let Type::Reference(r) = ty {
return extract_generic(name, &r.elem);
}
None
}
fn to_iter(e: TokenStream, ty: &Type, node_names: &[Ident]) -> Option<Expr> {
if let Some(ty) = extract_vec(ty) {
let inner_expr = to_iter(quote!(item), ty, node_names)?;
return Some(parse_quote!(#e.iter().flat_map(|item| #inner_expr)));
}
if let Some(ty) = extract_generic("Option", ty) {
let inner_expr = to_iter(quote!(item), ty, node_names)?;
return Some(parse_quote!(#e.iter().flat_map(|item| #inner_expr)));
}
if let Some(ty) = extract_generic("Box", ty) {
let inner_expr = to_iter(quote!(item), ty, node_names)?;
return Some(parse_quote!({
let item = &*#e;
#inner_expr
}));
}
if let Type::Path(p) = ty {
let ty = &p.path.segments.last().unwrap().ident;
if node_names.contains(ty) {
return Some(parse_quote!(::std::iter::once(NodeRef::#ty(&#e))));
}
None
} else {
todo!("to_iter for {:?}", ty);
}
}
fn define_fields(crate_name: &Ident, node_types: &[&Item]) -> Vec<Item> {
let mut items = Vec::<Item>::new();
let mut kind_enum_members = Vec::new();
let mut parent_enum_members = Vec::new();
let mut node_ref_enum_members = Vec::new();
let mut kind_set_index_arms = Vec::<Arm>::new();
let mut node_ref_set_index_arms = Vec::<Arm>::new();
let mut node_ref_kind_arms = Vec::<Arm>::new();
let mut node_ref_iter_next_arms = Vec::<Arm>::new();
let node_names = node_types
.iter()
.filter_map(|ty| match ty {
Item::Enum(data) => Some(data.ident.clone()),
Item::Struct(data) => Some(data.ident.clone()),
_ => None,
})
.collect::<Vec<_>>();
let is_node_ref_raw = |ty: &Type| match ty {
Type::Path(p) => node_names.contains(&p.path.segments.last().unwrap().ident),
_ => false,
};
let is_node_ref = |ty: &Type| {
if let Some(ty) = extract_generic("Box", ty) {
return is_node_ref_raw(ty);
}
is_node_ref_raw(ty)
};
{
let mut defs = Vec::<Item>::new();
defs.push(parse_quote!(
use #crate_name::*;
));
defs.push(parse_quote!(
#[inline(always)]
fn assert_initial_index(idx: usize, index: usize) {
#[cfg(debug_assertions)]
if !(idx == usize::MAX || index == usize::MAX) {
{
panic!("Should be usize::MAX");
}
}
}
));
for ty in node_types {
let type_name = match ty {
Item::Enum(data) => data.ident.clone(),
Item::Struct(data) => data.ident.clone(),
_ => continue,
};
let fields_enum_name = Ident::new(&format!("{type_name}Field"), Span::call_site());
let mut variants = Vec::new();
kind_set_index_arms.push(parse_quote!(
Self::#type_name(v) => v.set_index(index),
));
node_ref_kind_arms.push(parse_quote!(
Self::#type_name(_, __field_kind) => AstParentKind::#type_name(*__field_kind),
));
node_ref_set_index_arms.push(parse_quote!(
Self::#type_name(_, __field_kind) => __field_kind.set_index(index),
));
match ty {
Item::Enum(data) => {
for variant in &data.variants {
let orig_ident = &variant.ident;
let variant_name = Ident::new(
&variant.ident.to_string().to_pascal_case(),
Span::call_site(),
);
let variant_doc = doc(&format!("Represents [`{type_name}::{orig_ident}`]"));
variants.push(quote!(
#variant_doc
#variant_name
));
}
kind_enum_members.push(quote!(
#type_name(#fields_enum_name)
));
parent_enum_members.push(quote!(
#type_name(&'ast #type_name, #fields_enum_name)
));
node_ref_enum_members.push(quote!(
#type_name(&'ast #type_name)
));
items.push(parse_quote!(
impl<'ast> From<&'ast #type_name> for NodeRef<'ast> {
fn from(node: &'ast #type_name) -> Self {
NodeRef::#type_name(node)
}
}
));
{
let mut arms = Vec::<Arm>::new();
for variant in &data.variants {
let variant_name = &variant.ident;
if variant.fields.len() != 1 {
continue;
}
for f in variant.fields.iter().filter(|f| is_node_ref(&f.ty)) {
let mut ty = &f.ty;
if let Some(inner) = extract_generic("Box", ty) {
ty = inner;
}
arms.push(parse_quote!(
#type_name::#variant_name(v0) => {
Box::new(::std::iter::once(NodeRef::#ty(v0)))
},
));
}
}
node_ref_iter_next_arms.push(parse_quote!(
NodeRef::#type_name(node) => {
match node {
#(#arms)*
_ => Box::new(::std::iter::empty::<NodeRef<'ast>>())
}
}
));
}
defs.push(parse_quote!(
impl #fields_enum_name {
#[inline(always)]
pub(crate) fn set_index(&mut self, _: usize) {
swc_visit::wrong_ast_path();
}
}
));
}
Item::Struct(data) => {
let mut set_index_arms = Vec::<Arm>::new();
for field in &data.fields {
let opt = field_variant(&type_name, field);
let opt = match opt {
Some(v) => v,
None => continue,
};
variants.push(opt.0);
set_index_arms.extend(opt.1);
}
kind_enum_members.push(quote!(
#type_name(#fields_enum_name)
));
parent_enum_members.push(quote!(
#type_name(&'ast #type_name, #fields_enum_name)
));
node_ref_enum_members.push(quote!(
#type_name(&'ast #type_name)
));
items.push(parse_quote!(
impl<'ast> From<&'ast #type_name> for NodeRef<'ast> {
fn from(node: &'ast #type_name) -> Self {
NodeRef::#type_name(node)
}
}
));
{
let mut iter: Expr = parse_quote!(::std::iter::empty::<NodeRef<'ast>>());
match &data.fields {
Fields::Named(fields) => {
for f in fields.named.iter() {
let ident = &f.ident;
let iter_expr =
to_iter(quote!(node.#ident), &f.ty, &node_names);
if let Some(iter_expr) = iter_expr {
iter = parse_quote!(#iter.chain(#iter_expr));
}
}
}
Fields::Unnamed(_fields) => {
}
Fields::Unit => {}
}
node_ref_iter_next_arms.push(parse_quote!(
NodeRef::#type_name(node) => {
let iterator = #iter;
Box::new(iterator)
}
));
}
defs.push(parse_quote!(
impl #fields_enum_name {
pub(crate) fn set_index(&mut self, index: usize) {
match self {
#(#set_index_arms)*
_ => {
swc_visit::wrong_ast_path()
}
}
}
}
));
}
_ => continue,
}
defs.push(parse_quote!(
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde-impl", derive(serde::Serialize, serde::Deserialize))]
pub enum #fields_enum_name {
#(#variants),*
}
))
}
{
defs.push(parse_quote!(
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde-impl", derive(serde::Serialize, serde::Deserialize))]
pub enum AstParentKind {
#(#kind_enum_members),*
}
));
defs.push(parse_quote!(
impl ::swc_visit::ParentKind for AstParentKind {
#[inline]
fn set_index(&mut self, index: usize) {
match self {
#(#kind_set_index_arms)*
}
}
}
));
}
{
defs.push(parse_quote!(
#[derive(Debug, Clone, Copy)]
pub enum AstParentNodeRef<'ast> {
#(#parent_enum_members),*
}
));
items.push(parse_quote!(
#[derive(Debug, Clone, Copy)]
pub enum NodeRef<'ast> {
#(#node_ref_enum_members),*
}
));
defs.push(parse_quote!(
impl<'ast> ::swc_visit::NodeRef for AstParentNodeRef<'ast> {
type ParentKind = AstParentKind;
#[inline(always)]
fn kind(&self) -> AstParentKind {
self.kind()
}
fn set_index(&mut self, index: usize) {
match self {
#(#node_ref_set_index_arms)*
}
}
}
));
defs.push(parse_quote!(
#[cfg(any(docsrs, feature = "path"))]
impl<'ast> AstParentNodeRef<'ast> {
#[inline]
pub fn kind(&self) -> AstParentKind {
match self {
#(#node_ref_kind_arms)*
}
}
}
));
items.push(parse_quote!(
impl<'ast> NodeRef<'ast> {
#[allow(unreachable_patterns)]
pub fn experimental_raw_children<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = NodeRef<'ast>>> {
match self {
#(#node_ref_iter_next_arms)*
}
}
}
));
items.push(parse_quote!(
impl<'ast> NodeRef<'ast> {
pub fn experimental_traverse(
&'ast self,
) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
let mut queue = std::collections::VecDeque::<NodeRef<'ast>>::new();
queue.push_back(*self);
Box::new(std::iter::from_fn(move || {
let node: NodeRef<'ast> = queue.pop_front()?;
{
let children = node.experimental_raw_children();
queue.extend(children);
}
Some(node)
}))
}
}
));
}
items.insert(
0,
parse_quote!(
#[cfg(any(docsrs, feature = "path"))]
pub mod fields {
#(#defs)*
}
),
);
items.push(parse_quote!(
#[cfg(any(docsrs, feature = "path"))]
pub use self::fields::{AstParentKind, AstParentNodeRef};
));
}
items
}