From 7cbc365da8b904dc3e2c39a1fb75b17b98debd11 Mon Sep 17 00:00:00 2001
From: Weihang Lo <me@weihanglo.tw>
Date: Mon, 17 May 2021 00:49:12 +0800
Subject: [PATCH] feat: new `--depth` option for `cargo tree`

---
 src/bin/cargo/commands/tree.rs |  2 ++
 src/cargo/ops/tree/mod.rs      | 37 +++++++++++++++++++++-------------
 2 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/src/bin/cargo/commands/tree.rs b/src/bin/cargo/commands/tree.rs
index e6a1f21b2..ea244ec18 100644
--- a/src/bin/cargo/commands/tree.rs
+++ b/src/bin/cargo/commands/tree.rs
@@ -55,6 +55,7 @@ pub fn cli() -> App {
             )
             .short("i"),
         )
+        .arg(opt("depth", "Maximum display depth of the dependency tree").value_name("DEPTH"))
         // Deprecated, use --prefix=none instead.
         .arg(Arg::with_name("no-indent").long("no-indent").hidden(true))
         // Deprecated, use --prefix=depth instead.
@@ -201,6 +202,7 @@ subtree of the package given to -p.\n\
         charset,
         format: args.value_of("format").unwrap().to_string(),
         graph_features,
+        max_display_depth: args.value_of_u32("depth")?.unwrap_or(u32::MAX),
     };
 
     tree::build_and_print(&ws, &opts)?;
diff --git a/src/cargo/ops/tree/mod.rs b/src/cargo/ops/tree/mod.rs
index 3f7d65174..7b7ac9cda 100644
--- a/src/cargo/ops/tree/mod.rs
+++ b/src/cargo/ops/tree/mod.rs
@@ -43,6 +43,8 @@ pub struct TreeOptions {
     pub format: String,
     /// Includes features in the tree as separate nodes.
     pub graph_features: bool,
+    /// Maximum display depth of the dependency tree.
+    pub max_display_depth: u32,
 }
 
 #[derive(PartialEq)]
@@ -241,6 +243,7 @@ fn print(
             symbols,
             opts.prefix,
             opts.no_dedupe,
+            opts.max_display_depth,
             &mut visited_deps,
             &mut levels_continue,
             &mut print_stack,
@@ -259,6 +262,7 @@ fn print_node<'a>(
     symbols: &Symbols,
     prefix: Prefix,
     no_dedupe: bool,
+    max_display_depth: u32,
     visited_deps: &mut HashSet<usize>,
     levels_continue: &mut Vec<bool>,
     print_stack: &mut Vec<usize>,
@@ -316,6 +320,7 @@ fn print_node<'a>(
             symbols,
             prefix,
             no_dedupe,
+            max_display_depth,
             visited_deps,
             levels_continue,
             print_stack,
@@ -334,6 +339,7 @@ fn print_dependencies<'a>(
     symbols: &Symbols,
     prefix: Prefix,
     no_dedupe: bool,
+    max_display_depth: u32,
     visited_deps: &mut HashSet<usize>,
     levels_continue: &mut Vec<bool>,
     print_stack: &mut Vec<usize>,
@@ -364,19 +370,22 @@ fn print_dependencies<'a>(
 
     let mut it = deps.iter().peekable();
     while let Some(dependency) = it.next() {
-        levels_continue.push(it.peek().is_some());
-        print_node(
-            config,
-            graph,
-            *dependency,
-            format,
-            symbols,
-            prefix,
-            no_dedupe,
-            visited_deps,
-            levels_continue,
-            print_stack,
-        );
-        levels_continue.pop();
+        if levels_continue.len() + 1 <= max_display_depth as usize {
+            levels_continue.push(it.peek().is_some());
+            print_node(
+                config,
+                graph,
+                *dependency,
+                format,
+                symbols,
+                prefix,
+                no_dedupe,
+                max_display_depth,
+                visited_deps,
+                levels_continue,
+                print_stack,
+            );
+            levels_continue.pop();
+        }
     }
 }