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 tree: Option<tree_sitter::Tree>,
70 tree_revision: u64,
71}
72
73impl Highlighter {
74 pub fn invalidate(&mut self) {
79 self.tree = None;
80 }
81
82 pub fn apply_edits(&mut self, edits: &[tree_sitter::InputEdit], revision: u64) {
83 if revision == self.tree_revision {
84 return;
85 }
86 if let Some(tree) = &mut self.tree {
87 for e in edits {
88 tree.edit(e);
89 }
90 }
91 self.tree_revision = revision;
94 }
95
96 pub fn for_path(path: &std::path::Path) -> Option<Self> {
97 let spec = languages::detect(path, None).or_else(|| {
98 let line = first_line(path)?;
101 languages::detect(path, Some(&line))
102 })?;
103 Self::from_spec(spec)
104 }
105
106 fn from_spec(spec: languages::LanguageSpec) -> Option<Self> {
107 let mut parser = Parser::new();
108 parser.set_language(&spec.language).ok()?;
109 let query = Query::new(&spec.language, spec.highlights).ok()?;
110 let classes = query
111 .capture_names()
112 .iter()
113 .map(|n| Class::from_capture(n))
114 .collect();
115 Some(Self {
116 parser,
117 query,
118 classes,
119 source_hash: u64::MAX, spans: Vec::new(),
121 tree: None,
122 tree_revision: 0,
123 })
124 }
125
126 pub fn highlight(
131 &mut self,
132 rope: &ropey::Rope,
133 revision: u64,
134 first_byte: usize,
135 last_byte: usize,
136 ) -> Vec<Span> {
137 let hash = revision;
138 if hash != self.source_hash {
139 let text = rope.to_string();
142 let tree = {
143 let parse_result = self.parser.parse_with_options(
148 &mut |byte: usize, _| {
149 if byte >= rope.len_bytes() {
150 return "";
151 }
152 let (chunk, start, _, _) = rope.chunk_at_byte(byte);
153 &chunk[byte - start..]
154 },
155 self.tree.as_ref(),
156 None,
157 );
158 match parse_result {
159 Some(tree) => tree,
160 None => self.parser.parse(&text, None).unwrap(),
161 }
162 };
163 self.tree = Some(tree.clone());
164 self.tree_revision = revision;
165 let mut cursor = QueryCursor::new();
166 let mut by_byte: HashMap<usize, (usize, Class)> = HashMap::new();
167 let mut matches = cursor.matches(&self.query, tree.root_node(), text.as_bytes());
168 while let Some(m) = { StreamingIterator::next(&mut matches) } {
169 for cap in m.captures {
170 let node = cap.node;
171 let class = self.classes[cap.index as usize];
172 let entry = by_byte
174 .entry(node.start_byte())
175 .or_insert((node.end_byte(), class));
176 if node.end_byte() - node.start_byte() <= entry.0 - node.start_byte() {
177 *entry = (node.end_byte(), class);
178 }
179 }
180 }
181 let mut spans: Vec<Span> = by_byte
182 .into_iter()
183 .map(|(start, (end, class))| Span { start, end, class })
184 .collect();
185 spans.sort_by_key(|s| (s.start, s.end));
186 self.spans = spans;
187 self.source_hash = hash;
188 }
189 let lo = self.spans.partition_point(|s| s.end <= first_byte);
191 let hi = self.spans.partition_point(|s| s.start < last_byte);
192 self.spans[lo..hi.max(lo)].to_vec()
193 }
194}
195
196fn first_line(path: &std::path::Path) -> Option<String> {
200 use std::io::{BufRead, BufReader, Read};
201 let mut line = String::new();
202 BufReader::new(std::fs::File::open(path).ok()?)
203 .take(256)
204 .read_line(&mut line)
205 .ok()?;
206 Some(line)
207}
208
209#[cfg(test)]
210mod tests {
211 use super::*;
212
213 fn classes_for(path: &std::path::Path, src: &str) -> Vec<Class> {
214 let mut hl = Highlighter::for_path(path).expect("language");
215 let rope = ropey::Rope::from_str(src);
216 hl.highlight(&rope, 0, 0, src.len())
217 .iter()
218 .map(|s| s.class)
219 .collect()
220 }
221
222 #[test]
223 fn rust_keywords_and_strings() {
224 let classes = classes_for(
225 std::path::Path::new("x.rs"),
226 "fn main() { let s = \"hi\"; }\n",
227 );
228 assert!(classes.contains(&Class::Keyword), "{classes:?}");
229 assert!(classes.contains(&Class::String), "{classes:?}");
230 }
231
232 #[test]
233 fn cpp_highlights_with_cxx_scanner() {
234 let classes = classes_for(std::path::Path::new("x.cpp"), "auto edge = hone(blade);\n");
237 assert!(!classes.is_empty(), "cpp grammar produced no spans");
238 assert!(classes.contains(&Class::Type), "{classes:?}"); }
240
241 #[test]
242 fn python_and_go_and_ts() {
243 assert!(
244 classes_for(std::path::Path::new("x.py"), "def f(x):\n return x\n")
245 .contains(&Class::Keyword)
246 );
247 assert!(classes_for(
248 std::path::Path::new("x.go"),
249 "package main\nfunc main() {}\n"
250 )
251 .contains(&Class::Keyword));
252 assert!(!classes_for(std::path::Path::new("x.ts"), "const x: number = 1;\n").is_empty());
253 assert!(!classes_for(std::path::Path::new("x.json"), "{\"a\": 1}\n").is_empty());
254 assert!(!classes_for(std::path::Path::new("x.sh"), "#!/bin/sh\necho hi\n").is_empty());
255 }
256
257 #[test]
258 fn fish_lua_and_sql() {
259 assert!(!classes_for(std::path::Path::new("x.fish"), "set -l name rust\n").is_empty());
262 assert!(
263 classes_for(std::path::Path::new("x.lua"), "local x = 1\n").contains(&Class::Keyword)
264 );
265 assert!(
266 classes_for(std::path::Path::new("x.sql"), "SELECT * FROM users;\n")
267 .contains(&Class::Keyword)
268 );
269 }
270
271 #[test]
272 fn shebang_script_file_resolves() {
273 let path =
276 std::env::temp_dir().join(format!("strop-syntax-shebang-{}", std::process::id()));
277 std::fs::write(&path, "#!/usr/bin/env bash\necho hi\n").unwrap();
278 let resolved = Highlighter::for_path(&path);
279 std::fs::remove_file(&path).ok();
280 let mut hl = resolved.expect("bash via shebang");
281 let rope = ropey::Rope::from_str("#!/usr/bin/env bash\necho hi\n");
282 assert!(!hl.highlight(&rope, 0, 0, rope.len_bytes()).is_empty());
283 }
284 #[test]
285 fn highlight_survives_backtracking_requests() {
286 let mut big = String::from("namespace std {\n");
290 for i in 0..400 {
291 big.push_str(&format!(
292 "template <typename T{i}> struct O{i} {{ T{i} v; O{i} f() {{ return O{i}{{}}; }} }};\n"
293 ));
294 }
295 big.push_str("}\n");
296 let mut hl = Highlighter::for_path(std::path::Path::new("x.hpp")).unwrap();
297 let rope = ropey::Rope::from_str(&big);
298 let spans = hl.highlight(&rope, 1, 0, rope.len_bytes());
299 assert!(!spans.is_empty(), "the big file highlights");
300 let edited = big.replacen("namespace", "namespace extra_long_name_here", 1);
303 let rope2 = ropey::Rope::from_str(&edited);
304 let spans2 = hl.highlight(&rope2, 2, 0, rope2.len_bytes());
305 assert!(!spans2.is_empty());
306 }
307}