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