vyre-build-scan 0.1.0

//! Rust-source `OpSpec` registry generation.
//!
//! # build.rs example
//!
//! ```ignore
//! fn main() {
//!     vyre_build_scan::scan_rust_specs(&vyre_build_scan::RustSpecRegistry {
//!         scan_dirs: &["src/ops"],
//!         const_name: "GENERATED_REGISTRY",
//!         output_file: "ops_registry.rs",
//!     });
//! }
//! ```

use crate::config::RustSpecRegistry;
use crate::fatal::fatal;
use crate::paths::output_path_in_out_dir;
use std::collections::BTreeSet;
use std::env;
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};

/// Scan Rust source trees for public `OpSpec` constants and emit a registry.
///
/// The scanner accepts top-level `pub const SPEC: OpSpec` declarations and
/// associated `impl Type { pub const SPEC: OpSpec = ... }` declarations when
/// the source file is reachable from its parent module.
///
/// # build.rs example
///
/// ```ignore
/// fn main() {
///     vyre_build_scan::scan_rust_specs(&vyre_build_scan::RustSpecRegistry {
///         scan_dirs: &["src/ops/primitive", "src/ops/hash"],
///         const_name: "GENERATED_REGISTRY",
///         output_file: "ops_registry.rs",
///     });
/// }
/// ```
pub fn scan_rust_specs(registry: &RustSpecRegistry<'_>) {
    let manifest_dir = PathBuf::from(
        env::var("CARGO_MANIFEST_DIR")
            .expect("build_scan: CARGO_MANIFEST_DIR must be set by cargo"),
    );
    let out_dir =
        PathBuf::from(env::var("OUT_DIR").expect("build_scan: OUT_DIR must be set by cargo"));

    let mut specs = BTreeSet::new();
    for scan_dir in registry.scan_dirs {
        let scan_path = manifest_dir.join(scan_dir);
        println!("cargo:rerun-if-changed={}", scan_path.display());
        collect_rust_specs(&manifest_dir, &scan_path, &mut specs);
    }

    let mut out = String::new();
    out.push_str("// @generated by vyre-build-scan - do not edit.\n");
    out.push_str("// Source: public OpSpec constants discovered from Rust source files.\n\n");
    out.push_str(&format!(
        "static {}: &[&OpSpec] = &[\n",
        registry.const_name,
    ));
    for spec in specs {
        out.push_str(&format!("    &{spec},\n"));
    }
    out.push_str("];\n");

    let output_path = output_path_in_out_dir(&out_dir, registry.output_file)
        .unwrap_or_else(|message| fatal(message));
    let mut file = fs::File::create(&output_path)
        .unwrap_or_else(|err| panic!("build_scan: cannot create {}: {err}", output_path.display()));
    file.write_all(out.as_bytes())
        .unwrap_or_else(|err| panic!("build_scan: cannot write {}: {err}", output_path.display()));
}

fn collect_rust_specs(manifest_dir: &Path, dir: &Path, specs: &mut BTreeSet<String>) {
    let entries = fs::read_dir(dir)
        .unwrap_or_else(|err| panic!("build_scan: cannot read {}: {err}", dir.display()));
    for entry in entries {
        let entry = entry.unwrap_or_else(|err| {
            panic!(
                "build_scan: cannot read dir entry in {}: {err}",
                dir.display()
            )
        });
        let path = entry.path();
        if path.is_dir() {
            collect_rust_specs(manifest_dir, &path, specs);
            continue;
        }
        if !path.is_file() || path.extension().and_then(|ext| ext.to_str()) != Some("rs") {
            continue;
        }
        if path.file_name().and_then(|name| name.to_str()) == Some("mod.rs") {
            continue;
        }
        let contents = fs::read_to_string(&path)
            .unwrap_or_else(|err| panic!("build_scan: cannot read {}: {err}", path.display()));
        if let Some(spec_path) = rust_spec_path(manifest_dir, &path, &contents) {
            println!("cargo:rerun-if-changed={}", path.display());
            specs.insert(spec_path);
        }
    }
}

fn rust_spec_path(manifest_dir: &Path, path: &Path, contents: &str) -> Option<String> {
    if !is_reachable_module(manifest_dir, path) {
        return None;
    }
    let spec_pos = contents
        .find("pub const SPEC: OpSpec")
        .or_else(|| contents.find("pub static SPEC: OpSpec"))?;
    let module_path = source_module_path(manifest_dir, path)?;
    if let Some(ty) = enclosing_impl_type(&contents[..spec_pos]) {
        Some(format!("{module_path}::{ty}::SPEC"))
    } else {
        Some(format!("{module_path}::SPEC"))
    }
}

fn is_reachable_module(manifest_dir: &Path, path: &Path) -> bool {
    let src_root = manifest_dir.join("src");
    let Ok(relative) = path.strip_prefix(&src_root) else {
        return false;
    };
    let mut components = relative
        .with_extension("")
        .components()
        .filter_map(|component| component.as_os_str().to_str().map(str::to_owned))
        .collect::<Vec<_>>();
    if components.is_empty()
        || components
            .last()
            .is_some_and(|component| component == "mod")
    {
        return false;
    }
    let Some(leaf) = components.pop() else {
        return false;
    };
    if components.is_empty() {
        return true;
    }
    let Some(parent_source) = parent_module_source(&src_root, &components) else {
        return false;
    };
    let Ok(parent_contents) = fs::read_to_string(parent_source) else {
        return false;
    };
    parent_contents.contains("automod::dir!")
        || parent_contents.contains(&format!("pub mod {leaf};"))
        || parent_contents.contains(&format!("pub(crate) mod {leaf};"))
        || parent_contents.contains(&format!("mod {leaf};"))
}

fn parent_module_source(src_root: &Path, components: &[String]) -> Option<PathBuf> {
    let file_path = src_root.join(components.join("/")).with_extension("rs");
    if file_path.is_file() {
        return Some(file_path);
    }
    let mod_path = src_root.join(components.join("/")).join("mod.rs");
    mod_path.is_file().then_some(mod_path)
}

fn source_module_path(manifest_dir: &Path, path: &Path) -> Option<String> {
    let src_root = manifest_dir.join("src");
    let relative = path.strip_prefix(src_root).ok()?;
    let mut parts = Vec::new();
    for component in relative.with_extension("").components() {
        let part = component.as_os_str().to_str()?;
        if part != "mod" {
            parts.push(part.to_string());
        }
    }
    if parts.is_empty() {
        None
    } else {
        Some(format!("crate::{}", parts.join("::")))
    }
}

fn enclosing_impl_type(prefix: &str) -> Option<String> {
    let impl_pos = prefix.rfind("impl ")?;
    let after_impl = &prefix[impl_pos + "impl ".len()..];
    let ty = after_impl
        .trim_start()
        .split(['<', '{', ' ', '\n', '\t'])
        .next()?;
    if ty.is_empty() || ty.contains("::") {
        None
    } else {
        Some(ty.to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::rust_spec_path;
    use std::{env, fs, path::PathBuf};

    #[test]
    fn extracts_top_level_spec_path() {
        let tmp = tempdir_sibling("spec_top_level");
        let src = tmp.join("src/ops/hash");
        fs::create_dir_all(&src).unwrap();
        fs::write(
            tmp.join("src/ops/hash.rs"),
            "automod::dir!(pub \"src/ops/hash\");",
        )
        .unwrap();
        let path = src.join("crc32.rs");
        fs::write(&path, "pub const SPEC: OpSpec = OpSpec::intrinsic(...);").unwrap();

        let got = rust_spec_path(&tmp, &path, &fs::read_to_string(&path).unwrap());
        assert_eq!(got.as_deref(), Some("crate::ops::hash::crc32::SPEC"));

        let _ = fs::remove_dir_all(&tmp);
    }

    #[test]
    fn extracts_associated_spec_path() {
        let tmp = tempdir_sibling("spec_associated");
        let src = tmp.join("src/ops/primitive/math");
        fs::create_dir_all(&src).unwrap();
        fs::write(
            tmp.join("src/ops/primitive/math.rs"),
            "automod::dir!(pub \"src/ops/primitive/math\");",
        )
        .unwrap();
        let path = src.join("add.rs");
        fs::write(
            &path,
            "pub struct Add;\nimpl Add {\n    pub const SPEC: OpSpec = OpSpec::intrinsic(...);\n}\n",
        )
        .unwrap();

        let got = rust_spec_path(&tmp, &path, &fs::read_to_string(&path).unwrap());
        assert_eq!(
            got.as_deref(),
            Some("crate::ops::primitive::math::add::Add::SPEC")
        );

        let _ = fs::remove_dir_all(&tmp);
    }

    #[test]
    fn skips_specs_not_declared_by_parent_module() {
        let tmp = tempdir_sibling("spec_unreachable");
        let src = tmp.join("src/ops/primitive/float");
        fs::create_dir_all(&src).unwrap();
        fs::write(tmp.join("src/ops/primitive/float.rs"), "pub mod f32_abs;\n").unwrap();
        let path = src.join("f32_ceil.rs");
        fs::write(
            &path,
            "pub struct F32Ceil;\nimpl F32Ceil {\n    pub const SPEC: OpSpec = OpSpec::intrinsic(...);\n}\n",
        )
        .unwrap();

        let got = rust_spec_path(&tmp, &path, &fs::read_to_string(&path).unwrap());
        assert_eq!(got, None);

        let _ = fs::remove_dir_all(&tmp);
    }

    fn tempdir_sibling(suffix: &str) -> PathBuf {
        let mut base = env::temp_dir();
        base.push(format!("vyre-build-scan-{}-{}", std::process::id(), suffix));
        base
    }
}