1#![forbid(unsafe_code)]
29
30use std::collections::HashMap;
31use std::sync::Arc;
32
33use hikari_core::{
34 HighlightSpan as HlSpan, Highlighter, HlClass, Language, LanguagePlugin, Selector, SpanSink,
35};
36use serde::{Deserialize, Serialize};
37use thiserror::Error;
38use tree_sitter::{InputEdit, Language as TsLanguage, Parser, Point, Tree};
39use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter as TsHighlighter};
40
41pub use hikari_token::Semantic;
45
46#[derive(Debug, Error)]
49pub enum TsError {
50 #[error("grammar not registered: {0}")]
51 Unknown(String),
52 #[error("tree-sitter: {0}")]
53 Ts(String),
54}
55
56pub type Result<T> = std::result::Result<T, TsError>;
57
58pub struct Grammar {
62 pub name: String,
63 pub language: TsLanguage,
64 pub config: HighlightConfiguration,
65 pub extensions: Vec<String>,
69}
70
71pub struct GrammarRegistry {
73 grammars: HashMap<String, Grammar>,
74 pub highlight_names: Vec<&'static str>,
77}
78
79impl GrammarRegistry {
80 pub fn builtin() -> Result<Self> {
86 let highlight_names = canonical_highlight_names();
87 let mut reg = Self {
88 grammars: HashMap::new(),
89 highlight_names,
90 };
91 reg.register(
92 "rust",
93 &tree_sitter_rust::language(),
94 tree_sitter_rust::HIGHLIGHTS_QUERY,
95 tree_sitter_rust::INJECTIONS_QUERY,
96 &["rs"],
97 )?;
98 reg.register(
99 "python",
100 &tree_sitter_python::language(),
101 tree_sitter_python::HIGHLIGHTS_QUERY,
102 "",
103 &["py", "pyi"],
104 )?;
105 reg.register(
106 "json",
107 &tree_sitter_json::language(),
108 tree_sitter_json::HIGHLIGHTS_QUERY,
109 "",
110 &["json"],
111 )?;
112 reg.register(
113 "bash",
114 &tree_sitter_bash::language(),
115 tree_sitter_bash::HIGHLIGHT_QUERY,
116 "",
117 &["sh", "bash", "zsh"],
118 )?;
119 reg.register(
120 "go",
121 &tree_sitter_go::language(),
122 tree_sitter_go::HIGHLIGHTS_QUERY,
123 "",
124 &["go"],
125 )?;
126 reg.register(
127 "c",
128 &tree_sitter_c::language(),
129 tree_sitter_c::HIGHLIGHT_QUERY,
130 "",
131 &["c", "h"],
132 )?;
133 let cpp_hl = format!(
137 "{}\n{}",
138 tree_sitter_c::HIGHLIGHT_QUERY,
139 tree_sitter_cpp::HIGHLIGHT_QUERY,
140 );
141 reg.register(
142 "cpp",
143 &tree_sitter_cpp::language(),
144 &cpp_hl,
145 "",
146 &["cpp", "cc", "cxx", "hpp", "hh"],
147 )?;
148 reg.register(
149 "css",
150 &tree_sitter_css::language(),
151 tree_sitter_css::HIGHLIGHTS_QUERY,
152 "",
153 &["css", "scss"],
154 )?;
155 reg.register(
156 "html",
157 &tree_sitter_html::language(),
158 tree_sitter_html::HIGHLIGHTS_QUERY,
159 tree_sitter_html::INJECTIONS_QUERY,
160 &["html", "htm"],
161 )?;
162 reg.register(
163 "ruby",
164 &tree_sitter_ruby::language(),
165 tree_sitter_ruby::HIGHLIGHTS_QUERY,
166 "",
167 &["rb"],
168 )?;
169 Ok(reg)
170 }
171
172 fn register(
179 &mut self,
180 name: &str,
181 language: &TsLanguage,
182 highlights: &str,
183 injections: &str,
184 extensions: &[&str],
185 ) -> Result<()> {
186 let mut cfg =
187 HighlightConfiguration::new(language.clone(), name, highlights, injections, "")
188 .map_err(|e| TsError::Ts(format!("{name}: {e}")))?;
189 cfg.configure(&self.highlight_names);
190 self.grammars.insert(
191 name.to_string(),
192 Grammar {
193 name: name.to_string(),
194 language: language.clone(),
195 config: cfg,
196 extensions: extensions.iter().map(|s| (*s).to_string()).collect(),
197 },
198 );
199 Ok(())
200 }
201
202 #[must_use]
203 pub fn get(&self, language: &str) -> Option<&Grammar> {
204 self.grammars.get(language)
205 }
206
207 #[must_use]
209 pub fn from_extension(&self, ext: &str) -> Option<&Grammar> {
210 self.grammars
211 .values()
212 .find(|g| g.extensions.iter().any(|e| e == ext))
213 }
214
215 pub fn add_extension(&mut self, language: &str, ext: impl Into<String>) -> bool {
218 if let Some(g) = self.grammars.get_mut(language) {
219 let ext = ext.into();
220 if !g.extensions.iter().any(|e| *e == ext) {
221 g.extensions.push(ext);
222 }
223 true
224 } else {
225 false
226 }
227 }
228
229 pub fn languages(&self) -> impl Iterator<Item = &str> {
231 self.grammars.keys().map(String::as_str)
232 }
233}
234
235pub struct BufferParser {
239 language: String,
240 parser: Parser,
241 tree: Option<Tree>,
242}
243
244impl BufferParser {
245 pub fn new(language: &str, registry: &GrammarRegistry) -> Result<Self> {
250 let grammar = registry
251 .get(language)
252 .ok_or_else(|| TsError::Unknown(language.to_string()))?;
253 let mut parser = Parser::new();
254 parser
255 .set_language(&grammar.language)
256 .map_err(|e| TsError::Ts(e.to_string()))?;
257 Ok(Self {
258 language: language.to_string(),
259 parser,
260 tree: None,
261 })
262 }
263
264 #[must_use]
265 pub fn language(&self) -> &str {
266 &self.language
267 }
268
269 pub fn reparse(&mut self, src: &str) -> Result<()> {
281 self.tree = self.parser.parse(src, None);
282 Ok(())
283 }
284
285 pub fn reparse_edit(
295 &mut self,
296 old_src: &str,
297 new_src: &str,
298 start_byte: usize,
299 old_end_byte: usize,
300 ) -> Result<()> {
301 if self.tree.is_some() {
302 let edit = TsEdit::from_splice(old_src, new_src, start_byte, old_end_byte);
303 if let Some(tree) = self.tree.as_mut() {
304 tree.edit(&edit.to_input_edit());
305 }
306 self.tree = self.parser.parse(new_src, self.tree.as_ref());
307 } else {
308 self.tree = self.parser.parse(new_src, None);
309 }
310 Ok(())
311 }
312
313 #[must_use]
314 pub fn tree(&self) -> Option<&Tree> {
315 self.tree.as_ref()
316 }
317}
318
319#[derive(Debug, Clone, Copy, PartialEq, Eq)]
324pub struct TsEdit {
325 pub start_byte: usize,
326 pub old_end_byte: usize,
327 pub new_end_byte: usize,
328 pub start_point: (usize, usize),
330 pub old_end_point: (usize, usize),
331 pub new_end_point: (usize, usize),
332}
333
334impl TsEdit {
335 #[must_use]
339 pub fn from_splice(old: &str, new: &str, start_byte: usize, old_end_byte: usize) -> Self {
340 let new_end_byte = new.len() - (old.len() - old_end_byte);
341 Self {
342 start_byte,
343 old_end_byte,
344 new_end_byte,
345 start_point: byte_to_point(old, start_byte),
346 old_end_point: byte_to_point(old, old_end_byte),
347 new_end_point: byte_to_point(new, new_end_byte),
348 }
349 }
350
351 fn to_input_edit(self) -> InputEdit {
352 let pt = |(row, column): (usize, usize)| Point { row, column };
353 InputEdit {
354 start_byte: self.start_byte,
355 old_end_byte: self.old_end_byte,
356 new_end_byte: self.new_end_byte,
357 start_position: pt(self.start_point),
358 old_end_position: pt(self.old_end_point),
359 new_end_position: pt(self.new_end_point),
360 }
361 }
362}
363
364#[must_use]
367fn byte_to_point(text: &str, byte: usize) -> (usize, usize) {
368 let byte = byte.min(text.len());
369 let prefix = &text[..byte];
370 let row = prefix.bytes().filter(|&b| b == b'\n').count();
371 let col = prefix.len() - prefix.rfind('\n').map_or(0, |i| i + 1);
372 (row, col)
373}
374
375#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
379pub struct HighlightSpan {
380 pub start: usize,
381 pub end: usize,
382 pub semantic: Semantic,
383}
384
385pub fn highlight(
390 src: &str,
391 grammar: &Grammar,
392 registry: &GrammarRegistry,
393) -> Result<Vec<HighlightSpan>> {
394 let mut highlighter = TsHighlighter::new();
395 let events = highlighter
396 .highlight(&grammar.config, src.as_bytes(), None, |_| None)
397 .map_err(|e| TsError::Ts(e.to_string()))?;
398
399 let mut stack: Vec<usize> = Vec::new();
400 let mut spans: Vec<HighlightSpan> = Vec::new();
401 let mut run_start: Option<(usize, usize)> = None;
402
403 for ev in events {
404 let ev = ev.map_err(|e| TsError::Ts(e.to_string()))?;
405 match ev {
406 HighlightEvent::HighlightStart(h) => stack.push(h.0),
407 HighlightEvent::HighlightEnd => {
408 stack.pop();
409 run_start = None;
410 }
411 HighlightEvent::Source { start, end } => {
412 if let Some(&top) = stack.last() {
413 let sem = highlight_index_to_semantic(top, ®istry.highlight_names);
414 match run_start {
415 Some((rs, _)) if rs == start => {}
416 _ => {
417 spans.push(HighlightSpan {
418 start,
419 end,
420 semantic: sem,
421 });
422 run_start = Some((start, end));
423 }
424 }
425 }
426 }
427 }
428 }
429
430 Ok(spans)
431}
432
433fn canonical_highlight_names() -> Vec<&'static str> {
436 vec![
437 "keyword",
438 "function",
439 "function.call",
440 "function.method",
441 "type",
442 "type.builtin",
443 "constant",
444 "constant.builtin",
445 "string",
446 "string.special",
447 "number",
448 "boolean",
449 "comment",
450 "operator",
451 "punctuation",
452 "punctuation.bracket",
453 "punctuation.delimiter",
454 "variable",
455 "variable.parameter",
456 "variable.builtin",
457 "attribute",
458 "label",
459 "tag",
460 ]
461}
462
463fn highlight_index_to_semantic(index: usize, names: &[&'static str]) -> Semantic {
486 let name = names.get(index).copied().unwrap_or("");
487 match name {
488 n if n.starts_with("keyword") => Semantic::Keyword,
489 n if n.starts_with("function") => Semantic::Accent,
491 n if n.starts_with("type") => Semantic::Accent,
492 n if n.starts_with("constant.builtin") || n == "boolean" => Semantic::Literal,
493 n if n.starts_with("constant") => Semantic::Literal,
494 n if n.starts_with("string") => Semantic::String,
495 n if n == "number" => Semantic::Number,
496 n if n.starts_with("comment") => Semantic::Comment,
497 n if n.starts_with("operator") => Semantic::Symbol,
500 n if n.starts_with("punctuation") => Semantic::Symbol,
501 n if n.starts_with("variable") => Semantic::Unchanged,
503 n if n == "attribute" => Semantic::Hint,
504 n if n == "label" => Semantic::Hint,
505 n if n == "tag" => Semantic::Keyword,
506 _ => Semantic::Unchanged,
510 }
511}
512
513#[derive(Clone)]
519pub struct TreeSitterHost {
520 registry: Arc<GrammarRegistry>,
521}
522
523impl TreeSitterHost {
524 pub fn builtin() -> Result<Self> {
529 Ok(Self {
530 registry: Arc::new(GrammarRegistry::builtin()?),
531 })
532 }
533
534 pub fn languages(&self) -> impl Iterator<Item = &str> {
536 self.registry.languages()
537 }
538
539 #[must_use]
542 pub fn plugins(&self) -> Vec<Box<dyn LanguagePlugin>> {
543 let mut out: Vec<Box<dyn LanguagePlugin>> = Vec::new();
544 for name in self.registry.languages() {
545 let lang: &'static str = Box::leak(name.to_string().into_boxed_str());
548 let selectors: Vec<Selector> = self
549 .registry
550 .get(name)
551 .map(|g| {
552 g.extensions
553 .iter()
554 .map(|e| Selector::Extension(Box::leak(e.clone().into_boxed_str())))
555 .collect()
556 })
557 .unwrap_or_default();
558 out.push(Box::new(TreeSitterPlugin {
559 language: Language(lang),
560 selectors: selectors.leak(),
561 registry: self.registry.clone(),
562 grammar: lang,
563 }));
564 }
565 out
566 }
567
568 #[must_use]
570 pub fn highlighter(&self, grammar: &'static str) -> TreeSitterHighlighter {
571 TreeSitterHighlighter {
572 registry: self.registry.clone(),
573 grammar,
574 }
575 }
576}
577
578pub struct TreeSitterPlugin {
580 language: Language,
581 selectors: &'static [Selector],
582 registry: Arc<GrammarRegistry>,
583 grammar: &'static str,
584}
585
586impl LanguagePlugin for TreeSitterPlugin {
587 fn language(&self) -> Language {
588 self.language
589 }
590 fn selectors(&self) -> &'static [Selector] {
591 self.selectors
592 }
593 fn make_highlighter(&self) -> Box<dyn Highlighter> {
594 Box::new(TreeSitterHighlighter {
595 registry: self.registry.clone(),
596 grammar: self.grammar,
597 })
598 }
599}
600
601pub struct TreeSitterHighlighter {
605 registry: Arc<GrammarRegistry>,
606 grammar: &'static str,
607}
608
609impl Highlighter for TreeSitterHighlighter {
610 fn highlight(&self, text: &str) -> Vec<HlSpan> {
611 let len = u32::try_from(text.len()).unwrap_or(u32::MAX);
612 let mut sink = SpanSink::for_document(len);
613 if let Some(grammar) = self.registry.get(self.grammar)
614 && let Ok(spans) = highlight(text, grammar, &self.registry)
615 {
616 for s in spans {
617 let class: HlClass = s.semantic.into();
619 sink.push(
620 u32::try_from(s.start).unwrap_or(u32::MAX),
621 u32::try_from(s.end).unwrap_or(u32::MAX),
622 class,
623 );
624 }
625 }
626 sink.finish()
627 }
628}
629
630#[cfg(test)]
631mod tests {
632 use super::{BufferParser, GrammarRegistry, TreeSitterHost, TsEdit, byte_to_point};
633 use hikari_core::{Highlighter, HlClass};
634
635 #[test]
636 fn builtin_host_highlights_rust_with_coverage() {
637 let host = TreeSitterHost::builtin().expect("builtin grammars");
638 let hl = host.highlighter("rust");
639 let src = "fn main() {\n let x = 42;\n}\n";
640 let spans = hl.highlight(src);
641 let mut cursor = 0u32;
643 for s in &spans {
644 assert_eq!(s.span.start, cursor, "gap/overlap at {cursor}");
645 cursor = s.span.end;
646 }
647 assert_eq!(cursor as usize, src.len(), "partition must cover the text");
648 assert!(
649 spans.iter().any(|s| s.class != HlClass::Plain),
650 "tree-sitter should classify something in real Rust",
651 );
652 }
653
654 #[test]
655 fn rust_grammar_is_registered() {
656 let host = TreeSitterHost::builtin().expect("builtin grammars");
657 assert!(host.languages().any(|l| l == "rust"));
658 assert!(!host.plugins().is_empty());
659 }
660
661 #[test]
662 fn byte_to_point_counts_rows_and_byte_columns() {
663 assert_eq!(byte_to_point("abc", 2), (0, 2));
664 assert_eq!(byte_to_point("ab\ncd", 3), (1, 0));
665 assert_eq!(byte_to_point("x\ny\nz", 4), (2, 0));
666 }
667
668 #[test]
670 fn incremental_reparse_equals_full_parse() {
671 let r = GrammarRegistry::builtin().unwrap();
672 let old = "fn main() { let x = 1; }";
673 let new = "fn main() { let x = 42; }";
674 let start = old.find('1').unwrap();
675
676 let mut inc = BufferParser::new("rust", &r).unwrap();
677 inc.reparse(old).unwrap();
678 inc.reparse_edit(old, new, start, start + 1).unwrap();
679
680 let mut full = BufferParser::new("rust", &r).unwrap();
681 full.reparse(new).unwrap();
682
683 assert_eq!(
684 inc.tree().unwrap().root_node().to_sexp(),
685 full.tree().unwrap().root_node().to_sexp(),
686 );
687 }
688
689 #[test]
690 fn ts_edit_new_end_byte_accounts_for_length_delta() {
691 let old = "x = 1;";
692 let new = "x = 42;";
693 let start = old.find('1').unwrap();
694 let e = TsEdit::from_splice(old, new, start, start + 1);
695 assert_eq!(e.new_end_byte, start + 2);
696 }
697}
698
699#[cfg(test)]
700mod capture_semantics {
701 use super::*;
702 use hikari_token::hlclass_to_semantic;
703
704 fn sem(name: &'static str) -> Semantic {
713 highlight_index_to_semantic(0, &[name])
714 }
715
716 #[test]
717 fn identifiers_fold_to_the_identifier_bucket() {
718 assert_eq!(sem("function"), Semantic::Accent);
720 assert_eq!(sem("function.method"), Semantic::Accent);
721 assert_eq!(sem("type"), Semantic::Accent);
722 assert_eq!(hlclass_to_semantic(HlClass::Function), Semantic::Accent);
723 assert_eq!(hlclass_to_semantic(HlClass::Type), Semantic::Accent);
724 }
725
726 #[test]
727 fn symbolic_tokens_fold_to_the_symbolic_bucket() {
728 assert_eq!(sem("operator"), Semantic::Symbol);
730 assert_eq!(sem("punctuation.bracket"), Semantic::Symbol);
731 assert_eq!(hlclass_to_semantic(HlClass::Operator), Semantic::Symbol);
732 assert_eq!(hlclass_to_semantic(HlClass::Punctuation), Semantic::Symbol);
733 }
734
735 #[test]
736 fn a_function_name_is_never_classified_as_punctuation() {
737 let class: HlClass = sem("function").into();
739 assert_ne!(
740 class,
741 HlClass::Punctuation,
742 "a function NAME is not a symbolic token",
743 );
744 }
745
746 #[test]
747 fn an_unknown_capture_is_ordinary_text_not_punctuation() {
748 let class: HlClass = sem("something.nobody.mapped").into();
750 assert_ne!(class, HlClass::Punctuation);
751 }
752}