From 0b433108be3b34c42ed15517f8b430298efe7c37 Mon Sep 17 00:00:00 2001
From: rzvxa <3788964+rzvxa@users.noreply.github.com>
Date: Tue, 9 Jul 2024 03:42:11 +0000
Subject: [PATCH] feat(ast_codegen): add cli arguments to use with
 `oxc_ast_codegen` executable. (#4117)

```
Usage: oxc_ast_codegen [--dry-run] [--no-fmt] [--schema=ARG]

Available options:
        --dry-run     Runs all generators but won't write anything down.
        --no-fmt      Don't run cargo fmt at the end
        --schema=ARG  Path of output `schema.json`.
    -h, --help        Prints help information
```
---
 Cargo.lock                    |  2 ++
 tasks/ast_codegen/Cargo.toml  |  2 ++
 tasks/ast_codegen/src/main.rs | 40 +++++++++++++++++++++++++++--------
 tasks/ast_codegen/src/util.rs |  7 ++++++
 4 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 3c31cc71c0e37..56388adba9f6c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1337,6 +1337,7 @@ dependencies = [
 name = "oxc_ast_codegen"
 version = "0.0.0"
 dependencies = [
+ "bpaf",
  "convert_case",
  "itertools 0.13.0",
  "lazy_static",
@@ -1345,6 +1346,7 @@ dependencies = [
  "quote",
  "regex",
  "serde",
+ "serde_json",
  "syn",
 ]
 
diff --git a/tasks/ast_codegen/Cargo.toml b/tasks/ast_codegen/Cargo.toml
index 012c27454bd2c..16d7d56af6e4a 100644
--- a/tasks/ast_codegen/Cargo.toml
+++ b/tasks/ast_codegen/Cargo.toml
@@ -27,7 +27,9 @@ quote = { workspace = true }
 proc-macro2 = { workspace = true }
 itertools = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
 regex = { workspace = true }
 prettyplease = { workspace = true }
 lazy_static = { workspace = true }
 convert_case = { workspace = true }
+bpaf = { workspace = true, features = ["autocomplete", "bright-color", "derive"] }
diff --git a/tasks/ast_codegen/src/main.rs b/tasks/ast_codegen/src/main.rs
index 6e4a2c7e3c2bb..a5802b586c4b6 100644
--- a/tasks/ast_codegen/src/main.rs
+++ b/tasks/ast_codegen/src/main.rs
@@ -17,6 +17,7 @@ use std::{
     rc::Rc,
 };
 
+use bpaf::{Bpaf, Parser};
 use fmt::{cargo_fmt, pprint};
 use itertools::Itertools;
 use proc_macro2::TokenStream;
@@ -26,6 +27,7 @@ use defs::TypeDef;
 use generators::{AstBuilderGenerator, AstKindGenerator, VisitGenerator, VisitMutGenerator};
 use linker::{linker, Linker};
 use schema::{Inherit, Module, REnum, RStruct, RType, Schema};
+use util::write_all_to;
 
 use crate::generators::ImplGetSpanGenerator;
 
@@ -204,18 +206,29 @@ fn write_generated_streams(
     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())?;
+        write_all_to(content.as_bytes(), path)?;
     }
     Ok(())
 }
 
+#[derive(Debug, Bpaf)]
+pub struct CliOptions {
+    /// Runs all generators but won't write anything down.
+    #[bpaf(switch)]
+    dry_run: bool,
+    /// Don't run cargo fmt at the end
+    #[bpaf(switch)]
+    no_fmt: bool,
+    /// Path of output `schema.json`.
+    schema: Option<std::path::PathBuf>,
+}
+
 #[allow(clippy::print_stdout)]
 fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
-    let CodegenResult { outputs, .. } = files()
+    let cli_options = cli_options().run();
+
+    let CodegenResult { outputs, schema } = files()
         .fold(AstCodegen::default(), AstCodegen::add_file)
         .with(AstKindGenerator)
         .with(AstBuilderGenerator)
@@ -226,11 +239,20 @@ fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
 
     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(".")?;
+    if !cli_options.dry_run {
+        write_generated_streams(streams.into_iter().map(|it| it.1.into_stream()))?;
+    }
+
+    if !cli_options.no_fmt {
+        cargo_fmt(".")?;
+    }
+
+    if let CliOptions { schema: Some(schema_path), dry_run: false, .. } = cli_options {
+        let path = schema_path.to_str().expect("invalid path for schema output.");
+        let schema = serde_json::to_string_pretty(&schema).map_err(|e| e.to_string())?;
+        write_all_to(schema.as_bytes(), path)?;
+    }
 
-    // let schema = serde_json::to_string_pretty(&schema).map_err(|e| e.to_string())?;
-    // println!("{schema}");
     Ok(())
 }
diff --git a/tasks/ast_codegen/src/util.rs b/tasks/ast_codegen/src/util.rs
index 64b625059d11b..a48dd4e9e6e1a 100644
--- a/tasks/ast_codegen/src/util.rs
+++ b/tasks/ast_codegen/src/util.rs
@@ -220,3 +220,10 @@ impl TokenStreamExt for TokenStream {
             .collect()
     }
 }
+
+pub fn write_all_to<S: AsRef<str>>(data: &[u8], path: S) -> std::io::Result<()> {
+    use std::{fs, io::Write};
+    let mut file = fs::File::create(path.as_ref())?;
+    file.write_all(data)?;
+    Ok(())
+}