use std::{
collections::HashMap,
sync::OnceLock,
};
use tree_sitter_highlight::{
HighlightConfiguration,
Highlighter,
HtmlRenderer,
};
const CAPTURES: &[(&str, &str)] = &[
("attribute", "hl-attr"),
("comment", "hl-comment"),
("constant", "hl-constant"),
("constant.builtin", "hl-constant"),
("constructor", "hl-type"),
("escape", "hl-escape"),
("function", "hl-function"),
("function.builtin", "hl-function"),
("function.method", "hl-function"),
("keyword", "hl-keyword"),
("label", "hl-label"),
("module", "hl-type"),
("number", "hl-number"),
("operator", "hl-operator"),
("property", "hl-property"),
("punctuation", "hl-punctuation"),
("punctuation.bracket", "hl-punctuation"),
("punctuation.delimiter", "hl-punctuation"),
("string", "hl-string"),
("string.special", "hl-string"),
("tag", "hl-tag"),
("type", "hl-type"),
("type.builtin", "hl-type"),
("variable", "hl-variable"),
("variable.builtin", "hl-keyword"),
("variable.parameter", "hl-parameter"),
];
pub(crate) fn highlight(language: &str, source: &str) -> Option<String> {
let configuration = grammar(language)?;
let mut highlighter = Highlighter::new();
let events = highlighter
.highlight(configuration, source.as_bytes(), None, None, |name| {
grammar(name)
})
.ok()?;
let mut renderer = HtmlRenderer::new();
renderer
.render(events, source.as_bytes(), &|highlight, attributes| {
let class = CAPTURES
.get(highlight.0)
.map_or("hl-unknown", |(_, class)| *class);
attributes.extend_from_slice(format!("class=\"{class}\"").as_bytes());
})
.ok()?;
let mut html: String = renderer.lines().collect();
if html.ends_with('\n') {
html.pop();
}
Some(html)
}
fn canonical_name(language: &str) -> String {
let folded = language.to_lowercase();
match folded.as_str() {
"sh" | "shell" | "zsh" | "console" | "terminal" => "bash",
"js" | "mjs" | "cjs" | "node" => "javascript",
"ts" => "typescript",
"py" | "python3" => "python",
"rs" => "rust",
"golang" => "go",
"h" => "c",
"yml" => "yaml",
"htm" | "xhtml" => "html",
"jsonc" => "json",
_ => return folded,
}
.to_string()
}
type Grammar = (
&'static str,
fn() -> tree_sitter::Language,
&'static str,
&'static str,
&'static str,
);
fn table() -> [Grammar; 14] {
[
(
"rust",
|| tree_sitter_rust::LANGUAGE.into(),
tree_sitter_rust::HIGHLIGHTS_QUERY,
tree_sitter_rust::INJECTIONS_QUERY,
"",
),
(
"python",
|| tree_sitter_python::LANGUAGE.into(),
tree_sitter_python::HIGHLIGHTS_QUERY,
"",
"",
),
(
"javascript",
|| tree_sitter_javascript::LANGUAGE.into(),
tree_sitter_javascript::HIGHLIGHT_QUERY,
tree_sitter_javascript::INJECTIONS_QUERY,
tree_sitter_javascript::LOCALS_QUERY,
),
(
"typescript",
|| tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
tree_sitter_typescript::HIGHLIGHTS_QUERY,
"",
tree_sitter_typescript::LOCALS_QUERY,
),
(
"tsx",
|| tree_sitter_typescript::LANGUAGE_TSX.into(),
tree_sitter_typescript::HIGHLIGHTS_QUERY,
"",
tree_sitter_typescript::LOCALS_QUERY,
),
(
"go",
|| tree_sitter_go::LANGUAGE.into(),
tree_sitter_go::HIGHLIGHTS_QUERY,
"",
"",
),
(
"c",
|| tree_sitter_c::LANGUAGE.into(),
tree_sitter_c::HIGHLIGHT_QUERY,
"",
"",
),
(
"bash",
|| tree_sitter_bash::LANGUAGE.into(),
tree_sitter_bash::HIGHLIGHT_QUERY,
"",
"",
),
(
"java",
|| tree_sitter_java::LANGUAGE.into(),
tree_sitter_java::HIGHLIGHTS_QUERY,
"",
"",
),
(
"json",
|| tree_sitter_json::LANGUAGE.into(),
tree_sitter_json::HIGHLIGHTS_QUERY,
"",
"",
),
(
"toml",
|| tree_sitter_toml_ng::LANGUAGE.into(),
tree_sitter_toml_ng::HIGHLIGHTS_QUERY,
"",
"",
),
(
"yaml",
|| tree_sitter_yaml::LANGUAGE.into(),
tree_sitter_yaml::HIGHLIGHTS_QUERY,
"",
"",
),
(
"html",
|| tree_sitter_html::LANGUAGE.into(),
tree_sitter_html::HIGHLIGHTS_QUERY,
tree_sitter_html::INJECTIONS_QUERY,
"",
),
(
"css",
|| tree_sitter_css::LANGUAGE.into(),
tree_sitter_css::HIGHLIGHTS_QUERY,
"",
"",
),
]
}
struct Lazy {
entry: Grammar,
compiled: OnceLock<Option<HighlightConfiguration>>,
}
impl Lazy {
fn get(&self) -> Option<&HighlightConfiguration> {
self.compiled
.get_or_init(|| {
let (name, language, highlights, injections, locals) = self.entry;
let mut configuration =
HighlightConfiguration::new(language(), name, highlights, injections, locals)
.ok()?;
let names: Vec<&str> = CAPTURES.iter().map(|(name, _)| *name).collect();
configuration.configure(&names);
Some(configuration)
})
.as_ref()
}
}
fn grammars() -> &'static HashMap<&'static str, Lazy> {
static GRAMMARS: OnceLock<HashMap<&'static str, Lazy>> = OnceLock::new();
GRAMMARS.get_or_init(|| {
table()
.into_iter()
.map(|entry| {
(
entry.0,
Lazy {
entry,
compiled: OnceLock::new(),
},
)
})
.collect()
})
}
fn grammar(language: &str) -> Option<&'static HighlightConfiguration> {
grammars().get(canonical_name(language).as_str())?.get()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn every_grammar_compiles_its_queries() {
for name in [
"rust",
"python",
"javascript",
"typescript",
"tsx",
"go",
"c",
"bash",
"java",
"json",
"toml",
"yaml",
"html",
"css",
] {
assert!(
highlight(name, "x").is_some(),
"grammar `{name}` did not load"
);
}
}
#[test]
fn highlights_rust_keywords_and_strings() {
let html = highlight("rust", "fn main() { let x = \"hi\"; }").unwrap_or_default();
assert!(html.contains("hl-keyword"), "no keyword span: {html}");
assert!(html.contains("hl-string"), "no string span: {html}");
assert!(html.contains("main"));
}
#[test]
fn does_not_trail_a_newline() {
let html = highlight("rust", "fn main() {}").unwrap_or_default();
assert!(
!html.ends_with('\n'),
"a trailing newline shows as a blank line in the block: {html:?}"
);
}
#[test]
fn resolves_the_aliases_people_write() {
for (alias, code) in [("rs", "fn f() {}"), ("sh", "echo hi"), ("py", "x = 1")] {
assert!(
highlight(alias, code).is_some(),
"`{alias}` was not recognized"
);
}
}
#[test]
fn escapes_the_source_it_highlights() {
let html = highlight("rust", "let a = b < c && d > e;").unwrap_or_default();
assert!(html.contains("<"), "unescaped `<`: {html}");
assert!(html.contains(">"), "unescaped `>`: {html}");
assert!(html.contains("&"), "unescaped `&`: {html}");
}
#[test]
fn an_unknown_language_is_left_alone() {
assert!(highlight("brainfuck", "+++.").is_none());
}
#[test]
fn every_capture_has_a_class() {
assert!(
CAPTURES
.iter()
.all(|(name, class)| { !name.is_empty() && class.starts_with("hl-") })
);
}
}