Merge pull request #17682 from github/redsun82/rust-codegen

Rust: integrate rust code generation into `//rust/codegen`

Cargo.toml

@@ -7,7 +7,7 @@ members = [
     "ruby/extractor",
     "rust/extractor",
     "rust/extractor/macros",
-    "rust/generate-schema",
+    "rust/ast-generator",
 ]
 
 [patch.crates-io]

MODULE.bazel

@@ -61,7 +61,7 @@ r.from_cargo(
         "//ruby/extractor:Cargo.toml",
         "//rust/extractor:Cargo.toml",
         "//rust/extractor/macros:Cargo.toml",
-        "//rust/generate-schema:Cargo.toml",
+        "//rust/ast-generator:Cargo.toml",
         "//shared/tree-sitter-extractor:Cargo.toml",
     ],
 )

rust/generate-schema/BUILD.bazel → rust/ast-generator/BUILD.bazel

@@ -1,7 +1,8 @@
+load("@tree_sitter_extractors_deps//:defs.bzl", "aliases", "all_crate_deps")
 load("//misc/bazel:rust.bzl", "codeql_rust_binary")
 
 codeql_rust_binary(
-    name = "generate-schema",
+    name = "ast-generator",
     srcs = glob(["src/**/*.rs"]),
     aliases = aliases(),
     proc_macro_deps = all_crate_deps(
@@ -12,3 +13,9 @@ codeql_rust_binary(
         normal = True,
     ),
 )
+
+filegroup(
+    name = "manifest",
+    srcs = ["Cargo.toml"],
+    visibility = ["//rust:__subpackages__"],
+)

rust/generate-schema/Cargo.toml → rust/ast-generator/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "generate-schema"
+name = "ast-generator"
 version = "0.1.0"
 edition = "2021"
 
@@ -8,4 +8,3 @@ ungrammar = "1.16.1"
 proc-macro2 = "1.0.33"
 quote = "1.0.12"
 itertools = "0.10.1"
-

rust/generate-schema/src/codegen.rs → rust/ast-generator/src/codegen.rs

@@ -12,5 +12,5 @@ pub fn ensure_file_contents(
     contents: &String,
     _check: bool,
 ) {
-    std::fs::write(path, contents).expect("Unable to write file");
+    std::fs::write(path, contents).unwrap_or_else(|_| panic!("Unable to write {}", path.display()));
 }

rust/generate-schema/src/main.rs → rust/ast-generator/src/main.rs

@@ -55,29 +55,38 @@ fn write_schema(
     let mut buf: Vec<u8> = Vec::new();
     writeln!(
         buf,
-        "# Generated by `cargo generate-schema`, do not edit by hand.\n"
+        "# Generated by `ast-generator`, do not edit by hand.\n"
     )?;
-    writeln!(buf, "from .prelude import *\n")?;
+    writeln!(buf, "from .prelude import *")?;
 
     for node in &grammar.enums {
         let super_classses = if let Some(cls) = super_types.get(&node.name) {
-            let super_classes: Vec<String> = cls.iter().map(|x| class_name(x)).collect();
+            let super_classes: Vec<String> = cls.iter().map(class_name).collect();
             super_classes.join(",")
         } else {
             "AstNode".to_owned()
         };
-        writeln!(buf, "class {}({}):", class_name(&node.name), super_classses)?;
+        writeln!(
+            buf,
+            "\nclass {}({}):",
+            class_name(&node.name),
+            super_classses
+        )?;
         writeln!(buf, "   pass")?;
-        writeln!(buf, "")?;
     }
     for node in &grammar.nodes {
         let super_classses = if let Some(cls) = super_types.get(&node.name) {
-            let super_classes: Vec<String> = cls.iter().map(|x| class_name(x)).collect();
+            let super_classes: Vec<String> = cls.iter().map(class_name).collect();
             super_classes.join(",")
         } else {
             "AstNode".to_owned()
         };
-        writeln!(buf, "class {}({}):", class_name(&node.name), super_classses)?;
+        writeln!(
+            buf,
+            "\nclass {}({}):",
+            class_name(&node.name),
+            super_classses
+        )?;
         let mut empty = true;
         for field in get_fields(node) {
             if field.tp == "SyntaxToken" {
@@ -111,7 +120,6 @@ fn write_schema(
         if empty {
             writeln!(buf, "   pass")?;
         }
-        writeln!(buf, "")?;
     }
     Ok(String::from_utf8_lossy(&buf).to_string())
 }
@@ -408,7 +416,9 @@ fn write_extractor(grammar: &AstSrc) -> std::io::Result<String> {
     let mut buf: Vec<u8> = Vec::new();
     writeln!(
         buf,
-        "//! Generated by `cargo generate-schema`, do not edit by hand.\n
+        "//! Generated by `ast-generator`, do not edit by hand.\n
+#![cfg_attr(any(), rustfmt::skip)]
+
 use crate::generated;
 use super::base::{{TextValue, Translator}};
 use crate::trap::{{Label, TrapId}};
@@ -462,7 +472,7 @@ impl Translator {{
             type_name,
             class_name
         )?;
-        for field in get_fields(&node) {
+        for field in get_fields(node) {
             if &field.tp == "SyntaxToken" {
                 continue;
             }
@@ -500,7 +510,7 @@ impl Translator {{
             class_name
         )?;
         writeln!(buf, "            id: TrapId::Star,")?;
-        for field in get_fields(&node) {
+        for field in get_fields(node) {
             if field.tp == "SyntaxToken" {
                 continue;
             }
@@ -523,7 +533,7 @@ impl Translator {{
 }
 
 fn main() -> std::io::Result<()> {
-    let grammar: Grammar = fs::read_to_string(project_root().join("generate-schema/rust.ungram"))
+    let grammar: Grammar = fs::read_to_string(project_root().join("ast-generator/rust.ungram"))
         .unwrap()
         .parse()
         .unwrap();
@@ -537,9 +547,7 @@ fn main() -> std::io::Result<()> {
     let mut super_types: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
     for node in &grammar.enums {
         for variant in &node.variants {
-            let set = super_types
-                .entry(variant.to_owned())
-                .or_insert_with(|| BTreeSet::new());
+            let set = super_types.entry(variant.to_owned()).or_default();
             set.insert(node.name.to_owned());
         }
     }
@@ -550,7 +558,7 @@ fn main() -> std::io::Result<()> {
         super_class_x.cmp(&super_class_y).then(x.name.cmp(&y.name))
     });
     let schema = write_schema(&grammar, super_types)?;
-    let schema_path = PathBuf::from("../schema/ast.py");
+    let schema_path = project_root().join("schema/ast.py");
     codegen::ensure_file_contents(
         crate::flags::CodegenType::Grammar,
         &schema_path,
@@ -559,7 +567,7 @@ fn main() -> std::io::Result<()> {
     );
 
     let extractor = write_extractor(&grammar)?;
-    let extractor_path = PathBuf::from("../extractor/src/translate/generated.rs");
+    let extractor_path = project_root().join("extractor/src/translate/generated.rs");
     codegen::ensure_file_contents(
         crate::flags::CodegenType::Grammar,
         &extractor_path,

rust/codegen/BUILD.bazel

@@ -1,15 +1,17 @@
-load("@bazel_skylib//rules:native_binary.bzl", "native_binary")
+_args = [
+    "//rust/ast-generator",
+    "//rust/ast-generator:manifest",
+    "//misc/codegen",
+    "//rust:codegen-conf",
+]
 
-native_binary(
+sh_binary(
     name = "codegen",
-    src = "//misc/codegen",
-    out = "codegen",
-    args = [
-        "--configuration-file=$(location //rust:codegen-conf)",
-    ],
-    data = [
-        "//rust:codegen-conf",
-        "//rust:schema",
-    ],
+    srcs = ["codegen.sh"],
+    args = ["$(rlocationpath %s)" % a for a in _args],
+    data = _args,
     visibility = ["//rust:__subpackages__"],
+    deps = [
+        "//misc/bazel:sh_runfiles",
+    ],
 )

rust/codegen/codegen.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+set -eu
+
+source misc/bazel/runfiles.sh 2>/dev/null || source external/ql~/misc/bazel/runfiles.sh
+
+ast_generator="$(rlocation "$1")"
+ast_generator_manifest="$(rlocation "$2")"
+codegen="$(rlocation "$3")"
+codegen_conf="$(rlocation "$4")"
+
+CARGO_MANIFEST_DIR="$(dirname "$ast_generator_manifest")" "$ast_generator"
+"$codegen" --configuration-file="$codegen_conf"