From 91d2eb736e18e33972d20faab389257f578cbd7b Mon Sep 17 00:00:00 2001 From: Akshay Date: Tue, 4 Oct 2022 16:25:20 +0530 Subject: init --- tree-viz/.gitignore | 1 + tree-viz/Cargo.toml | 19 ++++ tree-viz/src/app.rs | 230 +++++++++++++++++++++++++++++++++++++++++++++++++ tree-viz/src/config.rs | 24 ++++++ tree-viz/src/main.rs | 96 +++++++++++++++++++++ 5 files changed, 370 insertions(+) create mode 100644 tree-viz/.gitignore create mode 100644 tree-viz/Cargo.toml create mode 100644 tree-viz/src/app.rs create mode 100644 tree-viz/src/config.rs create mode 100644 tree-viz/src/main.rs (limited to 'tree-viz') diff --git a/tree-viz/.gitignore b/tree-viz/.gitignore new file mode 100644 index 0000000..ea8c4bf --- /dev/null +++ b/tree-viz/.gitignore @@ -0,0 +1 @@ +/target diff --git a/tree-viz/Cargo.toml b/tree-viz/Cargo.toml new file mode 100644 index 0000000..04107d3 --- /dev/null +++ b/tree-viz/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "tree-viz" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +tree-sitter-elm = "5.6.3" +tree-sitter-rust = "0.20.1" +tree-sitter = "^0.20.8" +console = "^0.15" +once_cell = "1.14.0" +tree-sitter-go = "0.19.1" +tree-sitter-typescript = "0.20.1" + +[dependencies.notify] +version = "5.0.0" +default-features = false diff --git a/tree-viz/src/app.rs b/tree-viz/src/app.rs new file mode 100644 index 0000000..68d3641 --- /dev/null +++ b/tree-viz/src/app.rs @@ -0,0 +1,230 @@ +use crate::{config::Config, LANGUAGE}; + +use std::{ + collections::HashMap, + fmt::Write, + path::{Path, PathBuf}, +}; + +use console::{style, Style, Term}; +use tree_sitter::{Node, Parser, Query, QueryCursor, Range, Tree}; + +pub struct App { + config: Config, + path: PathBuf, + query: Option, + query_path: Option, + src: Vec, + tree: Tree, +} + +impl App { + pub fn new<'a, P: AsRef>(src: &'a [u8], path: P, query_path: Option

) -> Self { + let path = path.as_ref().to_owned(); + + let mut parser = Parser::new(); + parser.set_language(*LANGUAGE).unwrap(); + + let tree = parser.parse(&src, None).unwrap(); + let query_path = query_path.map(|q| q.as_ref().to_owned()); + let query = query_path.as_ref().map(|p| { + let query_src = std::fs::read_to_string(&p).expect("unable to read query"); + Query::new(*LANGUAGE, &query_src).expect("query parse error") + }); + + Self { + config: Default::default(), + path, + query, + query_path, + src: src.to_owned(), + tree, + } + } + + pub fn draw(&self) { + let term = Term::stdout(); + term.clear_screen().unwrap(); + let mut done = false; + let mut depth = 0; + let mut in_capture: Option = None; + let mut cursor = self.tree.walk(); + + let capture_names = self + .query + .as_ref() + .map(|q| q.capture_names()) + .unwrap_or_default(); + let capture_map = self + .query + .as_ref() + .map(|query| { + QueryCursor::new() + .matches(&query, self.tree.root_node(), self.src.as_slice()) + .flat_map(|match_| match_.captures) + .fold( + HashMap::new(), + |mut map: HashMap>, capture| { + map.entry(capture.node) + .and_modify(|idxs| idxs.push(capture.index)) + .or_insert_with(|| vec![capture.index]); + map + }, + ) + }) + .unwrap_or_default(); + + while !done { + let node = cursor.node(); + let mut tree_string = String::new(); + in_capture = match in_capture { + Some(range) + if !contains(&range, &node.range()) && capture_map.contains_key(&node) => + { + Some(node.range()) + } + Some(range) if !contains(&range, &node.range()) => None, + None if capture_map.contains_key(&node) => Some(node.range()), + i => i, + }; + + write!( + tree_string, + "{}", + (if in_capture.is_some() { + Style::new().on_yellow().on_bright() + } else { + Style::new() + }) + .bright() + .black() + .apply_to( + format!("{}{}", "|", " ".repeat(self.config.indent_level)) + .repeat(depth as usize) + ) + ) + .unwrap(); + + write!( + tree_string, + "{} ", + if node.is_error() { + Style::new().red() + } else if in_capture.is_some() { + Style::new().on_yellow().on_bright() + } else { + Style::new() + } + .apply_to(node.kind()), + ) + .unwrap(); + + if let Some(idxs) = capture_map.get(&node) { + for index in idxs { + write!( + tree_string, + "@{} ", + style(capture_names[*index as usize].as_str()).magenta() + ) + .unwrap(); + } + } + + if self.config.show_ranges { + let range = node.range(); + write!( + tree_string, + " {}", + style(format!("{:?}..{:?}", range.start_byte, range.end_byte,)) + .bright() + .black() + ) + .unwrap(); + } + + if self.config.show_src { + write!( + tree_string, + " {:.?}", + style(node.utf8_text(&self.src).unwrap()).cyan() + ) + .unwrap(); + } + + term.write_line(&tree_string).unwrap(); + term.clear_to_end_of_screen().unwrap(); + + if cursor.goto_first_child() { + depth += 1; + continue; + } + if cursor.goto_next_sibling() { + continue; + } + + loop { + if !cursor.goto_parent() { + done = true; + break; + } else { + depth -= 1; + } + + if cursor.goto_next_sibling() { + break; + } + } + } + + // see https://github.com/console-rs/console/issues/36#issuecomment-624731432 + // for the reasoning behing this hackjob + + term.write_line("\n(>) increase indent").unwrap(); + term.clear_to_end_of_screen().unwrap(); + + term.write_line("(<) decrease indent ").unwrap(); + term.clear_to_end_of_screen().unwrap(); + + term.write_line("(n) toggle ranges").unwrap(); + term.clear_to_end_of_screen().unwrap(); + + term.write_line("(s) toggle source text").unwrap(); + term.clear_to_end_of_screen().unwrap(); + + term.write_line("(r) reload from disk").unwrap(); + term.clear_to_end_of_screen().unwrap(); + + term.write_line("(C-c) quit").unwrap(); + term.clear_to_end_of_screen().unwrap(); + } + + pub fn increase_indent(&mut self) { + self.config.indent_level = self.config.indent_level.saturating_add(1); + } + + pub fn decrease_indent(&mut self) { + self.config.indent_level = self.config.indent_level.saturating_sub(1); + } + + pub fn toggle_ranges(&mut self) { + self.config.show_ranges = !self.config.show_ranges; + } + + pub fn toggle_source(&mut self) { + self.config.show_src = !self.config.show_src; + } + + pub fn reload(&mut self) { + let src = std::fs::read_to_string(&self.path).unwrap(); + let new = Self::new(src.as_bytes(), &self.path, self.query_path.as_ref()); + *self = Self { + config: self.config, + ..new + }; + } +} + +// does a encompass b +fn contains(a: &Range, b: &Range) -> bool { + a.start_byte <= b.start_byte && a.end_byte >= b.end_byte +} diff --git a/tree-viz/src/config.rs b/tree-viz/src/config.rs new file mode 100644 index 0000000..6f34291 --- /dev/null +++ b/tree-viz/src/config.rs @@ -0,0 +1,24 @@ +use std::default::Default; + +#[derive(Clone, Copy)] +pub struct Config { + pub indent_level: usize, + pub show_ranges: bool, + pub show_src: bool, +} + +impl Default for Config { + fn default() -> Self { + Config::new() + } +} + +impl Config { + fn new() -> Self { + Self { + indent_level: 2, + show_ranges: true, + show_src: true, + } + } +} diff --git a/tree-viz/src/main.rs b/tree-viz/src/main.rs new file mode 100644 index 0000000..79964cb --- /dev/null +++ b/tree-viz/src/main.rs @@ -0,0 +1,96 @@ +mod app; +mod config; + +use std::{ + env, fs, + path::Path, + sync::{mpsc, Arc, RwLock}, + thread, + time::Duration, +}; + +use app::App; +use console::{Key, Term}; +use notify::{Event as WatchEvent, EventKind as WatchEventKind, RecursiveMode, Watcher}; +use once_cell::sync::Lazy; +use tree_sitter::Language; + +//pub static LANGUAGE: Lazy = Lazy::new(tree_sitter_rust::language); +pub static LANGUAGE: Lazy = Lazy::new(tree_sitter_typescript::language_typescript); + +fn main() { + let mut args = env::args(); + let _ = args.next(); + + let path = args.next().expect("no arg passed"); + let query_path = args.next(); + let src = fs::read_to_string(&path).expect("unable to read file"); + + let app = Arc::new(RwLock::new(App::new( + src.as_bytes(), + &path, + query_path.as_ref(), + ))); + + let watch_fn = |watcher_app: Arc>| { + move |ev| { + if let Ok(WatchEvent { + kind: WatchEventKind::Modify(..), + .. + }) = ev + { + if let Ok(mut locked) = watcher_app.try_write() { + locked.reload(); + locked.draw(); + }; + } + } + }; + + let mut watcher1 = notify::recommended_watcher(watch_fn(Arc::clone(&app))).unwrap(); + watcher1 + .watch(Path::new(&path), RecursiveMode::NonRecursive) + .unwrap(); + + let mut watcher2 = notify::recommended_watcher(watch_fn(Arc::clone(&app))).unwrap(); + if let Some(query_path) = query_path { + watcher2 + .watch(Path::new(&query_path), RecursiveMode::NonRecursive) + .unwrap(); + } + + let (tx, rx) = mpsc::channel(); + let tx0 = tx.clone(); + thread::spawn(move || { + let term = Term::stdout(); + loop { + if let Ok(Key::Char(ev)) = term.read_key() { + tx0.send(ev).unwrap(); + } + } + }); + + if let Ok(locked) = app.try_read() { + locked.draw(); + } + + loop { + match rx.try_recv() { + Ok(ev) => { + if let Ok(mut locked) = app.try_write() { + match ev { + '>' => locked.increase_indent(), + '<' => locked.decrease_indent(), + 'n' => locked.toggle_ranges(), + 's' => locked.toggle_source(), + 'r' => locked.reload(), + _ => (), + } + locked.draw(); + } + } + _ => (), + } + thread::sleep(Duration::from_millis(10)); + } +} -- cgit v1.2.3