From facbbb13dcd389450956b45e980817e6f7f9adef Mon Sep 17 00:00:00 2001
From: Ada Alakbarova <ada.alakbarova@proton.me>
Date: Fri, 25 Jul 2025 20:00:46 +0200
Subject: [PATCH] feat(parse): add an option to limit tree depth

useful for getting a cursory overview of large trees
---
 mgf_dev/src/main.rs |  7 ++++--
 src/ast.rs          | 58 +++++++++++++++++++++++++++++++++++++++------
 2 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/mgf_dev/src/main.rs b/mgf_dev/src/main.rs
index 993ec29..43e8693 100644
--- a/mgf_dev/src/main.rs
+++ b/mgf_dev/src/main.rs
@@ -33,6 +33,9 @@ enum Command {
     Parse {
         /// Path to the file to parse. Its type will be guessed from its extension.
         path: PathBuf,
+        /// Limit the depth of the tree
+        #[arg(short, long)]
+        max_depth: Option<usize>,
     },
     /// Compare two files, returning exit code 0 if their trees are isomorphic, and 1 otherwise
     Compare {
@@ -75,7 +78,7 @@ fn real_main(args: &CliArgs) -> Result<i32, String> {
     };
 
     let exit_code = match &args.command {
-        Command::Parse { path } => {
+        Command::Parse { path, max_depth } => {
             let lang_profile = lang_profile(path)?;
 
             let contents = contents(path)?;
@@ -83,7 +86,7 @@ fn real_main(args: &CliArgs) -> Result<i32, String> {
             let tree = AstNode::parse(&contents, lang_profile, &arena, &ref_arena)
                 .map_err(|err| format!("File has parse errors: {err}"))?;
 
-            print!("{}", tree.ascii_tree());
+            print!("{}", tree.ascii_tree(*max_depth));
             0
         }
         Command::Compare {
diff --git a/src/ast.rs b/src/ast.rs
index 09dab7c..42c285e 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -822,17 +822,30 @@ impl<'a> AstNode<'a> {
         }
     }
 
-    /// Represents the node and its sub-structure in ASCII art
-    pub fn ascii_tree(&'a self) -> String {
-        self.internal_ascii_tree(&Color::DarkGray.prefix().to_string(), true, None)
+    /// Represents the node and its sub-structure in ASCII art, optionally printing only the nodes
+    /// up to a given depth
+    pub fn ascii_tree(&'a self, max_depth: Option<usize>) -> String {
+        self.internal_ascii_tree(
+            0,
+            max_depth,
+            &Color::DarkGray.prefix().to_string(),
+            true,
+            None,
+        )
     }
 
     fn internal_ascii_tree(
         &'a self,
+        depth: usize,
+        max_depth: Option<usize>,
         prefix: &str,
         last_child: bool,
         parent: Option<&CommutativeParent>,
     ) -> String {
+        if max_depth == Some(depth) {
+            return String::new();
+        }
+
         let num_children = self.children.len();
         let next_parent = self.commutative_parent_definition();
 
@@ -883,7 +896,13 @@ impl<'a> AstNode<'a> {
                 .filter(|(_, child)| child.grammar_name != "@virtual_line@")
                 .map(|(index, child)| {
                     let new_prefix = format!("{prefix}{} ", if last_child { " " } else { "│" });
-                    child.internal_ascii_tree(&new_prefix, index == num_children - 1, next_parent)
+                    child.internal_ascii_tree(
+                        depth + 1,
+                        max_depth,
+                        &new_prefix,
+                        index == num_children - 1,
+                        next_parent,
+                    )
                 }),
         )
         .collect()
@@ -1513,8 +1532,6 @@ mod tests {
         let ctx = ctx();
         let tree = ctx.parse("a.json", "{\"foo\": 3, \"bar\": 4}");
 
-        let ascii_tree = tree.ascii_tree();
-
         let expected = "\
 \u{1b}[90m└\u{1b}[0mdocument
 \u{1b}[90m  └\u{1b}[0mobject\u{1b}[95m Commutative\u{1b}[0m
@@ -1537,7 +1554,34 @@ mod tests {
 \u{1b}[90m    └\u{1b}[0m\u{1b}[31m}\u{1b}[0m
 ";
 
-        assert_eq!(ascii_tree, expected);
+        assert_eq!(tree.ascii_tree(None), expected);
+        assert_eq!(tree.ascii_tree(Some(5)), expected);
+
+        let expected = "\
+\u{1b}[90m└\u{1b}[0mdocument
+\u{1b}[90m  └\u{1b}[0mobject\u{1b}[95m Commutative\u{1b}[0m
+\u{1b}[90m    ├\u{1b}[0m\u{1b}[31m{\u{1b}[0m
+\u{1b}[90m    ├\u{1b}[0mpair \u{1b}[96mSignature [[\"foo\"]]\u{1b}[0m
+\u{1b}[90m    │ ├key: \u{1b}[0mstring
+\u{1b}[90m    │ ├\u{1b}[0m\u{1b}[31m:\u{1b}[0m
+\u{1b}[90m    │ └value: \u{1b}[0mnumber \u{1b}[31m3\u{1b}[0m
+\u{1b}[90m    ├\u{1b}[0m\u{1b}[31m,\u{1b}[0m
+\u{1b}[90m    ├\u{1b}[0mpair \u{1b}[96mSignature [[\"bar\"]]\u{1b}[0m
+\u{1b}[90m    │ ├key: \u{1b}[0mstring
+\u{1b}[90m    │ ├\u{1b}[0m\u{1b}[31m:\u{1b}[0m
+\u{1b}[90m    │ └value: \u{1b}[0mnumber \u{1b}[31m4\u{1b}[0m
+\u{1b}[90m    └\u{1b}[0m\u{1b}[31m}\u{1b}[0m
+";
+
+        assert_eq!(tree.ascii_tree(Some(4)), expected);
+
+        let expected = "\
+\u{1b}[90m└\u{1b}[0mdocument
+";
+
+        assert_eq!(tree.ascii_tree(Some(1)), expected);
+
+        assert_eq!(tree.ascii_tree(Some(0)), "");
     }
 
     #[test]
-- 
2.47.3