improvement(ast_codegen): use visit_as attribute to create aliased …

…`AstKind`s. (#4011)

crates/oxc_ast/src/generated/ast_kind.rs

@@ -71,6 +71,7 @@ pub enum AstType {
     LabeledStatement,
     ThrowStatement,
     TryStatement,
+    FinallyClause,
     CatchClause,
     CatchParameter,
     DebuggerStatement,
@@ -85,6 +86,7 @@ pub enum AstType {
     ArrowFunctionExpression,
     YieldExpression,
     Class,
+    ClassHeritage,
     ClassBody,
     MethodDefinition,
     PropertyDefinition,
@@ -164,8 +166,6 @@ pub enum AstType {
     JSXSpreadAttribute,
     JSXIdentifier,
     JSXText,
-    FinallyClause,
-    ClassHeritage,
     ExpressionArrayElement,
 }
 
@@ -237,6 +237,7 @@ pub enum AstKind<'a> {
     LabeledStatement(&'a LabeledStatement<'a>),
     ThrowStatement(&'a ThrowStatement<'a>),
     TryStatement(&'a TryStatement<'a>),
+    FinallyClause(&'a BlockStatement<'a>),
     CatchClause(&'a CatchClause<'a>),
     CatchParameter(&'a CatchParameter<'a>),
     DebuggerStatement(&'a DebuggerStatement),
@@ -251,6 +252,7 @@ pub enum AstKind<'a> {
     ArrowFunctionExpression(&'a ArrowFunctionExpression<'a>),
     YieldExpression(&'a YieldExpression<'a>),
     Class(&'a Class<'a>),
+    ClassHeritage(&'a Expression<'a>),
     ClassBody(&'a ClassBody<'a>),
     MethodDefinition(&'a MethodDefinition<'a>),
     PropertyDefinition(&'a PropertyDefinition<'a>),
@@ -330,8 +332,6 @@ pub enum AstKind<'a> {
     JSXSpreadAttribute(&'a JSXSpreadAttribute<'a>),
     JSXIdentifier(&'a JSXIdentifier<'a>),
     JSXText(&'a JSXText<'a>),
-    FinallyClause(&'a BlockStatement<'a>),
-    ClassHeritage(&'a Expression<'a>),
     ExpressionArrayElement(&'a Expression<'a>),
 }
 
@@ -404,6 +404,7 @@ impl<'a> GetSpan for AstKind<'a> {
             Self::LabeledStatement(it) => it.span(),
             Self::ThrowStatement(it) => it.span(),
             Self::TryStatement(it) => it.span(),
+            Self::FinallyClause(it) => it.span(),
             Self::CatchClause(it) => it.span(),
             Self::CatchParameter(it) => it.span(),
             Self::DebuggerStatement(it) => it.span(),
@@ -418,6 +419,7 @@ impl<'a> GetSpan for AstKind<'a> {
             Self::ArrowFunctionExpression(it) => it.span(),
             Self::YieldExpression(it) => it.span(),
             Self::Class(it) => it.span(),
+            Self::ClassHeritage(it) => it.span(),
             Self::ClassBody(it) => it.span(),
             Self::MethodDefinition(it) => it.span(),
             Self::PropertyDefinition(it) => it.span(),
@@ -497,8 +499,6 @@ impl<'a> GetSpan for AstKind<'a> {
             Self::JSXSpreadAttribute(it) => it.span(),
             Self::JSXIdentifier(it) => it.span(),
             Self::JSXText(it) => it.span(),
-            Self::FinallyClause(it) => it.span(),
-            Self::ClassHeritage(it) => it.span(),
             Self::ExpressionArrayElement(it) => it.span(),
         }
     }

tasks/ast_codegen/src/generators/ast_kind.rs

@@ -2,7 +2,7 @@ use itertools::Itertools;
 use quote::quote;
 use syn::{parse_quote, Arm, Ident, Type, Variant};
 
-use crate::{schema::RType, CodegenCtx, Generator, GeneratorOutput};
+use crate::{schema::RType, util::TypeExt, CodegenCtx, Generator, GeneratorOutput, TypeRef};
 
 use super::generated_header;
 
@@ -84,13 +84,64 @@ pub fn blacklist((ident, _): &(Ident, Type)) -> bool {
     !BLACK_LIST.contains(&ident.to_string().as_str())
 }
 
-pub fn aliased_nodes() -> [(Ident, Type); 3] {
+pub fn aliased_nodes() -> [(Ident, Type); 1] {
     use syn::parse_quote as pq;
-    [
-        (pq!(FinallyClause), pq!(BlockStatement<'a>)),
-        (pq!(ClassHeritage), pq!(Expression<'a>)),
-        (pq!(ExpressionArrayElement), pq!(Expression<'a>)),
-    ]
+    [(pq!(ExpressionArrayElement), pq!(Expression<'a>))]
+}
+
+pub fn process_types(ty: &TypeRef) -> Vec<(Ident, Type)> {
+    let aliases = match &*ty.borrow() {
+        RType::Enum(enum_) => enum_
+            .item
+            .variants
+            .iter()
+            .filter_map(|it| {
+                it.attrs
+                    .iter()
+                    .find(|it| it.path().is_ident("visit_as"))
+                    .map(|attr| (it, attr))
+                    .map(|(it, attr)| {
+                        assert!(
+                            it.fields.len() == 1,
+                            "visit_as only supports single argument fields."
+                        );
+                        let field = it.fields.iter().next().unwrap();
+                        let type_name = field.ty.get_ident().inner_ident();
+                        (attr.parse_args().unwrap(), parse_quote!(#type_name<'a>))
+                    })
+            })
+            .collect_vec(),
+        RType::Struct(struct_) => struct_
+            .item
+            .fields
+            .iter()
+            .filter_map(|it| {
+                it.attrs
+                    .iter()
+                    .find(|it| it.path().is_ident("visit_as"))
+                    .map(|attr| (it, attr))
+                    .map(|(field, attr)| {
+                        let type_name = field.ty.get_ident().inner_ident();
+                        (attr.parse_args().unwrap(), parse_quote!(#type_name<'a>))
+                    })
+            })
+            .collect_vec(),
+        _ => panic!(),
+    };
+
+    Some(ty)
+        .into_iter()
+        .map(|kind| {
+            if let kind @ (RType::Enum(_) | RType::Struct(_)) = &*kind.borrow() {
+                let ident = kind.ident().unwrap().clone();
+                let typ = kind.as_type().unwrap();
+                (ident, typ)
+            } else {
+                panic!()
+            }
+        })
+        .chain(aliases)
+        .collect()
 }
 
 impl Generator for AstKindGenerator {
@@ -103,14 +154,10 @@ impl Generator for AstKindGenerator {
             .ty_table
             .iter()
             .filter(|it| it.borrow().visitable())
-            .filter_map(|maybe_kind| match &*maybe_kind.borrow() {
-                kind @ (RType::Enum(_) | RType::Struct(_)) if kind.visitable() => {
-                    let ident = kind.ident().unwrap().clone();
-                    let typ = kind.as_type().unwrap();
-                    Some((ident, typ))
-                }
-                _ => None,
-            })
+            .filter(
+                |maybe_kind| matches!(&*maybe_kind.borrow(), kind @ (RType::Enum(_) | RType::Struct(_)) if kind.visitable())
+            )
+            .flat_map(process_types)
             .filter(blacklist)
             .chain(aliased_nodes())
             .collect();