diffler_core/syntax/

registry.rs

//! Maps a file path to its tree-sitter grammar, highlight configuration and
//! tags query.

use std::borrow::Cow;
use std::collections::HashMap;
use std::ops::Range;
use std::path::Path;
use std::sync::{LazyLock, OnceLock};

use tree_sitter::{Language, Parser, Query, QueryCursor, StreamingIterator, Tree};
use tree_sitter_highlight::HighlightConfiguration;

use crate::syntax::MAX_PARSE_BYTES;

/// A grammar capture like `function.method` resolves to its longest prefix
/// here, so general categories color every grammar.
pub const HIGHLIGHT_NAMES: &[&str] = &[
    "attribute",
    "boolean",
    "comment",
    "conditional",
    "constant",
    "constant.builtin",
    "constructor",
    "field",
    "function",
    "function.builtin",
    "keyword",
    "label",
    "number",
    "operator",
    "parameter",
    "property",
    "punctuation",
    "punctuation.bracket",
    "punctuation.delimiter",
    "spell",
    "storageclass",
    "string",
    "string.escape",
    "string.special",
    "tag",
    "text.emphasis",
    "text.literal",
    "text.reference",
    "text.strong",
    "text.title",
    "text.uri",
    "type",
    "type.builtin",
    "variable",
    "variable.builtin",
    "variable.parameter",
];

pub struct LangEntry {
    pub name: &'static str,
    pub language: Language,
    highlights: Cow<'static, str>,
    injections: Cow<'static, str>,
    tags_query: Option<Cow<'static, str>>,
    /// Indentation is syntax here, so a re-indent is never a reformat.
    pub layout_significant: bool,
    /// Compiling a query costs ~15ms, so we compile on first use. `None` when
    /// the query fails to compile, and the file renders plain.
    config: OnceLock<Option<HighlightConfiguration>>,
    tags: OnceLock<Option<Query>>,
}

impl LangEntry {
    pub fn config(&self) -> Option<&HighlightConfiguration> {
        self.config
            .get_or_init(|| {
                HighlightConfiguration::new(
                    self.language.clone(),
                    self.name,
                    &self.highlights,
                    &self.injections,
                    "",
                )
                .ok()
                .map(|mut config| {
                    config.configure(HIGHLIGHT_NAMES);
                    config
                })
            })
            .as_ref()
    }

    pub fn tags(&self) -> Option<&Query> {
        self.tags
            .get_or_init(|| Query::new(&self.language, self.tags_query.as_deref()?).ok())
            .as_ref()
    }
}

pub static REGISTRY: LazyLock<LanguageRegistry> = LazyLock::new(LanguageRegistry::build);

/// Suffixes after a file's real extension, as in `config.yml.example`.
const TEMPLATE_SUFFIXES: &[&str] = &[
    "example", "sample", "template", "tmpl", "tpl", "dist", "default", "local", "orig", "bak",
];

pub struct LanguageRegistry {
    entries: Vec<LangEntry>,
    by_ext: HashMap<&'static str, usize>,
    by_name: HashMap<&'static str, usize>,
    by_filename: HashMap<&'static str, usize>,
    /// We apply this by hand over the block grammar's `(inline)` nodes, since
    /// tree-sitter's injection does not drive the split markdown grammar.
    markdown_inline_query: Option<Query>,
}

impl LanguageRegistry {
    // flat per-language registration table
    #[allow(clippy::too_many_lines)]
    pub fn build() -> Self {
        let mut r = Self {
            entries: Vec::new(),
            by_ext: HashMap::new(),
            by_name: HashMap::new(),
            by_filename: HashMap::new(),
            markdown_inline_query: None,
        };

        r.add(
            "rust",
            &["rs"],
            tree_sitter_rust::LANGUAGE.into(),
            tree_sitter_rust::HIGHLIGHTS_QUERY,
            Some(tree_sitter_rust::TAGS_QUERY),
        );
        r.add(
            "python",
            &["py", "pyi"],
            tree_sitter_python::LANGUAGE.into(),
            tree_sitter_python::HIGHLIGHTS_QUERY,
            Some(tree_sitter_python::TAGS_QUERY),
        );
        r.layout_significant();
        r.add(
            "javascript",
            &["js", "jsx", "mjs", "cjs"],
            tree_sitter_javascript::LANGUAGE.into(),
            format!(
                "{}\n{}",
                tree_sitter_javascript::HIGHLIGHT_QUERY,
                tree_sitter_javascript::JSX_HIGHLIGHT_QUERY
            ),
            Some(tree_sitter_javascript::TAGS_QUERY),
        );
        r.add(
            "typescript",
            &["ts", "mts", "cts"],
            tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
            format!(
                "{}\n{}",
                tree_sitter_javascript::HIGHLIGHT_QUERY,
                tree_sitter_typescript::HIGHLIGHTS_QUERY
            ),
            None,
        );
        r.tags_over_javascript(tree_sitter_typescript::TAGS_QUERY);
        r.add(
            "tsx",
            &["tsx"],
            tree_sitter_typescript::LANGUAGE_TSX.into(),
            format!(
                "{}\n{}\n{}",
                tree_sitter_javascript::HIGHLIGHT_QUERY,
                tree_sitter_javascript::JSX_HIGHLIGHT_QUERY,
                tree_sitter_typescript::HIGHLIGHTS_QUERY
            ),
            None,
        );
        r.tags_over_javascript(tree_sitter_typescript::TAGS_QUERY);
        r.add(
            "go",
            &["go"],
            tree_sitter_go::LANGUAGE.into(),
            tree_sitter_go::HIGHLIGHTS_QUERY,
            Some(tree_sitter_go::TAGS_QUERY),
        );
        r.add(
            "c",
            &["c", "h"],
            tree_sitter_c::LANGUAGE.into(),
            tree_sitter_c::HIGHLIGHT_QUERY,
            Some(tree_sitter_c::TAGS_QUERY),
        );
        r.add(
            "cpp",
            &["cpp", "cc", "cxx", "hpp", "hh", "hxx"],
            tree_sitter_cpp::LANGUAGE.into(),
            // the C++ query extends C's; alone it matches only C++ constructs
            format!(
                "{}\n{}",
                tree_sitter_c::HIGHLIGHT_QUERY,
                tree_sitter_cpp::HIGHLIGHT_QUERY
            ),
            Some(tree_sitter_cpp::TAGS_QUERY),
        );
        r.add(
            "java",
            &["java"],
            tree_sitter_java::LANGUAGE.into(),
            tree_sitter_java::HIGHLIGHTS_QUERY,
            Some(tree_sitter_java::TAGS_QUERY),
        );
        r.add(
            "c-sharp",
            &["cs"],
            tree_sitter_c_sharp::LANGUAGE.into(),
            tree_sitter_c_sharp::HIGHLIGHTS_QUERY,
            Some(tree_sitter_c_sharp::TAGS_QUERY),
        );
        r.add(
            "ruby",
            &["rb"],
            tree_sitter_ruby::LANGUAGE.into(),
            tree_sitter_ruby::HIGHLIGHTS_QUERY,
            Some(tree_sitter_ruby::TAGS_QUERY),
        );
        r.add(
            "php",
            &["php"],
            tree_sitter_php::LANGUAGE_PHP.into(),
            tree_sitter_php::HIGHLIGHTS_QUERY,
            Some(tree_sitter_php::TAGS_QUERY),
        );
        r.add(
            "bash",
            &["sh", "bash", "zsh", "env"],
            tree_sitter_bash::LANGUAGE.into(),
            tree_sitter_bash::HIGHLIGHT_QUERY,
            None,
        );
        r.add(
            "json",
            &["json"],
            tree_sitter_json::LANGUAGE.into(),
            tree_sitter_json::HIGHLIGHTS_QUERY,
            None,
        );
        r.register(
            "html",
            &["html", "htm"],
            tree_sitter_html::LANGUAGE.into(),
            tree_sitter_html::HIGHLIGHTS_QUERY,
            tree_sitter_html::INJECTIONS_QUERY,
            None,
        );
        r.add(
            "css",
            &["css"],
            tree_sitter_css::LANGUAGE.into(),
            tree_sitter_css::HIGHLIGHTS_QUERY,
            None,
        );
        r.add(
            "yaml",
            &["yml", "yaml"],
            tree_sitter_yaml::LANGUAGE.into(),
            tree_sitter_yaml::HIGHLIGHTS_QUERY,
            None,
        );
        r.layout_significant();
        // the grammar guards its number patterns with Lua-style `%d`
        // predicates tree-sitter never matches, so we append a pattern that
        // wins over them
        r.add(
            "sql",
            &["sql"],
            tree_sitter_sequel::LANGUAGE.into(),
            format!(
                "{}\n((literal) @number (#match? @number \"^[-+]?[0-9][0-9.]*$\"))\n",
                tree_sitter_sequel::HIGHLIGHTS_QUERY
            ),
            None,
        );
        // inline emphasis, code spans and links come from markdown_inline_spans
        r.register(
            "markdown",
            &["md", "markdown"],
            tree_sitter_md::LANGUAGE.into(),
            tree_sitter_md::HIGHLIGHT_QUERY_BLOCK,
            tree_sitter_md::INJECTION_QUERY_BLOCK,
            None,
        );
        let md_inline: Language = tree_sitter_md::INLINE_LANGUAGE.into();
        r.markdown_inline_query =
            Query::new(&md_inline, tree_sitter_md::HIGHLIGHT_QUERY_INLINE).ok();
        r.register(
            "markdown_inline",
            &[],
            md_inline,
            tree_sitter_md::HIGHLIGHT_QUERY_INLINE,
            "",
            None,
        );

        r.add(
            "toml",
            &["toml"],
            tree_sitter_toml_ng::LANGUAGE.into(),
            tree_sitter_toml_ng::HIGHLIGHTS_QUERY,
            None,
        );
        // the grammar crate ships a parser only; the query is vendored
        r.add(
            "hcl",
            &["tf", "tfvars", "hcl"],
            tree_sitter_hcl::LANGUAGE.into(),
            include_str!("../../queries/hcl/highlights.scm"),
            None,
        );
        r.add(
            "dockerfile",
            &["dockerfile", "containerfile"],
            arborium_dockerfile::language().into(),
            arborium_dockerfile::HIGHLIGHTS_QUERY,
            None,
        );
        r.add(
            "make",
            &["mk"],
            tree_sitter_make::LANGUAGE.into(),
            tree_sitter_make::HIGHLIGHTS_QUERY,
            None,
        );
        r.layout_significant();
        r.add(
            "lua",
            &["lua"],
            tree_sitter_lua::LANGUAGE.into(),
            tree_sitter_lua::HIGHLIGHTS_QUERY,
            Some(tree_sitter_lua::TAGS_QUERY),
        );
        r.add(
            "nix",
            &["nix"],
            tree_sitter_nix::LANGUAGE.into(),
            tree_sitter_nix::HIGHLIGHTS_QUERY,
            None,
        );
        r.add(
            "xml",
            &["xml", "xsd", "xslt", "svg"],
            tree_sitter_xml::LANGUAGE_XML.into(),
            tree_sitter_xml::XML_HIGHLIGHT_QUERY,
            None,
        );
        r.add(
            "swift",
            &["swift"],
            tree_sitter_swift::LANGUAGE.into(),
            tree_sitter_swift::HIGHLIGHTS_QUERY,
            Some(tree_sitter_swift::TAGS_QUERY),
        );
        r.add(
            "scala",
            &["scala", "sbt"],
            tree_sitter_scala::LANGUAGE.into(),
            tree_sitter_scala::HIGHLIGHTS_QUERY,
            None,
        );
        r.layout_significant();
        r.add(
            "elixir",
            &["ex", "exs"],
            tree_sitter_elixir::LANGUAGE.into(),
            tree_sitter_elixir::HIGHLIGHTS_QUERY,
            Some(tree_sitter_elixir::TAGS_QUERY),
        );
        r.add(
            "zig",
            &["zig", "zon"],
            tree_sitter_zig::LANGUAGE.into(),
            tree_sitter_zig::HIGHLIGHTS_QUERY,
            None,
        );
        r.add(
            "haskell",
            &["hs"],
            tree_sitter_haskell::LANGUAGE.into(),
            tree_sitter_haskell::HIGHLIGHTS_QUERY,
            None,
        );
        r.layout_significant();
        r.add(
            "dart",
            &["dart"],
            tree_sitter_dart::LANGUAGE.into(),
            tree_sitter_dart::HIGHLIGHTS_QUERY,
            Some(tree_sitter_dart::TAGS_QUERY),
        );
        r.add(
            "powershell",
            &["ps1", "psm1", "psd1"],
            tree_sitter_powershell::LANGUAGE.into(),
            tree_sitter_powershell::HIGHLIGHTS_QUERY,
            None,
        );
        r.register(
            "svelte",
            &["svelte"],
            tree_sitter_svelte_ng::LANGUAGE.into(),
            // the svelte query extends html's, and its injections carry
            // `<script>` and `<style>` into the js and css grammars
            format!(
                "{}\n{}",
                tree_sitter_html::HIGHLIGHTS_QUERY,
                tree_sitter_svelte_ng::HIGHLIGHTS_QUERY
            ),
            tree_sitter_svelte_ng::INJECTIONS_QUERY,
            None,
        );
        r.name_files("make", &["makefile", "gnumakefile"]);
        r.name_files("dockerfile", &["dockerfile", "containerfile"]);
        r.name_files(
            "bash",
            &[
                ".env",
                ".envrc",
                ".bashrc",
                ".bash_profile",
                ".bash_aliases",
                ".bash_logout",
                ".profile",
                ".zshrc",
                ".zprofile",
                ".zshenv",
                ".zlogin",
            ],
        );

        r
    }

    fn name_files(&mut self, name: &'static str, filenames: &'static [&'static str]) {
        let Some(&idx) = self.by_name.get(name) else {
            return;
        };
        for file in filenames {
            self.by_filename.insert(file, idx);
        }
    }

    fn add(
        &mut self,
        name: &'static str,
        extensions: &'static [&'static str],
        language: Language,
        highlights: impl Into<Cow<'static, str>>,
        tags: Option<&'static str>,
    ) {
        self.register(name, extensions, language, highlights, "", tags);
    }

    fn register(
        &mut self,
        name: &'static str,
        extensions: &'static [&'static str],
        language: Language,
        highlights: impl Into<Cow<'static, str>>,
        injections: impl Into<Cow<'static, str>>,
        tags: Option<&'static str>,
    ) {
        let idx = self.entries.len();
        self.entries.push(LangEntry {
            name,
            language,
            highlights: highlights.into(),
            injections: injections.into(),
            tags_query: tags.map(Cow::Borrowed),
            layout_significant: false,
            config: OnceLock::new(),
            tags: OnceLock::new(),
        });
        self.by_name.insert(name, idx);
        for ext in extensions {
            self.by_ext.insert(ext, idx);
        }
    }

    /// Applies to the language registered last.
    fn layout_significant(&mut self) {
        if let Some(entry) = self.entries.last_mut() {
            entry.layout_significant = true;
        }
    }

    /// TypeScript's tags query covers only what it adds to JavaScript, so we
    /// prepend JavaScript's to the language registered last.
    fn tags_over_javascript(&mut self, tags: &'static str) {
        if let Some(entry) = self.entries.last_mut() {
            entry.tags_query = Some(Cow::Owned(format!(
                "{}\n{tags}",
                tree_sitter_javascript::TAGS_QUERY
            )));
        }
    }

    /// By exact file name, then extension.
    pub fn for_path(&self, path: &str) -> Option<&LangEntry> {
        let name = Path::new(path).file_name()?.to_str()?.to_ascii_lowercase();
        if let Some(&idx) = self.by_filename.get(name.as_str()) {
            return self.entries.get(idx);
        }
        let ext = Path::new(path).extension()?.to_str()?;
        let &idx = self.by_ext.get(ext)?;
        self.entries.get(idx)
    }

    /// [`Self::for_path`], then a naming convention, then the `#!` line.
    pub fn for_file(&self, path: &str, content: &str) -> Option<&LangEntry> {
        self.for_path(path)
            .or_else(|| {
                let name = Path::new(path).file_name()?.to_str()?;
                self.by_convention(&name.to_ascii_lowercase())
            })
            .or_else(|| self.by_name(shebang_language(content.lines().next()?)?))
    }

    /// A known name with a suffix (`dockerfile.prod`, `.env.local`), or a
    /// template suffix (`config.yml.example`). `name` is lowercase.
    fn by_convention(&self, name: &str) -> Option<&LangEntry> {
        let prefix = name
            .get(1..)
            .and_then(|rest| rest.find('.'))
            .and_then(|dot| name.get(..=dot));
        if let Some(&idx) = prefix.and_then(|prefix| self.by_filename.get(prefix)) {
            return self.entries.get(idx);
        }
        let (stem, suffix) = name.rsplit_once('.')?;
        if !TEMPLATE_SUFFIXES.contains(&suffix) {
            return None;
        }
        self.for_path(stem).or_else(|| self.by_convention(stem))
    }

    pub fn by_name(&self, name: &str) -> Option<&LangEntry> {
        let &idx = self.by_name.get(name)?;
        self.entries.get(idx)
    }

    pub fn names(&self) -> Vec<&'static str> {
        let mut names: Vec<&'static str> = self.entries.iter().map(|entry| entry.name).collect();
        names.sort_unstable();
        names
    }

    /// A markdown fence token, by grammar name then extension.
    pub fn for_token(&self, token: &str) -> Option<&LangEntry> {
        let token = token.trim().to_ascii_lowercase();
        let token = match token.as_str() {
            "c++" => "cpp",
            "c#" | "csharp" => "cs",
            "shell" => "bash",
            "golang" => "go",
            "terraform" => "hcl",
            "docker" => "dockerfile",
            "makefile" => "make",
            "pwsh" | "ps" => "powershell",
            other => other,
        };
        let &idx = self.by_name.get(token).or_else(|| self.by_ext.get(token))?;
        self.entries.get(idx)
    }

    pub fn config_for_injection(&self, lang: &str) -> Option<&HighlightConfiguration> {
        self.for_token(lang)?.config()
    }

    /// Inline markdown captures with their highlight name, narrowest first so
    /// a first-match renderer picks the most specific.
    pub fn markdown_inline_spans(&self, content: &str) -> Vec<(Range<usize>, &'static str)> {
        if content.len() > MAX_PARSE_BYTES {
            return Vec::new();
        }
        let (Some(query), Some(block), Some(inline)) = (
            self.markdown_inline_query.as_ref(),
            self.by_name
                .get("markdown")
                .and_then(|&i| self.entries.get(i)),
            self.by_name
                .get("markdown_inline")
                .and_then(|&i| self.entries.get(i)),
        ) else {
            return Vec::new();
        };

        let mut bp = Parser::new();
        if bp.set_language(&block.language).is_err() {
            return Vec::new();
        }
        let Some(block_tree) = bp.parse(content, None) else {
            return Vec::new();
        };
        let ranges = inline_node_ranges(&block_tree);
        if ranges.is_empty() {
            return Vec::new();
        }

        let mut ip = Parser::new();
        if ip.set_included_ranges(&ranges).is_err() || ip.set_language(&inline.language).is_err() {
            return Vec::new();
        }
        let Some(inline_tree) = ip.parse(content, None) else {
            return Vec::new();
        };

        let names = query.capture_names();
        let mut cursor = QueryCursor::new();
        let mut spans = Vec::new();
        let mut matches = cursor.matches(query, inline_tree.root_node(), content.as_bytes());
        while let Some(m) = matches.next() {
            for cap in m.captures {
                let cname = names.get(cap.index as usize).copied().unwrap_or("");
                if let Some(name) = recognized_highlight(cname) {
                    spans.push((cap.node.byte_range(), name));
                }
            }
        }
        spans.sort_by_key(|(r, _)| r.end - r.start);
        spans
    }
}

fn inline_node_ranges(tree: &Tree) -> Vec<tree_sitter::Range> {
    let mut ranges = Vec::new();
    let mut cursor = tree.walk();
    loop {
        let node = cursor.node();
        if node.kind() == "inline" && node.end_byte() > node.start_byte() {
            ranges.push(node.range());
        }
        if cursor.goto_first_child() {
            continue;
        }
        loop {
            if cursor.goto_next_sibling() {
                break;
            }
            if !cursor.goto_parent() {
                return ranges;
            }
        }
    }
}

/// The longest `HIGHLIGHT_NAMES` entry that is a dotted prefix of `capture`,
/// the way tree-sitter resolves it.
fn recognized_highlight(capture: &str) -> Option<&'static str> {
    HIGHLIGHT_NAMES
        .iter()
        .copied()
        .filter(|name| {
            capture == *name
                || capture
                    .strip_prefix(name)
                    .is_some_and(|r| r.starts_with('.'))
        })
        .max_by_key(|name| name.len())
}

impl Default for LanguageRegistry {
    fn default() -> Self {
        Self::build()
    }
}

fn shebang_language(line: &str) -> Option<&'static str> {
    let command = line.strip_prefix("#!")?.trim();
    let mut words = command.split_whitespace();
    let program = words.next()?;
    let program = if program.ends_with("/env") {
        words.find(|word| !word.starts_with('-'))?
    } else {
        program
    };
    let name = program.rsplit('/').next()?;
    let name = name.trim_end_matches(|c: char| c.is_ascii_digit() || c == '.');
    let lang = match name {
        "python" | "pypy" => "python",
        "bash" | "sh" | "zsh" | "dash" | "ksh" | "ash" => "bash",
        "node" | "nodejs" | "deno" | "bun" => "javascript",
        "ruby" => "ruby",
        "php" => "php",
        "lua" | "luajit" => "lua",
        "pwsh" => "powershell",
        "elixir" => "elixir",
        _ => return None,
    };
    Some(lang)
}

#[cfg(test)]
mod tests {
    use super::REGISTRY;

    fn lang(path: &str) -> Option<&'static str> {
        REGISTRY.for_file(path, "").map(|entry| entry.name)
    }

    #[test]
    fn convention_names_pick_their_grammar() {
        assert_eq!(lang("deploy/Dockerfile.prod"), Some("dockerfile"));
        assert_eq!(lang("api.dockerfile"), Some("dockerfile"));
        assert_eq!(lang(".env"), Some("bash"));
        assert_eq!(lang("app/.env.example"), Some("bash"));
        assert_eq!(lang("Makefile.local"), Some("make"));
        assert_eq!(lang(".zshrc"), Some("bash"));
    }

    #[test]
    fn a_template_suffix_falls_back_to_the_real_extension() {
        assert_eq!(lang("config.yml.example"), Some("yaml"));
        assert_eq!(lang("settings.json.dist"), Some("json"));
        assert_eq!(lang("notes.example"), None, "nothing under the suffix");
        assert_eq!(lang(".env.example"), Some("bash"));
        assert_eq!(
            REGISTRY.for_path("config.yml.example").map(|e| e.name),
            None,
            "the file kinds and Stats read the name and extension alone"
        );
    }

    #[test]
    fn a_shebang_names_the_language_of_a_script_with_no_extension() {
        let script = "#!/usr/bin/env python3\nprint(1)\n";
        assert_eq!(
            REGISTRY.for_file("bin/tool", script).map(|e| e.name),
            Some("python")
        );
        let shell = "#!/bin/sh -e\necho hi\n";
        assert_eq!(
            REGISTRY.for_file("run", shell).map(|e| e.name),
            Some("bash")
        );
        assert_eq!(
            REGISTRY.for_file("README", "plain text\n").map(|e| e.name),
            None
        );
    }

    #[test]
    fn every_interpreter_a_shebang_names_is_a_bundled_grammar() {
        let inferred = [
            "python", "sh", "node", "ruby", "php", "lua", "pwsh", "elixir",
        ]
        .into_iter()
        .filter_map(|program| super::shebang_language(&format!("#!/bin/{program}")));
        for name in inferred {
            assert!(REGISTRY.by_name(name).is_some(), "{name} is no grammar");
        }
    }

    #[test]
    fn exactly_the_layout_languages_are_layout_significant() {
        let flagged: Vec<_> = super::REGISTRY
            .entries
            .iter()
            .filter(|entry| entry.layout_significant)
            .map(|entry| entry.name)
            .collect();
        assert_eq!(flagged, ["python", "yaml", "make", "scala", "haskell"]);
    }
}