ae-tree-sitter-bundle 0.1.0

//! Build script for `tree-sitter-bundle`.
//!
//! For every grammar whose Cargo feature is enabled AND whose source has been
//! vendored into `grammars/<name>/`, this script:
//!   1. compiles the grammar's `parser.c` (and `scanner.c` / `scanner.cc` if present),
//!   2. resolves its highlight/injection/locals queries (with local overrides),
//!   3. emits `$OUT_DIR/generated.rs` registering the language for the runtime API.
//!
//! Grammars whose feature is on but whose source is missing are skipped with a
//! warning, so you can vendor lazily. Run `scripts/fetch-grammars.sh` to vendor.

use std::collections::BTreeSet;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::{env, fs};

#[derive(Default)]
struct Grammar {
    name: String,
    repo: String,
    rev: String,
    /// Symbol override for the single-language case.
    symbol: Option<String>,
    extensions: Vec<String>,
    /// Explicit sub-languages for multi-grammar repos (typescript, ocaml, ...).
    languages: Vec<Lang>,
}

#[derive(Default)]
struct Lang {
    id: String,
    /// Source dir relative to the grammar root. Defaults to "src".
    src: Option<String>,
    symbol: Option<String>,
    extensions: Vec<String>,
}

/// Minimal parser for our own `grammars.toml`. We control the format (flat
/// `key = "string"` and `key = ["a", "b"]` lines under `[[grammar]]` /
/// `[[grammar.language]]` tables), so a tiny reader avoids pulling a full TOML
/// crate into build-dependencies and keeps the crate buildable on old toolchains.
fn parse_manifest(text: &str) -> Vec<Grammar> {
    fn unquote(s: &str) -> String {
        s.trim().trim_matches('"').to_string()
    }
    fn parse_array(s: &str) -> Vec<String> {
        s.trim()
            .trim_start_matches('[')
            .trim_end_matches(']')
            .split(',')
            .map(|p| p.trim().trim_matches('"').to_string())
            .filter(|p| !p.is_empty())
            .collect()
    }

    let mut grammars: Vec<Grammar> = Vec::new();
    let mut in_language = false;

    for raw in text.lines() {
        let line = raw.split('#').next().unwrap_or("").trim();
        if line.is_empty() {
            continue;
        }
        if line == "[[grammar]]" {
            grammars.push(Grammar::default());
            in_language = false;
            continue;
        }
        if line == "[[grammar.language]]" {
            if let Some(g) = grammars.last_mut() {
                g.languages.push(Lang::default());
            }
            in_language = true;
            continue;
        }
        let Some((key, value)) = line.split_once('=') else {
            continue;
        };
        let key = key.trim();
        let value = value.trim();
        let Some(g) = grammars.last_mut() else { continue };

        if in_language {
            let Some(l) = g.languages.last_mut() else { continue };
            match key {
                "id" => l.id = unquote(value),
                "src" => l.src = Some(unquote(value)),
                "symbol" => l.symbol = Some(unquote(value)),
                "extensions" => l.extensions = parse_array(value),
                _ => {}
            }
        } else {
            match key {
                "name" => g.name = unquote(value),
                "repo" => g.repo = unquote(value),
                "rev" => g.rev = unquote(value),
                "symbol" => g.symbol = Some(unquote(value)),
                "extensions" => g.extensions = parse_array(value),
                _ => {}
            }
        }
    }
    grammars
}

/// One compiled, ready-to-register language.
struct Built {
    id: String,
    symbol: String,
    extensions: Vec<String>,
}

fn default_symbol(id: &str) -> String {
    let mut s = String::from("tree_sitter_");
    for ch in id.chars() {
        s.push(if ch.is_ascii_alphanumeric() { ch } else { '_' });
    }
    s
}

fn feature_enabled(name: &str) -> bool {
    let key: String = name
        .chars()
        .map(|c| if c.is_ascii_alphanumeric() { c.to_ascii_uppercase() } else { '_' })
        .collect();
    env::var(format!("CARGO_FEATURE_{key}")).is_ok()
}

/// Whether build-time fetching is permitted. On by default; turned off for
/// offline/reproducible builds via `TS_BUNDLE_NO_FETCH=1` or when Cargo is in
/// offline mode (`cargo build --offline` / `--frozen`).
fn fetch_allowed() -> bool {
    if env::var_os("TS_BUNDLE_NO_FETCH").is_some() {
        return false;
    }
    if env::var("CARGO_NET_OFFLINE").map(|v| v == "true").unwrap_or(false) {
        return false;
    }
    true
}

/// Resolve a grammar's source directory, fetching it if necessary.
///
/// Order of preference:
///   1. a committed/vendored copy at `grammars/<name>/` (offline, reproducible);
///   2. a previously-fetched copy in the cache, keyed by revision;
///   3. a fresh `git` fetch into the cache (unless fetching is disabled).
///
/// The cache defaults to `$OUT_DIR/grammars` (persists across incremental builds,
/// re-fetched after `cargo clean`); set `TS_BUNDLE_GRAMMAR_CACHE` to a stable
/// directory to keep it across cleans.
fn ensure_grammar_source(
    name: &str,
    repo: &str,
    rev: &str,
    manifest_dir: &Path,
    out_dir: &Path,
) -> Option<PathBuf> {
    // 1. committed/vendored copy wins.
    let local = manifest_dir.join("grammars").join(name);
    if dir_has_content(&local) {
        return Some(local);
    }

    // 2. cached fetch, keyed by revision so re-pinning re-fetches.
    let cache_root = env::var_os("TS_BUNDLE_GRAMMAR_CACHE")
        .map(PathBuf::from)
        .unwrap_or_else(|| out_dir.join("grammars"));
    let short = &rev[..rev.len().min(12)];
    let dest = cache_root.join(format!("{name}-{short}"));
    if dest.join(".fetched").exists() {
        return Some(dest);
    }

    // 3. fetch.
    if !fetch_allowed() {
        return None;
    }
    if repo.is_empty() || rev.is_empty() {
        println!("cargo:warning=grammar '{name}' has no repo/rev to fetch");
        return None;
    }
    println!("cargo:warning=fetching grammar '{name}' @ {short}");
    match git_fetch(repo, rev, &dest) {
        Ok(()) => {
            let _ = fs::write(dest.join(".fetched"), rev);
            Some(dest)
        }
        Err(e) => {
            println!("cargo:warning=fetch failed for '{name}': {e}");
            let _ = fs::remove_dir_all(&dest);
            None
        }
    }
}

fn dir_has_content(p: &Path) -> bool {
    p.is_dir()
        && fs::read_dir(p)
            .map(|mut it| it.next().is_some())
            .unwrap_or(false)
}

/// Shallow-clone `repo` at `rev` into `dest`, then drop the `.git` directory.
fn git_fetch(repo: &str, rev: &str, dest: &Path) -> Result<(), String> {
    let _ = fs::remove_dir_all(dest);
    fs::create_dir_all(dest).map_err(|e| e.to_string())?;
    git(&["init", "-q"], dest)?;
    git(&["remote", "add", "origin", repo], dest)?;
    // Try fetching the exact revision shallowly; fall back to a full fetch for
    // servers that don't allow fetch-by-sha.
    if git(&["fetch", "--depth", "1", "origin", rev], dest).is_err() {
        git(&["fetch", "origin"], dest)?;
    }
    git(&["checkout", "-q", rev], dest)?;
    let _ = fs::remove_dir_all(dest.join(".git"));
    Ok(())
}

fn git(args: &[&str], dir: &Path) -> Result<(), String> {
    let status = Command::new("git")
        .args(args)
        .current_dir(dir)
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status()
        .map_err(|e| format!("`git` not available: {e}"))?;
    if status.success() {
        Ok(())
    } else {
        Err(format!("git {:?} exited with {status}", args))
    }
}

fn main() {
    let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
    let queries_override = manifest_dir.join("queries");
    let manifest_path = manifest_dir.join("grammars.toml");

    println!("cargo:rerun-if-changed=grammars.toml");
    println!("cargo:rerun-if-changed=grammars");
    println!("cargo:rerun-if-changed=queries");
    println!("cargo:rerun-if-env-changed=TS_BUNDLE_GRAMMAR_CACHE");
    println!("cargo:rerun-if-env-changed=TS_BUNDLE_NO_FETCH");

    let manifest = parse_manifest(
        &fs::read_to_string(&manifest_path).expect("grammars.toml not found"),
    );

    let mut built: Vec<Built> = Vec::new();
    let mut seen_symbols: BTreeSet<String> = BTreeSet::new();

    for grammar in &manifest {
        if !feature_enabled(&grammar.name) {
            continue;
        }
        let gdir = match ensure_grammar_source(
            &grammar.name,
            &grammar.repo,
            &grammar.rev,
            &manifest_dir,
            &out_dir,
        ) {
            Some(d) => d,
            None => {
                println!(
                    "cargo:warning=grammar '{}' enabled but unavailable \
                     (not vendored and fetch disabled/failed); skipping",
                    grammar.name
                );
                continue;
            }
        };

        // Normalize into a list of language entries.
        let langs: Vec<Lang> = if grammar.languages.is_empty() {
            vec![Lang {
                id: grammar.name.clone(),
                src: Some("src".into()),
                symbol: grammar.symbol.clone(),
                extensions: grammar.extensions.clone(),
            }]
        } else {
            grammar.languages.iter().map(|l| Lang {
                id: l.id.clone(),
                src: Some(l.src.clone().unwrap_or_else(|| "src".into())),
                symbol: l.symbol.clone(),
                extensions: if l.extensions.is_empty() {
                    grammar.extensions.clone()
                } else {
                    l.extensions.clone()
                },
            }).collect()
        };

        for lang in &langs {
            let src_rel = lang.src.clone().unwrap_or_else(|| "src".into());
            let src_dir = gdir.join(&src_rel);
            let parser_c = src_dir.join("parser.c");
            if !parser_c.exists() {
                println!(
                    "cargo:warning=language '{}': no parser.c at {}; skipping",
                    lang.id,
                    parser_c.display()
                );
                continue;
            }
            let symbol = lang.symbol.clone().unwrap_or_else(|| default_symbol(&lang.id));
            if !seen_symbols.insert(symbol.clone()) {
                println!("cargo:warning=duplicate symbol '{symbol}'; skipping '{}'", lang.id);
                continue;
            }

            compile_language(&lang.id, &src_dir);

            built.push(Built {
                id: lang.id.clone(),
                symbol,
                extensions: lang.extensions.clone(),
            });

            resolve_queries(&lang.id, &gdir, &src_dir, &queries_override, &out_dir);
        }
    }

    if built.is_empty() {
        println!(
            "cargo:warning=tree-sitter-bundle: no grammars compiled. Enable language \
             features (e.g. --features rust,python or --features full) and vendor sources."
        );
    }

    write_generated(&out_dir, &built);
}

/// Compile parser.c (+ scanner) for one language into its own static archive.
fn compile_language(id: &str, src_dir: &Path) {
    let safe: String = id.chars().map(|c| if c.is_ascii_alphanumeric() { c } else { '_' }).collect();

    // C sources: parser.c + optional scanner.c
    let mut c = cc::Build::new();
    c.include(src_dir)
        .warnings(false)
        .flag_if_supported("-w")
        .flag_if_supported("-std=c11")
        .file(src_dir.join("parser.c"));
    let scanner_c = src_dir.join("scanner.c");
    if scanner_c.exists() {
        c.file(scanner_c);
    }
    c.compile(&format!("ts_{safe}_c"));

    // C++ scanner, if any, compiled separately and linked against libstdc++.
    let scanner_cc = src_dir.join("scanner.cc");
    let scanner_cpp = src_dir.join("scanner.cpp");
    let cc_file = if scanner_cc.exists() {
        Some(scanner_cc)
    } else if scanner_cpp.exists() {
        Some(scanner_cpp)
    } else {
        None
    };
    if let Some(file) = cc_file {
        let mut cpp = cc::Build::new();
        cpp.cpp(true)
            .include(src_dir)
            .warnings(false)
            .flag_if_supported("-w")
            .flag_if_supported("-std=c++14")
            .file(file)
            .compile(&format!("ts_{safe}_cc"));
        // Ensure the C++ runtime is linked.
        println!("cargo:rustc-link-lib=stdc++");
    }
}

/// Copy the best-matching query files into OUT_DIR (or write empty placeholders).
fn resolve_queries(id: &str, gdir: &Path, src_dir: &Path, overrides: &Path, out_dir: &Path) {
    let dst_dir = out_dir.join("queries").join(id);
    fs::create_dir_all(&dst_dir).unwrap();

    for kind in ["highlights", "injections", "locals"] {
        let file = format!("{kind}.scm");
        let src_parent = src_dir.parent().unwrap_or(src_dir);
        let candidates = [
            overrides.join(id).join(&file),          // local override wins
            src_dir.join("queries").join(&file),     // queries inside src (rare)
            src_parent.join("queries").join(&file),  // queries beside src (typescript/markdown subdirs)
            gdir.join("queries").join(&file),        // repo-root queries
            gdir.join("queries").join(id).join(&file),
        ];
        let dst = dst_dir.join(&file);
        let found = candidates.iter().find(|p| p.exists());
        match found {
            Some(p) => {
                fs::copy(p, &dst).unwrap();
            }
            None => {
                // empty placeholder so include_str! always succeeds
                fs::write(&dst, "").unwrap();
            }
        }
    }
}

fn write_generated(out_dir: &Path, built: &[Built]) {
    let mut s = String::new();
    s.push_str("// @generated by build.rs - do not edit\n");
    s.push_str("use tree_sitter_language::LanguageFn;\n\n");
    s.push_str("pub(crate) struct RawLanguage {\n");
    s.push_str("    pub name: &'static str,\n");
    s.push_str("    pub language: LanguageFn,\n");
    s.push_str("    pub highlights: &'static str,\n");
    s.push_str("    pub injections: &'static str,\n");
    s.push_str("    pub locals: &'static str,\n");
    s.push_str("    pub extensions: &'static [&'static str],\n");
    s.push_str("}\n\n");

    s.push_str("extern \"C\" {\n");
    for b in built {
        s.push_str(&format!("    fn {}() -> *const ();\n", b.symbol));
    }
    s.push_str("}\n\n");

    s.push_str("pub(crate) static RAW_LANGUAGES: &[RawLanguage] = &[\n");
    for b in built {
        let exts = b
            .extensions
            .iter()
            .map(|e| format!("\"{e}\""))
            .collect::<Vec<_>>()
            .join(", ");
        s.push_str("    RawLanguage {\n");
        s.push_str(&format!("        name: \"{}\",\n", b.id));
        s.push_str(&format!(
            "        language: unsafe {{ LanguageFn::from_raw({}) }},\n",
            b.symbol
        ));
        s.push_str(&format!(
            "        highlights: include_str!(concat!(env!(\"OUT_DIR\"), \"/queries/{}/highlights.scm\")),\n",
            b.id
        ));
        s.push_str(&format!(
            "        injections: include_str!(concat!(env!(\"OUT_DIR\"), \"/queries/{}/injections.scm\")),\n",
            b.id
        ));
        s.push_str(&format!(
            "        locals: include_str!(concat!(env!(\"OUT_DIR\"), \"/queries/{}/locals.scm\")),\n",
            b.id
        ));
        s.push_str(&format!("        extensions: &[{exts}],\n"));
        s.push_str("    },\n");
    }
    s.push_str("];\n");

    fs::write(out_dir.join("generated.rs"), s).unwrap();
}