1use std::{
28 collections::{HashMap, HashSet},
29 fs,
30 ops::RangeBounds,
31 path::{Path, PathBuf},
32 sync::{LazyLock, Mutex},
33};
34
35use duat_core::{
36 Plugins,
37 buffer::Buffer,
38 context::{self, Handle},
39 data::Pass,
40 form::{self, Form},
41 text::{Builder, Text, txt},
42};
43use tree_sitter::{Language, Node, Query};
44
45use crate::languages::get_language;
46pub use crate::parser::Parser;
47
48mod cursor;
49mod languages;
50mod parser;
51mod tree;
52
53#[derive(Default)]
65pub struct TreeSitter;
66
67impl duat_core::Plugin for TreeSitter {
68 fn plug(self, _: &Plugins) {
69 fn copy_dir_all(src: &include_dir::Dir, dst: impl AsRef<Path>) -> std::io::Result<()> {
70 fs::create_dir_all(&dst)?;
71 for entry in src.entries() {
72 if let Some(dir) = entry.as_dir() {
73 copy_dir_all(dir, dst.as_ref().join(entry.path().file_name().unwrap()))?;
74 } else {
75 fs::write(
76 dst.as_ref().join(entry.path().file_name().unwrap()),
77 entry.as_file().unwrap().contents(),
78 )?
79 }
80 }
81 Ok(())
82 }
83
84 static QUERIES: include_dir::Dir = include_dir::include_dir!("$CARGO_MANIFEST_DIR/queries");
85
86 let Ok(plugin_dir) = duat_core::utils::plugin_dir("duat-treesitter") else {
87 context::error!("No local directory, queries aren't installed");
88 return;
89 };
90
91 let dest = plugin_dir.join("queries");
92 match dest.try_exists() {
93 Ok(false) => match copy_dir_all(&QUERIES, &dest) {
94 Ok(_) => {
95 context::info!("Installed tree-sitter queries at [buffer]{dest}");
96 }
97 Err(err) => {
98 context::info!(
99 "Failed to install tree-sitter queries at [buffer]{dest}: {err}"
100 );
101 }
102 },
103 Ok(true) => {}
104 Err(err) => {
105 context::warn!("Coudn't confirm existance of [buffer]{dest}: {err}")
106 }
107 }
108
109 form::set_many_weak!(
110 ("variable", Form::white()),
111 ("variable.builtin", Form::dark_yellow()),
112 ("constant", Form::grey()),
113 ("constant.builtin", Form::dark_yellow()),
114 ("module", Form::blue().italic()),
115 ("label", Form::green()),
116 ("string", Form::green()),
117 ("character", Form::dark_yellow()),
118 ("boolean", Form::dark_yellow()),
119 ("number", Form::dark_yellow()),
120 ("type", Form::yellow().italic()),
121 ("type.builtin", Form::yellow().reset()),
122 ("attribute", Form::green()),
123 ("property", Form::green()),
124 ("function", Form::blue().reset()),
125 ("constructor", Form::dark_yellow().reset()),
126 ("operator", Form::cyan()),
127 ("keyword", Form::magenta()),
128 ("punctuation.bracket", Form::grey()),
129 ("punctuation.delimiter", Form::grey()),
130 ("comment", Form::grey()),
131 ("comment.documentation", Form::grey().bold()),
132 ("markup.strong", Form::bold()),
133 ("markup.italic", Form::italic()),
134 ("markup.strikethrough", Form::crossed_out()),
135 ("markup.underline", Form::underlined()),
136 ("markup.heading", Form::blue().bold()),
137 ("markup.math", Form::yellow()),
138 ("markup.quote", Form::grey().italic()),
139 ("markup.link", Form::blue().underlined()),
140 ("markup.raw", Form::cyan()),
141 ("markup.list", Form::yellow()),
142 ("markup.list.checked", Form::green()),
143 ("markup.list.unchecked", Form::grey()),
144 ("diff.plus", Form::red()),
145 ("diff.delta", Form::blue()),
146 ("diff.minus", Form::green()),
147 ("node.field", "variable.member"),
148 );
149
150 parser::add_parser_hook();
151 }
152}
153
154type LangParts<'a> = (&'a str, &'a Language, Queries<'a>);
155
156#[derive(Clone, Copy)]
157struct Queries<'a> {
158 highlights: &'a Query,
159 indents: &'a Query,
160 injections: &'a Query,
161}
162
163fn lang_parts_of(lang: &str, handle: &Handle) -> Option<LangParts<'static>> {
164 static MAPS: LazyLock<Mutex<HashMap<&str, LangParts<'static>>>> = LazyLock::new(Mutex::default);
165 static FAILED_PARTS: LazyLock<Mutex<HashSet<String>>> = LazyLock::new(Mutex::default);
166
167 let mut maps = MAPS.lock().unwrap();
168
169 if let Some(lang_parts) = maps.get(lang).copied() {
170 Some(lang_parts)
171 } else if FAILED_PARTS.lock().unwrap().contains(lang) {
172 None
173 } else {
174 let language: &'static Language = Box::leak(Box::new(get_language(lang, handle)?));
175
176 let get_queries = || {
177 let highlights = query_from_path(lang, "highlights", language).ok()?;
178 let indents = query_from_path(lang, "indents", language).ok()?;
179 let injections = query_from_path(lang, "injections", language).ok()?;
180 Some(Queries { highlights, indents, injections })
181 };
182
183 let Some(queries) = get_queries() else {
184 FAILED_PARTS.lock().unwrap().insert(lang.to_string());
185 return None;
186 };
187
188 let lang: &'static str = lang.to_string().leak();
189
190 maps.insert(lang, (lang, language, queries));
191
192 Some((lang, language, queries))
193 }
194}
195
196fn query_from_path(name: &str, kind: &str, language: &Language) -> Result<&'static Query, Text> {
201 static QUERIES: LazyLock<Mutex<HashMap<PathBuf, &'static Query>>> =
202 LazyLock::new(Mutex::default);
203
204 let queries_dir = duat_core::utils::plugin_dir("duat-treesitter")?.join("queries");
205
206 let path = queries_dir.join(name).join(kind).with_extension("scm");
207
208 let mut queries = QUERIES.lock().unwrap();
209
210 Ok(if let Some(query) = queries.get(&path) {
211 query
212 } else {
213 let Ok(mut query) = fs::read_to_string(&path) else {
214 let query = Box::leak(Box::new(Query::new(language, "").unwrap()));
215 queries.insert(path, query);
216 return Ok(query);
217 };
218
219 let Some(first_line) = query.lines().map(String::from).next() else {
220 context::warn!("Query is empty");
221 let query = Box::leak(Box::new(Query::new(language, "").unwrap()));
222 queries.insert(path, query);
223 return Ok(query);
224 };
225
226 if let Some(langs) = first_line.strip_prefix("; inherits: ") {
227 for name in langs.split(',') {
228 let path = queries_dir.join(name).join(kind).with_extension("scm");
229 match fs::read_to_string(&path) {
230 Ok(inherited_query) => {
231 if inherited_query.is_empty() {
232 context::warn!("Inherited query is empty");
233 }
234
235 query = format!("{inherited_query}\n{query}");
236 }
237 Err(err) => context::error!("{err}"),
238 }
239 }
240 }
241
242 let query = Box::leak(Box::new(match Query::new(language, &query) {
243 Ok(query) => query,
244 Err(err) => return Err(txt!("{err}")),
245 }));
246
247 queries.insert(path, query);
248
249 query
250 })
251}
252
253pub trait TsHandle {
255 fn get_ts_parser<'p>(&'p self, pa: &'p mut Pass) -> Option<(&'p Parser, &'p Buffer)>;
256
257 fn ts_get_indentations(
268 &self,
269 pa: &mut Pass,
270 selections: impl RangeBounds<usize> + Clone,
271 ) -> Option<Vec<usize>>;
272}
273
274impl TsHandle for Handle {
275 fn get_ts_parser<'p>(&'p self, pa: &'p mut Pass) -> Option<(&'p Parser, &'p Buffer)> {
276 parser::sync_parse(pa, self)
277 }
278
279 fn ts_get_indentations(
280 &self,
281 pa: &mut Pass,
282 selections: impl RangeBounds<usize> + Clone,
283 ) -> Option<Vec<usize>> {
284 let range = duat_core::utils::get_range(selections, self.selections(pa).len());
285
286 let carets: Vec<usize> = self
287 .selections(pa)
288 .iter()
289 .enumerate()
290 .take(range.end)
291 .skip(range.start)
292 .map(|(_, (sel, _))| sel.caret().byte())
293 .collect();
294
295 let (parser, buffer) = parser::sync_parse(pa, self)?;
296
297 carets
298 .into_iter()
299 .map(|byte| {
300 let bytes = buffer.bytes();
301 parser.indent_on(bytes.point_at_byte(byte), bytes, buffer.opts)
302 })
303 .collect()
304 }
305}
306
307#[allow(unused)]
308fn format_root(node: Node) -> Text {
309 fn format_range(node: Node, builder: &mut Builder) {
310 let mut first = true;
311 for point in [node.start_position(), node.end_position()] {
312 builder.push(txt!(
313 "[punctuation.bracket.TreeView][[[coords.TreeView]{}\
314 [punctuation.delimiter.TreeView],[] [coords.TreeView]{}\
315 [punctuation.bracket.TreeView]]]",
316 point.row,
317 point.column
318 ));
319
320 if first {
321 first = false;
322 builder.push(txt!("[punctuation.delimiter],[] "));
323 }
324 }
325 builder.push("\n");
326 }
327
328 fn format_node(
329 node: Node,
330 depth: usize,
331 pars: usize,
332 builder: &mut Builder,
333 name: Option<&str>,
334 ) {
335 builder.push(" ".repeat(depth));
336
337 if let Some(name) = name {
338 builder.push(txt!("[node.field]{name}[punctuation.delimiter.TreeView]: "));
339 }
340
341 builder.push(txt!("[punctuation.bracket.TreeView]("));
342 builder.push(txt!("[node.name]{}", node.grammar_name()));
343
344 let mut cursor = node.walk();
345 let named_children = node.named_children(&mut cursor);
346 let len = named_children.len();
347
348 if len == 0 {
349 builder.push(txt!(
350 "[punctuation.bracket.TreeView]{}[] ",
351 ")".repeat(pars)
352 ));
353 format_range(node, builder);
354 } else {
355 builder.push(" ");
356 format_range(node, builder);
357
358 let mut i = 0;
359
360 for (i, child) in named_children.enumerate() {
361 let name = node.field_name_for_named_child(i as u32);
362 let pars = if i == len - 1 { pars + 1 } else { 1 };
363 format_node(child, depth + 1, pars, builder, name);
364 }
365 }
366 }
367
368 let mut cursor = node.walk();
369 let mut builder = Text::builder();
370
371 format_node(node, 0, 1, &mut builder, None);
372
373 builder.build()
374}