Codegen: Multi-pass refactoring & support for --check (#3918)

- Adds a new `CodeFormatter` trait.
- Adds a new trait-based generic `generate_code` helper, and reimplement
all language-specific backends in terms of it.
- Adds `--check` flag.
- Use `--check` instead of git dirtyness checks on CI.

---

- Fixes #3495

.github/workflows/contrib_checks.yml

@@ -92,11 +92,7 @@ jobs:
         uses: actions-rs/cargo@v1
         with:
           command: codegen
-          args: --force
-
-      - name: No Diffs From Running Codegen
-        run: |
-          git diff --exit-code
+          args: --force --check
 
   rs-lints:
     name: Rust lints (fmt, check, cranky, tests, doc)

.github/workflows/reusable_checks.yml

@@ -134,11 +134,7 @@ jobs:
         uses: actions-rs/cargo@v1
         with:
           command: codegen
-          args: --force
-
-      - name: No Diffs From Running Codegen
-        run: |
-          git diff --exit-code
+          args: --force --check
 
   rs-lints:
     name: Rust lints (fmt, check, cranky, tests, doc)

crates/re_types_builder/src/bin/build_re_types.rs

@@ -44,17 +44,20 @@ fn main() {
 
     let mut profiler = re_tracing::Profiler::default();
     let mut always_run = false;
+    let mut check = false;
 
     for arg in std::env::args().skip(1) {
         match arg.as_str() {
             "--help" => {
                 println!("Usage: [--help] [--force] [--profile]");
                 return;
             }
-            "--force" => {
+            "--force" => always_run = true,
+            "--profile" => profiler.start(),
+            "--check" => {
                 always_run = true;
+                check = true;
             }
-            "--profile" => profiler.start(),
             _ => {
                 eprintln!("Unknown argument: {arg:?}");
                 return;
@@ -114,26 +117,30 @@ fn main() {
             &reporter,
             cpp_output_dir_path,
             &objects,
-            &arrow_registry
+            &arrow_registry,
+            check,
         ),
         || re_types_builder::generate_rust_code(
             &reporter,
             workspace_dir,
             &objects,
-            &arrow_registry
+            &arrow_registry,
+            check,
         ),
         || re_types_builder::generate_python_code(
             &reporter,
             python_output_dir_path,
             python_testing_output_dir_path,
             &objects,
             &arrow_registry,
+            check,
         ),
         || re_types_builder::generate_docs(
             &reporter,
             docs_content_dir_path,
             &objects,
-            &arrow_registry
+            &arrow_registry,
+            check,
         ),
     );
 

crates/re_types_builder/src/codegen/mod.rs

@@ -1,5 +1,3 @@
-pub type GeneratedFiles = std::collections::BTreeMap<camino::Utf8PathBuf, String>;
-
 /// Implements the codegen pass.
 pub trait CodeGenerator {
     /// Generates user-facing code from [`crate::Objects`].
@@ -12,7 +10,7 @@ pub trait CodeGenerator {
         reporter: &crate::Reporter,
         objects: &crate::Objects,
         arrow_registry: &crate::ArrowRegistry,
-    ) -> GeneratedFiles;
+    ) -> crate::GeneratedFiles;
 }
 
 // ---

crates/re_types_builder/src/format/cpp.rs

@@ -0,0 +1,26 @@
+use crate::CodeFormatter;
+
+// ---
+
+pub struct CppCodeFormatter;
+
+impl CodeFormatter for CppCodeFormatter {
+    fn format(&mut self, _reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
+        use rayon::prelude::*;
+
+        re_tracing::profile_function!();
+
+        files.par_iter_mut().for_each(|(filepath, contents)| {
+            *contents = if matches!(filepath.extension(), Some("cpp" | "hpp")) {
+                format_code(contents)
+            } else {
+                contents.clone()
+            };
+        });
+    }
+}
+
+fn format_code(code: &str) -> String {
+    clang_format::clang_format_with_style(code, &clang_format::ClangFormatStyle::File)
+        .expect("Failed to run clang-format")
+}

crates/re_types_builder/src/format/mod.rs

@@ -0,0 +1,21 @@
+/// Implements the formatting pass.
+pub trait CodeFormatter {
+    /// Formats generated files in-place.
+    fn format(&mut self, reporter: &crate::Reporter, files: &mut crate::GeneratedFiles);
+}
+
+pub struct NoopCodeFormatter;
+
+impl CodeFormatter for NoopCodeFormatter {
+    fn format(&mut self, _reporter: &crate::Reporter, _files: &mut crate::GeneratedFiles) {}
+}
+
+// ---
+
+mod cpp;
+mod python;
+mod rust;
+
+pub use self::cpp::CppCodeFormatter;
+pub use self::python::PythonCodeFormatter;
+pub use self::rust::RustCodeFormatter;

crates/re_types_builder/src/format/python.rs

@@ -0,0 +1,145 @@
+use anyhow::Context;
+use camino::{Utf8Path, Utf8PathBuf};
+
+use crate::CodeFormatter;
+
+// ---
+
+pub struct PythonCodeFormatter {
+    pub pkg_path: Utf8PathBuf,
+    pub testing_pkg_path: Utf8PathBuf,
+}
+
+impl PythonCodeFormatter {
+    pub fn new(pkg_path: impl Into<Utf8PathBuf>, testing_pkg_path: impl Into<Utf8PathBuf>) -> Self {
+        Self {
+            pkg_path: pkg_path.into(),
+            testing_pkg_path: testing_pkg_path.into(),
+        }
+    }
+}
+
+impl CodeFormatter for PythonCodeFormatter {
+    fn format(&mut self, _reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
+        use rayon::prelude::*;
+
+        re_tracing::profile_function!();
+
+        // Running `black` once for each file is very slow, so we write all
+        // files to a temporary folder, format it, and copy back the results.
+
+        let tempdir = tempfile::tempdir().unwrap();
+        let tempdir_path = Utf8PathBuf::try_from(tempdir.path().to_owned()).unwrap();
+
+        files.par_iter().for_each(|(filepath, contents)| {
+            let formatted_source_path = format_path_for_tmp_dir(
+                &self.pkg_path,
+                &self.testing_pkg_path,
+                filepath,
+                &tempdir_path,
+            );
+            crate::codegen::common::write_file(&formatted_source_path, contents);
+        });
+
+        format_python_dir(&tempdir_path).unwrap();
+
+        // Read back and copy to the final destination:
+        files.par_iter_mut().for_each(|(filepath, contents)| {
+            let formatted_source_path = format_path_for_tmp_dir(
+                &self.pkg_path,
+                &self.testing_pkg_path,
+                filepath,
+                &tempdir_path,
+            );
+            *contents = std::fs::read_to_string(formatted_source_path).unwrap();
+        });
+    }
+}
+
+fn format_path_for_tmp_dir(
+    pkg_path: &Utf8Path,
+    testing_pkg_path: &Utf8Path,
+    filepath: &Utf8Path,
+    tempdir_path: &Utf8Path,
+) -> Utf8PathBuf {
+    // If the prefix is pkg_path, strip it, and then append to tempdir
+    // However, if the prefix is testing_pkg_path, strip it and insert an extra
+    // "testing" to avoid name collisions.
+    filepath.strip_prefix(pkg_path).map_or_else(
+        |_| {
+            tempdir_path
+                .join("testing")
+                .join(filepath.strip_prefix(testing_pkg_path).unwrap())
+        },
+        |f| tempdir_path.join(f),
+    )
+}
+
+fn format_python_dir(dir: &Utf8PathBuf) -> anyhow::Result<()> {
+    re_tracing::profile_function!();
+
+    // The order below is important and sadly we need to call black twice. Ruff does not yet
+    // fix line-length (See: https://github.com/astral-sh/ruff/issues/1904).
+    //
+    // 1) Call black, which among others things fixes line-length
+    // 2) Call ruff, which requires line-lengths to be correct
+    // 3) Call black again to cleanup some whitespace issues ruff might introduce
+
+    run_black_on_dir(dir).context("black")?;
+    run_ruff_on_dir(dir).context("ruff")?;
+    run_black_on_dir(dir).context("black")?;
+    Ok(())
+}
+
+fn python_project_path() -> Utf8PathBuf {
+    let path = crate::rerun_workspace_path()
+        .join("rerun_py")
+        .join("pyproject.toml");
+    assert!(path.exists(), "Failed to find {path:?}");
+    path
+}
+
+fn run_black_on_dir(dir: &Utf8PathBuf) -> anyhow::Result<()> {
+    re_tracing::profile_function!();
+    use std::process::{Command, Stdio};
+
+    let proc = Command::new("black")
+        .arg(format!("--config={}", python_project_path()))
+        .arg(dir)
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()?;
+
+    let output = proc.wait_with_output()?;
+
+    if output.status.success() {
+        Ok(())
+    } else {
+        let stdout = String::from_utf8_lossy(&output.stdout).trim().to_owned();
+        let stderr = String::from_utf8_lossy(&output.stderr).trim().to_owned();
+        anyhow::bail!("{stdout}\n{stderr}")
+    }
+}
+
+fn run_ruff_on_dir(dir: &Utf8PathBuf) -> anyhow::Result<()> {
+    re_tracing::profile_function!();
+    use std::process::{Command, Stdio};
+
+    let proc = Command::new("ruff")
+        .arg(format!("--config={}", python_project_path()))
+        .arg("--fix")
+        .arg(dir)
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()?;
+
+    let output = proc.wait_with_output()?;
+
+    if output.status.success() {
+        Ok(())
+    } else {
+        let stdout = String::from_utf8_lossy(&output.stdout).trim().to_owned();
+        let stderr = String::from_utf8_lossy(&output.stderr).trim().to_owned();
+        anyhow::bail!("{stdout}\n{stderr}")
+    }
+}

crates/re_types_builder/src/format/rust.rs

@@ -0,0 +1,48 @@
+use crate::CodeFormatter;
+
+// ---
+
+pub struct RustCodeFormatter;
+
+impl CodeFormatter for RustCodeFormatter {
+    fn format(&mut self, _reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
+        use rayon::prelude::*;
+
+        re_tracing::profile_function!();
+
+        files.par_iter_mut().for_each(|(filepath, contents)| {
+            *contents = if matches!(filepath.extension(), Some("rs")) {
+                format_code(contents)
+            } else {
+                contents.clone()
+            };
+        });
+    }
+}
+
+fn format_code(contents: &str) -> String {
+    re_tracing::profile_function!();
+
+    let mut contents = contents.replace(" :: ", "::"); // Fix `bytemuck :: Pod` -> `bytemuck::Pod`.
+
+    // Even though we already have used `prettyplease` we also
+    // need to run `cargo fmt`, since it catches some things `prettyplease` missed.
+    // We need to run `cago fmt` several times because it is not idempotent;
+    // see https://github.com/rust-lang/rustfmt/issues/5824
+    for _ in 0..2 {
+        // NOTE: We're purposefully ignoring the error here.
+        //
+        // In the very unlikely chance that the user doesn't have the `fmt` component installed,
+        // there's still no good reason to fail the build.
+        //
+        // The CI will catch the unformatted file at PR time and complain appropriately anyhow.
+
+        re_tracing::profile_scope!("rust-fmt");
+        use rust_format::Formatter as _;
+        if let Ok(formatted) = rust_format::RustFmt::default().format_str(&contents) {
+            contents = formatted;
+        }
+    }
+
+    contents
+}