refactor(ast_codegen): cleanup usages of the generated outputs.

tasks/ast_codegen/src/generators/ast_builder.rs

@@ -37,39 +37,42 @@ impl Generator for AstBuilderGenerator {
 
         let header = generated_header!();
 
-        GeneratorOutput::One(quote! {
-            #header
-            insert!("#![allow(clippy::default_trait_access, clippy::too_many_arguments, clippy::fn_params_excessive_bools)]");
-            endl!();
-
-            use oxc_allocator::{Allocator, Box, IntoIn, Vec};
-            use oxc_span::{Atom, SourceType, Span};
-            use oxc_syntax::{
-                number::{BigintBase, NumberBase},
-                operator::{
-                    AssignmentOperator, BinaryOperator, LogicalOperator, UnaryOperator, UpdateOperator,
-                },
-            };
-
-            endl!();
-
-            #[allow(clippy::wildcard_imports)]
-            use crate::ast::*;
-
-            endl!();
-
-            /// AST builder for creating AST nodes
-            #[derive(Clone, Copy)]
-            pub struct AstBuilder<'a> {
-                pub allocator: &'a Allocator,
-            }
+        GeneratorOutput::Stream((
+            "ast_builder",
+            quote! {
+                #header
+                insert!("#![allow(clippy::default_trait_access, clippy::too_many_arguments, clippy::fn_params_excessive_bools)]");
+                endl!();
+
+                use oxc_allocator::{Allocator, Box, IntoIn, Vec};
+                use oxc_span::{Atom, SourceType, Span};
+                use oxc_syntax::{
+                    number::{BigintBase, NumberBase},
+                    operator::{
+                        AssignmentOperator, BinaryOperator, LogicalOperator, UnaryOperator, UpdateOperator,
+                    },
+                };
+
+                endl!();
+
+                #[allow(clippy::wildcard_imports)]
+                use crate::ast::*;
+
+                endl!();
+
+                /// AST builder for creating AST nodes
+                #[derive(Clone, Copy)]
+                pub struct AstBuilder<'a> {
+                    pub allocator: &'a Allocator,
+                }
 
-            endl!();
+                endl!();
 
-            impl<'a> AstBuilder<'a> {
-                #(#fns)*
-            }
-        })
+                impl<'a> AstBuilder<'a> {
+                    #(#fns)*
+                }
+            },
+        ))
     }
 }
 

tasks/ast_codegen/src/generators/ast_kind.rs

@@ -175,37 +175,40 @@ impl Generator for AstKindGenerator {
 
         let header = generated_header!();
 
-        GeneratorOutput::One(quote! {
-            #header
+        GeneratorOutput::Stream((
+            "ast_kind",
+            quote! {
+                #header
 
-            use crate::ast::*;
-            use oxc_span::{GetSpan, Span};
+                use crate::ast::*;
+                use oxc_span::{GetSpan, Span};
 
-            endl!();
+                endl!();
 
-            #[derive(Debug, Clone, Copy)]
-            pub enum AstType {
-                #(#types),*,
-            }
+                #[derive(Debug, Clone, Copy)]
+                pub enum AstType {
+                    #(#types),*,
+                }
 
-            endl!();
+                endl!();
 
-            /// Untyped AST Node Kind
-            #[derive(Debug, Clone, Copy)]
-            pub enum AstKind<'a> {
-                #(#kinds),*,
-            }
+                /// Untyped AST Node Kind
+                #[derive(Debug, Clone, Copy)]
+                pub enum AstKind<'a> {
+                    #(#kinds),*,
+                }
 
-            endl!();
+                endl!();
 
-            impl<'a> GetSpan for AstKind<'a> {
-                #[allow(clippy::match_same_arms)]
-                fn span(&self) -> Span {
-                    match self {
-                        #(#span_matches),*,
+                impl<'a> GetSpan for AstKind<'a> {
+                    #[allow(clippy::match_same_arms)]
+                    fn span(&self) -> Span {
+                        match self {
+                            #(#span_matches),*,
+                        }
                     }
                 }
-            }
-        })
+            },
+        ))
     }
 }

tasks/ast_codegen/src/generators/impl_get_span.rs

@@ -36,16 +36,19 @@ impl Generator for ImplGetSpanGenerator {
 
         let header = generated_header!();
 
-        GeneratorOutput::One(quote! {
-            #header
-            insert!("#![allow(clippy::match_same_arms)]");
-            endl!();
+        GeneratorOutput::Stream((
+            "span",
+            quote! {
+                #header
+                insert!("#![allow(clippy::match_same_arms)]");
+                endl!();
 
-            use crate::ast::*;
-            use oxc_span::{GetSpan, Span};
+                use crate::ast::*;
+                use oxc_span::{GetSpan, Span};
 
-            #(#impls)*
-        })
+                #(#impls)*
+            },
+        ))
     }
 }
 

tasks/ast_codegen/src/generators/mod.rs

@@ -1,4 +1,3 @@
-mod ast;
 mod ast_builder;
 mod ast_kind;
 mod impl_get_span;
@@ -40,8 +39,7 @@ macro_rules! generated_header {
 pub(crate) use generated_header;
 pub(crate) use insert;
 
-pub use ast::AstGenerator;
 pub use ast_builder::AstBuilderGenerator;
 pub use ast_kind::AstKindGenerator;
 pub use impl_get_span::ImplGetSpanGenerator;
-pub use visit::VisitGenerator;
+pub use visit::{VisitGenerator, VisitMutGenerator};

tasks/ast_codegen/src/generators/visit.rs

@@ -29,17 +29,25 @@ use crate::{
 use super::generated_header;
 
 pub struct VisitGenerator;
+pub struct VisitMutGenerator;
 
 impl Generator for VisitGenerator {
     fn name(&self) -> &'static str {
         "VisitGenerator"
     }
 
     fn generate(&mut self, ctx: &CodegenCtx) -> GeneratorOutput {
-        let visit = (String::from("visit"), generate_visit(ctx));
-        let visit_mut = (String::from("visit_mut"), generate_visit_mut(ctx));
+        GeneratorOutput::Stream(("visit", generate_visit(ctx)))
+    }
+}
+
+impl Generator for VisitMutGenerator {
+    fn name(&self) -> &'static str {
+        "VisitMutGenerator"
+    }
 
-        GeneratorOutput::Many(HashMap::from_iter(vec![visit, visit_mut]))
+    fn generate(&mut self, ctx: &CodegenCtx) -> GeneratorOutput {
+        GeneratorOutput::Stream(("visit_mut", generate_visit_mut(ctx)))
     }
 }
 

tasks/ast_codegen/src/main.rs

@@ -23,7 +23,7 @@ use proc_macro2::TokenStream;
 use syn::parse_file;
 
 use defs::TypeDef;
-use generators::{AstBuilderGenerator, AstGenerator, AstKindGenerator, VisitGenerator};
+use generators::{AstBuilderGenerator, AstKindGenerator, VisitGenerator, VisitMutGenerator};
 use linker::{linker, Linker};
 use schema::{Inherit, Module, REnum, RStruct, RType, Schema};
 
@@ -47,42 +47,55 @@ trait Generator {
     fn generate(&mut self, ctx: &CodegenCtx) -> GeneratorOutput;
 }
 
+type GeneratedStream = (&'static str, TokenStream);
+
 #[derive(Debug, Clone)]
 enum GeneratorOutput {
     None,
-    One(TokenStream),
-    Many(HashMap<String, TokenStream>),
     Info(String),
+    Stream(GeneratedStream),
 }
 
 impl GeneratorOutput {
-    pub fn as_none(&self) {
-        assert!(matches!(self, Self::None));
+    pub fn is_none(&self) -> bool {
+        matches!(self, Self::None)
+    }
+
+    pub fn expect_none(&self) {
+        assert!(self.is_none());
     }
 
-    pub fn as_one(&self) -> &TokenStream {
-        if let Self::One(it) = self {
+    pub fn to_info(&self) -> &String {
+        if let Self::Info(it) = self {
             it
         } else {
             panic!();
         }
     }
 
-    pub fn as_many(&self) -> &HashMap<String, TokenStream> {
-        if let Self::Many(it) = self {
+    pub fn to_stream(&self) -> &GeneratedStream {
+        if let Self::Stream(it) = self {
             it
         } else {
             panic!();
         }
     }
 
-    pub fn as_info(&self) -> &String {
+    pub fn into_info(self) -> String {
         if let Self::Info(it) = self {
             it
         } else {
             panic!();
         }
     }
+
+    pub fn into_stream(self) -> GeneratedStream {
+        if let Self::Stream(it) = self {
+            it
+        } else {
+            panic!();
+        }
+    }
 }
 
 struct CodegenCtx {
@@ -179,71 +192,41 @@ fn files() -> impl std::iter::Iterator<Item = String> {
     vec![path("literal"), path("js"), path("ts"), path("jsx")].into_iter()
 }
 
-fn output_dir() -> Result<String> {
+fn output_dir() -> std::io::Result<String> {
     let dir = format!("{AST_ROOT_DIR}/src/generated");
-    fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
+    fs::create_dir_all(&dir)?;
     Ok(dir)
 }
 
+fn write_generated_streams(
+    streams: impl IntoIterator<Item = GeneratedStream>,
+) -> std::io::Result<()> {
+    let output_dir = output_dir()?;
+    for (name, stream) in streams {
+        let content = pprint(&stream);
+
+        let path = format!("{output_dir}/{name}.rs");
+        let mut file = fs::File::create(path)?;
+
+        file.write_all(content.as_bytes())?;
+    }
+    Ok(())
+}
+
 #[allow(clippy::print_stdout)]
 fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
     let CodegenResult { outputs, .. } = files()
         .fold(AstCodegen::default(), AstCodegen::add_file)
-        .with(AstGenerator)
         .with(AstKindGenerator)
         .with(AstBuilderGenerator)
         .with(ImplGetSpanGenerator)
         .with(VisitGenerator)
+        .with(VisitMutGenerator)
         .generate()?;
 
-    let output_dir = output_dir()?;
-    let outputs: HashMap<_, _> = outputs.into_iter().collect();
-
-    {
-        // write `span.rs` file
-        let output = outputs[ImplGetSpanGenerator.name()].as_one();
-        let span_content = pprint(output);
-
-        let path = format!("{output_dir}/span.rs");
-        let mut file = fs::File::create(path)?;
-
-        file.write_all(span_content.as_bytes())?;
-    }
-
-    {
-        // write `ast_kind.rs` file
-        let output = outputs[AstKindGenerator.name()].as_one();
-        let span_content = pprint(output);
-
-        let path = format!("{output_dir}/ast_kind.rs");
-        let mut file = fs::File::create(path)?;
-
-        file.write_all(span_content.as_bytes())?;
-    }
-
-    {
-        // write `ast_builder.rs` file
-        let output = outputs[AstBuilderGenerator.name()].as_one();
-        let span_content = pprint(output);
-
-        let path = format!("{output_dir}/ast_builder.rs");
-        let mut file = fs::File::create(path)?;
-
-        file.write_all(span_content.as_bytes())?;
-    }
-
-    {
-        // write `visit.rs` and `visit_mut.rs` files
-        let output = outputs[VisitGenerator.name()].as_many();
-        let content = pprint(&output["visit"]);
-        let content_mut = pprint(&output["visit_mut"]);
-
-        let mut visit = fs::File::create(format!("{output_dir}/visit.rs"))?;
-        let mut visit_mut = fs::File::create(format!("{output_dir}/visit_mut.rs"))?;
-
-        visit.write_all(content.as_bytes())?;
-        visit_mut.write_all(content_mut.as_bytes())?;
-    }
+    let (streams, _): (Vec<_>, Vec<_>) =
+        outputs.into_iter().partition(|it| matches!(it.1, GeneratorOutput::Stream(_)));
+    write_generated_streams(streams.into_iter().map(|it| it.1.into_stream()))?;
 
     cargo_fmt(".")?;