oppi.li
visualize tree-sitter trees and queries
1use crate::config::Config;
2
3use std::{
4 collections::HashMap,
5 fmt::Write,
6 path::{Path, PathBuf},
7};
8
9use console::{style, Style, Term};
10use tree_sitter::{Node, Parser, Query, QueryCursor, Range, Tree};
11
12pub struct App {
13 config: Config,
14 language: tree_sitter::Language,
15 path: PathBuf,
16 query: Option<Query>,
17 query_path: Option<PathBuf>,
18 query_error: Option<String>,
19 src: Vec<u8>,
20 tree: Tree,
21}
22
23impl App {
24 pub fn new<'a, P: AsRef<Path>>(
25 src: &'a [u8],
26 path: P,
27 query_path: Option<P>,
28 language: tree_sitter::Language,
29 ) -> Self {
30 let path = path.as_ref().to_owned();
31
32 let mut parser = Parser::new();
33 parser.set_language(language).unwrap();
34
35 let tree = parser.parse(&src, None).unwrap();
36 let query_path = query_path.map(|q| q.as_ref().to_owned());
37 let mut query_error = None;
38 let query = query_path.as_ref().and_then(|p| {
39 let query_src = std::fs::read_to_string(&p).expect("unable to read query");
40 match Query::new(language, &query_src) {
41 Ok(q) => Some(q),
42 Err(e) => {
43 query_error = Some(e.to_string());
44 None
45 }
46 }
47 });
48
49 Self {
50 config: Default::default(),
51 path,
52 query,
53 query_path,
54 query_error,
55 src: src.to_owned(),
56 tree,
57 language,
58 }
59 }
60
61 pub fn draw(&self) {
62 let term = Term::stdout();
63 term.clear_screen().unwrap();
64 let mut done = false;
65 let mut depth = 0;
66 let mut in_capture: Option<Range> = None;
67 let mut cursor = self.tree.walk();
68
69 let capture_names = self
70 .query
71 .as_ref()
72 .map(|q| q.capture_names())
73 .unwrap_or_default();
74 let capture_map = self
75 .query
76 .as_ref()
77 .map(|query| {
78 QueryCursor::new()
79 .matches(&query, self.tree.root_node(), self.src.as_slice())
80 .flat_map(|match_| match_.captures)
81 .fold(
82 HashMap::new(),
83 |mut map: HashMap<Node, Vec<u32>>, capture| {
84 map.entry(capture.node)
85 .and_modify(|idxs| idxs.push(capture.index))
86 .or_insert_with(|| vec![capture.index]);
87 map
88 },
89 )
90 })
91 .unwrap_or_default();
92
93 while !done {
94 let node = cursor.node();
95 let mut tree_string = String::new();
96 in_capture = match in_capture {
97 Some(range)
98 if !contains(&range, &node.range()) && capture_map.contains_key(&node) =>
99 {
100 Some(node.range())
101 }
102 Some(range) if !contains(&range, &node.range()) => None,
103 None if capture_map.contains_key(&node) => Some(node.range()),
104 i => i,
105 };
106
107 write!(
108 tree_string,
109 "{}",
110 (if in_capture.is_some() {
111 Style::new().on_yellow().on_bright()
112 } else {
113 Style::new()
114 })
115 .bright()
116 .black()
117 .apply_to(
118 format!("{}{}", "|", " ".repeat(self.config.indent_level))
119 .repeat(depth as usize)
120 )
121 )
122 .unwrap();
123
124 if self.config.show_field_name {
125 if let Some(f) = cursor.field_name() {
126 write!(
127 tree_string,
128 "{} ",
129 if in_capture.is_some() {
130 Style::new().on_yellow().on_bright()
131 } else {
132 Style::new()
133 }
134 .yellow()
135 .apply_to(f)
136 )
137 .unwrap()
138 }
139 }
140
141 write!(
142 tree_string,
143 "{} ",
144 if node.is_error() {
145 Style::new().red()
146 } else if in_capture.is_some() {
147 Style::new().on_yellow().on_bright()
148 } else {
149 Style::new()
150 }
151 .apply_to(node.kind()),
152 )
153 .unwrap();
154
155 if let Some(idxs) = capture_map.get(&node) {
156 for index in idxs {
157 write!(
158 tree_string,
159 "@{} ",
160 style(&capture_names[*index as usize]).magenta()
161 )
162 .unwrap();
163 }
164 }
165
166 if self.config.show_ranges {
167 let range = node.range();
168 write!(
169 tree_string,
170 " {}",
171 style(format!("{:?}..{:?}", range.start_byte, range.end_byte,))
172 .bright()
173 .black()
174 )
175 .unwrap();
176 }
177
178 if self.config.show_src {
179 write!(
180 tree_string,
181 " {:.?}",
182 style(node.utf8_text(&self.src).unwrap()).cyan()
183 )
184 .unwrap();
185 }
186
187 term.write_line(&tree_string).unwrap();
188 term.clear_to_end_of_screen().unwrap();
189
190 if cursor.goto_first_child() {
191 depth += 1;
192 continue;
193 }
194 if cursor.goto_next_sibling() {
195 continue;
196 }
197
198 loop {
199 if !cursor.goto_parent() {
200 done = true;
201 break;
202 } else {
203 depth -= 1;
204 }
205
206 if cursor.goto_next_sibling() {
207 break;
208 }
209 }
210 }
211
212 // see https://github.com/console-rs/console/issues/36#issuecomment-624731432
213 // for the reasoning behing this hackjob
214
215 term.write_line("\n(>) increase indent").unwrap();
216 term.clear_to_end_of_screen().unwrap();
217
218 term.write_line("(<) decrease indent ").unwrap();
219 term.clear_to_end_of_screen().unwrap();
220
221 term.write_line("(n) toggle ranges").unwrap();
222 term.clear_to_end_of_screen().unwrap();
223
224 term.write_line("(s) toggle source text").unwrap();
225 term.clear_to_end_of_screen().unwrap();
226
227 term.write_line("(r) reload from disk").unwrap();
228 term.clear_to_end_of_screen().unwrap();
229
230 term.write_line("(C-c) quit").unwrap();
231
232 if let Some(err) = self.query_error.as_ref() {
233 term.write_line(&format!("{}: {err}", style("query error").red()))
234 .unwrap();
235 }
236 term.clear_to_end_of_screen().unwrap();
237 }
238
239 pub fn increase_indent(&mut self) {
240 self.config.indent_level = self.config.indent_level.saturating_add(1);
241 }
242
243 pub fn decrease_indent(&mut self) {
244 self.config.indent_level = self.config.indent_level.saturating_sub(1);
245 }
246
247 pub fn toggle_ranges(&mut self) {
248 self.config.show_ranges = !self.config.show_ranges;
249 }
250
251 pub fn toggle_source(&mut self) {
252 self.config.show_src = !self.config.show_src;
253 }
254
255 pub fn reload(&mut self) {
256 let src = std::fs::read_to_string(&self.path).unwrap();
257 let new = Self::new(
258 src.as_bytes(),
259 &self.path,
260 self.query_path.as_ref(),
261 self.language.clone(),
262 );
263 *self = Self {
264 config: self.config,
265 ..new
266 };
267 }
268}
269
270// does a encompass b
271fn contains(a: &Range, b: &Range) -> bool {
272 a.start_byte <= b.start_byte && a.end_byte >= b.end_byte
273}