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 mut chunk_iter = rope.chunks();
138 let mut offset = 0usize;
139 self.parser
140 .parse_with_options(
141 &mut |byte: usize, _| {
142 for chunk in chunk_iter.by_ref() {
143 if byte < offset + chunk.len() {
144 return &chunk[byte - offset..];
145 }
146 offset += chunk.len();
147 }
148 ""
149 },
150 self.tree.as_ref(),
151 None,
152 )
153 .unwrap_or_else(|| self.parser.parse(&text, None).unwrap())
154 };
155 self.tree = Some(tree.clone());
156 self.tree_revision = revision;
157 let mut cursor = QueryCursor::new();
158 let mut by_byte: HashMap<usize, (usize, Class)> = HashMap::new();
159 let mut matches = cursor.matches(&self.query, tree.root_node(), text.as_bytes());
160 while let Some(m) = { StreamingIterator::next(&mut matches) } {
161 for cap in m.captures {
162 let node = cap.node;
163 let class = self.classes[cap.index as usize];
164 let entry = by_byte
166 .entry(node.start_byte())
167 .or_insert((node.end_byte(), class));
168 if node.end_byte() - node.start_byte() <= entry.0 - node.start_byte() {
169 *entry = (node.end_byte(), class);
170 }
171 }
172 }
173 let mut spans: Vec<Span> = by_byte
174 .into_iter()
175 .map(|(start, (end, class))| Span { start, end, class })
176 .collect();
177 spans.sort_by_key(|s| (s.start, s.end));
178 self.spans = spans;
179 self.source_hash = hash;
180 }
181 let lo = self.spans.partition_point(|s| s.end <= first_byte);
183 let hi = self.spans.partition_point(|s| s.start < last_byte);
184 self.spans[lo..hi.max(lo)].to_vec()
185 }
186}
187
188fn first_line(path: &str) -> Option<String> {
192 use std::io::{BufRead, BufReader, Read};
193 let mut line = String::new();
194 BufReader::new(std::fs::File::open(path).ok()?)
195 .take(256)
196 .read_line(&mut line)
197 .ok()?;
198 Some(line)
199}
200
201#[cfg(test)]
202mod tests {
203 use super::*;
204
205 fn classes_for(path: &str, src: &str) -> Vec<Class> {
206 let mut hl = Highlighter::for_path(path).expect("language");
207 let rope = ropey::Rope::from_str(src);
208 hl.highlight(&rope, 0, 0, src.len())
209 .iter()
210 .map(|s| s.class)
211 .collect()
212 }
213
214 #[test]
215 fn rust_keywords_and_strings() {
216 let classes = classes_for("x.rs", "fn main() { let s = \"hi\"; }\n");
217 assert!(classes.contains(&Class::Keyword), "{classes:?}");
218 assert!(classes.contains(&Class::String), "{classes:?}");
219 }
220
221 #[test]
222 fn cpp_highlights_with_cxx_scanner() {
223 let classes = classes_for("x.cpp", "auto edge = hone(blade);\n");
226 assert!(!classes.is_empty(), "cpp grammar produced no spans");
227 assert!(classes.contains(&Class::Type), "{classes:?}"); }
229
230 #[test]
231 fn python_and_go_and_ts() {
232 assert!(classes_for("x.py", "def f(x):\n return x\n").contains(&Class::Keyword));
233 assert!(classes_for("x.go", "package main\nfunc main() {}\n").contains(&Class::Keyword));
234 assert!(!classes_for("x.ts", "const x: number = 1;\n").is_empty());
235 assert!(!classes_for("x.json", "{\"a\": 1}\n").is_empty());
236 assert!(!classes_for("x.sh", "#!/bin/sh\necho hi\n").is_empty());
237 }
238
239 #[test]
240 fn fish_lua_and_sql() {
241 assert!(!classes_for("x.fish", "set -l name rust\n").is_empty());
244 assert!(classes_for("x.lua", "local x = 1\n").contains(&Class::Keyword));
245 assert!(classes_for("x.sql", "SELECT * FROM users;\n").contains(&Class::Keyword));
246 }
247
248 #[test]
249 fn shebang_script_file_resolves() {
250 let path =
253 std::env::temp_dir().join(format!("strop-syntax-shebang-{}", std::process::id()));
254 std::fs::write(&path, "#!/usr/bin/env bash\necho hi\n").unwrap();
255 let resolved = Highlighter::for_path(path.to_str().unwrap());
256 std::fs::remove_file(&path).ok();
257 let mut hl = resolved.expect("bash via shebang");
258 let rope = ropey::Rope::from_str("#!/usr/bin/env bash\necho hi\n");
259 assert!(!hl.highlight(&rope, 0, 0, rope.len_bytes()).is_empty());
260 }
261}