use std::ops::Range;
use std::sync::OnceLock;
use syntect::parsing::{ParseState, Scope, ScopeStack, SyntaxSet};
use crate::style::Token;
pub type LineTokens = Vec<(Range<usize>, Token)>;
fn syntax_set() -> &'static SyntaxSet {
static SET: OnceLock<SyntaxSet> = OnceLock::new();
SET.get_or_init(two_face::syntax::extra_newlines)
}
fn atoms() -> &'static [(u16, Token)] {
static ATOMS: OnceLock<Vec<(u16, Token)>> = OnceLock::new();
ATOMS.get_or_init(|| {
[
("punctuation", Token::Punctuation),
("keyword", Token::Keyword),
("storage", Token::Keyword),
("entity", Token::Entity),
("variable", Token::Entity),
("support", Token::Support),
("constant", Token::Constant),
("string", Token::String),
("comment", Token::Comment),
("invalid", Token::Invalid),
]
.into_iter()
.map(|(name, t)| {
(
Scope::new(name).expect("a bare atom is a scope").atom_at(0),
t,
)
})
.collect()
})
}
fn token_of_scope(scope: Scope) -> Option<Token> {
let atoms = atoms();
(0..scope.len() as usize)
.map(|i| scope.atom_at(i))
.filter_map(|a| atoms.iter().find(|(n, _)| *n == a).map(|(_, t)| *t))
.max_by_key(|t| t.index())
}
fn token_of_stack(stack: &ScopeStack) -> Option<Token> {
stack
.as_slice()
.iter()
.rev()
.find_map(|s| token_of_scope(*s))
}
pub fn knows(lang: &str) -> bool {
syntax_set().find_syntax_by_token(lang).is_some()
}
pub fn highlight(lang: &str, lines: &[&str]) -> Option<Vec<LineTokens>> {
let set = syntax_set();
let syntax = set.find_syntax_by_token(lang)?;
let mut state = ParseState::new(syntax);
let mut stack = ScopeStack::new();
let mut out = Vec::with_capacity(lines.len());
let mut buf = String::new();
for line in lines {
buf.clear();
buf.push_str(line);
buf.push('\n');
let mut spans: LineTokens = Vec::new();
let Ok(ops) = state.parse_line(&buf, set) else {
out.push(spans);
break;
};
let mut last = 0usize;
for (pos, op) in &ops {
let pos = (*pos).min(line.len());
if pos > last {
push_span(&mut spans, last..pos, token_of_stack(&stack));
last = pos;
}
let _ = stack.apply(op);
}
if line.len() > last {
push_span(&mut spans, last..line.len(), token_of_stack(&stack));
}
out.push(spans);
}
out.resize_with(lines.len(), Vec::new);
Some(out)
}
fn push_span(spans: &mut LineTokens, range: Range<usize>, token: Option<Token>) {
let Some(token) = token else { return };
if let Some((last, t)) = spans.last_mut()
&& *t == token
&& last.end == range.start
{
last.end = range.end;
return;
}
spans.push((range, token));
}
#[cfg(test)]
mod tests {
use super::*;
fn at(spans: &[LineTokens], line: usize, at: usize) -> Option<Token> {
spans[line]
.iter()
.find(|(r, _)| r.contains(&at))
.map(|(_, t)| *t)
}
#[test]
fn a_rust_block_is_classed_the_way_a_reader_expects() {
let lines = [
"fn main() {",
" let s = \"hi\"; // greet",
" println!(\"{}\", 42);",
"}",
];
let spans = highlight("rust", &lines).expect("rust is a known language");
assert_eq!(spans.len(), lines.len());
assert_eq!(at(&spans, 0, 0), Some(Token::Keyword));
assert_eq!(at(&spans, 0, 3), Some(Token::Entity));
assert_eq!(at(&spans, 1, 4), Some(Token::Keyword));
let quote = lines[1].find('"').unwrap();
assert_eq!(at(&spans, 1, quote), Some(Token::String));
assert_eq!(at(&spans, 1, quote + 1), Some(Token::String));
let slash = lines[1].find("//").unwrap();
assert_eq!(at(&spans, 1, slash), Some(Token::Comment));
assert_eq!(at(&spans, 1, slash + 4), Some(Token::Comment));
let num = lines[2].find("42").unwrap();
assert_eq!(at(&spans, 2, num), Some(Token::Constant));
assert_eq!(at(&spans, 3, 0), Some(Token::Punctuation));
}
#[test]
fn spans_are_ascending_and_within_their_line() {
let lines = [
"const x: [u8; 3] = [1, 2, 3]; // n",
"",
"fn f() -> u8 { x[0] }",
];
let spans = highlight("rust", &lines).unwrap();
for (i, line) in spans.iter().enumerate() {
let mut end = 0;
for (r, _) in line {
assert!(r.start >= end, "line {i}: {r:?} overlaps or goes backwards");
assert!(
r.end <= lines[i].len(),
"line {i}: {r:?} runs past the line"
);
assert!(r.start < r.end, "line {i}: {r:?} is empty");
end = r.end;
}
}
assert!(spans[1].is_empty(), "an empty line has nothing to class");
}
#[test]
fn state_carries_from_line_to_line() {
let lines = ["/* a", " b */ let c = 1;"];
let spans = highlight("rust", &lines).unwrap();
assert_eq!(at(&spans, 1, 3), Some(Token::Comment));
assert_eq!(at(&spans, 1, 8), Some(Token::Keyword));
}
#[test]
fn the_org_languages_are_known() {
for lang in [
"rust",
"rs",
"zig",
"swift",
"toml",
"typescript",
"ts",
"js",
"sh",
"md",
] {
assert!(knows(lang), "{lang} should resolve to a grammar");
}
assert!(!knows(""));
assert!(!knows("no-such-language"));
assert!(highlight("no-such-language", &["x"]).is_none());
}
#[test]
fn adjacent_alike_spans_merge() {
let spans = highlight("rust", &["let s = \"a b c\";"]).unwrap();
let strings: Vec<_> = spans[0]
.iter()
.filter(|(_, t)| *t == Token::String)
.collect();
assert_eq!(strings.len(), 1, "{:?}", spans[0]);
assert_eq!(strings[0].0, 8..15);
}
}