5 files changed, 370 insertions, 0 deletions

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>,
+    query_path: Option<PathBuf>,
+    src: Vec<u8>,
+    tree: Tree,
+}
+
+impl App {
+    pub fn new<'a, P: AsRef<Path>>(src: &'a [u8], path: P, query_path: Option<P>) -> 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<Range> = 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<Node, Vec<u32>>, 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<Language> = Lazy::new(tree_sitter_rust::language);
+pub static LANGUAGE: Lazy<Language> = 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<RwLock<App>>| {
+        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));
+    }
+}