Pretty print `tree-sitter-subtree` expression (#4295)

2 files changed, 111 insertions, 1 deletions

diff --git a/helix-core/src/syntax.rs b/helix-core/src/syntax.rs
index 61d382fd..b5083f55 100644
--- a/helix-core/src/syntax.rs
+++ b/helix-core/src/syntax.rs
@@ -2030,6 +2030,57 @@ impl<I: Iterator<Item = HighlightEvent>> Iterator for Merge<I> {
     }
 }
 
+pub fn pretty_print_tree<W: fmt::Write>(fmt: &mut W, node: Node) -> fmt::Result {
+    pretty_print_tree_impl(fmt, node, true, None, 0)
+}
+
+fn pretty_print_tree_impl<W: fmt::Write>(
+    fmt: &mut W,
+    node: Node,
+    is_root: bool,
+    field_name: Option<&str>,
+    depth: usize,
+) -> fmt::Result {
+    fn is_visible(node: Node) -> bool {
+        node.is_missing()
+            || (node.is_named() && node.language().node_kind_is_visible(node.kind_id()))
+    }
+
+    if is_visible(node) {
+        let indentation_columns = depth * 2;
+        write!(fmt, "{:indentation_columns$}", "")?;
+
+        if let Some(field_name) = field_name {
+            write!(fmt, "{}: ", field_name)?;
+        }
+
+        write!(fmt, "({}", node.kind())?;
+    } else if is_root {
+        write!(fmt, "(\"{}\")", node.kind())?;
+    }
+
+    for child_idx in 0..node.child_count() {
+        if let Some(child) = node.child(child_idx) {
+            if is_visible(child) {
+                fmt.write_char('\n')?;
+            }
+
+            pretty_print_tree_impl(
+                fmt,
+                child,
+                false,
+                node.field_name_for_child(child_idx as u32),
+                depth + 1,
+            )?;
+        }
+    }
+
+    if is_visible(node) {
+        write!(fmt, ")")?;
+    }
+
+    Ok(())
+}
 #[cfg(test)]
 mod test {
     use super::*;
@@ -2201,6 +2252,63 @@ mod test {
         );
     }
 
+    #[track_caller]
+    fn assert_pretty_print(source: &str, expected: &str, start: usize, end: usize) {
+        let source = Rope::from_str(source);
+
+        let loader = Loader::new(Configuration { language: vec![] });
+        let language = get_language("Rust").unwrap();
+
+        let config = HighlightConfiguration::new(language, "", "", "").unwrap();
+        let syntax = Syntax::new(&source, Arc::new(config), Arc::new(loader));
+
+        let root = syntax
+            .tree()
+            .root_node()
+            .descendant_for_byte_range(start, end)
+            .unwrap();
+
+        let mut output = String::new();
+        pretty_print_tree(&mut output, root).unwrap();
+
+        assert_eq!(expected, output);
+    }
+
+    #[test]
+    fn test_pretty_print() {
+        let source = r#"/// Hello"#;
+        assert_pretty_print(source, "(line_comment)", 0, source.len());
+
+        // A large tree should be indented with fields:
+        let source = r#"fn main() {
+            println!("Hello, World!");
+        }"#;
+        assert_pretty_print(
+            source,
+            concat!(
+                "(function_item\n",
+                "  name: (identifier)\n",
+                "  parameters: (parameters)\n",
+                "  body: (block\n",
+                "    (expression_statement\n",
+                "      (macro_invocation\n",
+                "        macro: (identifier)\n",
+                "        (token_tree\n",
+                "          (string_literal))))))",
+            ),
+            0,
+            source.len(),
+        );
+
+        // Selecting a token should print just that token:
+        let source = r#"fn main() {}"#;
+        assert_pretty_print(source, r#"("fn")"#, 0, 1);
+
+        // Error nodes are printed as errors:
+        let source = r#"}{"#;
+        assert_pretty_print(source, "(ERROR)", 0, source.len());
+    }
+
     #[test]
     fn test_load_runtime_file() {
         // Test to make sure we can load some data from the runtime directory.
diff --git a/helix-term/src/commands/typed.rs b/helix-term/src/commands/typed.rs
index 1bfc8153..13a0adcf 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -1487,7 +1487,9 @@ fn tree_sitter_subtree(
             .root_node()
             .descendant_for_byte_range(from, to)
         {
-            let contents = format!("```tsq\n{}\n```", selected_node.to_sexp());
+            let mut contents = String::from("```tsq\n");
+            helix_core::syntax::pretty_print_tree(&mut contents, selected_node)?;
+            contents.push_str("\n```");
 
             let callback = async move {
                 let call: job::Callback =