1use std::collections::HashMap;
6
7use streaming_iterator::StreamingIterator;
8pub mod languages;
9
10use tree_sitter::{Parser, Query, QueryCursor};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
15pub enum Class {
16 Keyword,
17 Function,
18 Type,
19 String,
20 Comment,
21 Number,
22 Operator,
23 Punctuation,
24 Constant,
25 Variable,
26 Attribute,
27}
28
29impl Class {
30 fn from_capture(name: &str) -> Self {
31 let head = name.split('.').next().unwrap_or(name);
32 match head {
33 "keyword" => Class::Keyword,
34 "function" | "constructor" => Class::Function,
35 "type" => Class::Type,
36 "string" | "character" => Class::String,
37 "comment" => Class::Comment,
38 "number" | "float" => Class::Number,
39 "operator" => Class::Operator,
40 "punctuation" => Class::Punctuation,
41 "constant" | "boolean" => Class::Constant,
42 "attribute" | "property" => Class::Attribute,
43 _ => Class::Variable,
44 }
45 }
46}
47
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub struct Span {
51 pub start: usize,
52 pub end: usize,
53 pub class: Class,
54}
55
56pub struct Highlighter {
61 parser: Parser,
62 query: Query,
63 classes: Vec<Class>,
65 source_hash: u64,
66 spans: Vec<Span>,
67}
68
69impl Highlighter {
70 pub fn for_path(path: &str) -> Option<Self> {
71 let ext = std::path::Path::new(path)
72 .extension()
73 .map(|e| format!(".{}", e.to_string_lossy()))?;
74 let spec = languages::for_extension(&ext)?;
75 let mut parser = Parser::new();
76 parser.set_language(&spec.language).ok()?;
77 let query = Query::new(&spec.language, spec.highlights).ok()?;
78 let classes = query
79 .capture_names()
80 .iter()
81 .map(|n| Class::from_capture(n))
82 .collect();
83 Some(Self {
84 parser,
85 query,
86 classes,
87 source_hash: 0,
88 spans: Vec::new(),
89 })
90 }
91
92 pub fn highlight(
96 &mut self,
97 rope: &ropey::Rope,
98 first_byte: usize,
99 last_byte: usize,
100 ) -> Vec<Span> {
101 let mut hasher = std::hash::DefaultHasher::new();
102 std::hash::Hash::hash(&rope.len_bytes(), &mut hasher);
103 if let (Some(first), Some(last)) = (
106 rope.get_byte(0),
107 rope.len_bytes()
108 .checked_sub(1)
109 .and_then(|i| rope.get_byte(i)),
110 ) {
111 std::hash::Hash::hash(&(first, last), &mut hasher);
112 }
113 let hash = std::hash::Hasher::finish(&hasher);
114 if hash != self.source_hash {
115 let text = rope.to_string(); let Some(tree) = self.parser.parse(&text, None) else {
117 return Vec::new();
118 };
119 let mut cursor = QueryCursor::new();
120 let mut by_byte: HashMap<usize, (usize, Class)> = HashMap::new();
121 let mut matches = cursor.matches(&self.query, tree.root_node(), text.as_bytes());
122 while let Some(m) = { StreamingIterator::next(&mut matches) } {
123 for cap in m.captures {
124 let node = cap.node;
125 let class = self.classes[cap.index as usize];
126 let entry = by_byte
128 .entry(node.start_byte())
129 .or_insert((node.end_byte(), class));
130 if node.end_byte() - node.start_byte() <= entry.0 - node.start_byte() {
131 *entry = (node.end_byte(), class);
132 }
133 }
134 }
135 let mut spans: Vec<Span> = by_byte
136 .into_iter()
137 .map(|(start, (end, class))| Span { start, end, class })
138 .collect();
139 spans.sort_by_key(|s| (s.start, s.end));
140 self.spans = spans;
141 self.source_hash = hash;
142 }
143 let lo = self.spans.partition_point(|s| s.end <= first_byte);
145 let hi = self.spans.partition_point(|s| s.start < last_byte);
146 self.spans[lo..hi.max(lo)].to_vec()
147 }
148}
149
150#[cfg(test)]
151mod tests {
152 use super::*;
153
154 fn classes_for(path: &str, src: &str) -> Vec<Class> {
155 let mut hl = Highlighter::for_path(path).expect("language");
156 let rope = ropey::Rope::from_str(src);
157 hl.highlight(&rope, 0, src.len())
158 .iter()
159 .map(|s| s.class)
160 .collect()
161 }
162
163 #[test]
164 fn rust_keywords_and_strings() {
165 let classes = classes_for("x.rs", "fn main() { let s = \"hi\"; }\n");
166 assert!(classes.contains(&Class::Keyword), "{classes:?}");
167 assert!(classes.contains(&Class::String), "{classes:?}");
168 }
169
170 #[test]
171 fn cpp_highlights_with_cxx_scanner() {
172 let classes = classes_for("x.cpp", "auto edge = hone(blade);\n");
177 assert!(!classes.is_empty(), "cpp grammar produced no spans");
178 assert!(classes.contains(&Class::Type), "{classes:?}"); }
180
181 #[test]
182 fn python_and_go_and_ts() {
183 assert!(classes_for("x.py", "def f(x):\n return x\n").contains(&Class::Keyword));
184 assert!(classes_for("x.go", "package main\nfunc main() {}\n").contains(&Class::Keyword));
185 assert!(!classes_for("x.ts", "const x: number = 1;\n").is_empty());
186 assert!(!classes_for("x.json", "{\"a\": 1}\n").is_empty());
187 assert!(!classes_for("x.sh", "#!/bin/sh\necho hi\n").is_empty());
188 }
189}