From 2f7a88e8dcd9354384047dd609e18f3ddc6b9be3 Mon Sep 17 00:00:00 2001
From: Bob <bootandy@gmail.com>
Date: Tue, 1 Oct 2019 21:51:34 +0100
Subject: [PATCH] Fix issues from running on root directory

clean up: 80338f4

Fixes -d flag to work again. Add test to stop regression
---
 src/display.rs   | 26 ++++++++++++--------------
 src/tests.rs     | 26 ++++++++++++--------------
 src/utils/mod.rs |  5 ++---
 3 files changed, 26 insertions(+), 31 deletions(-)

diff --git a/src/display.rs b/src/display.rs
index bac6632..63e7418 100644
--- a/src/display.rs
+++ b/src/display.rs
@@ -3,7 +3,7 @@ extern crate ansi_term;
 use self::ansi_term::Colour::Fixed;
 use self::ansi_term::Style;
 use std::collections::HashSet;
-use utils::ensure_end_slash;
+use utils::{ensure_end_slash, strip_end_slash};
 
 static UNITS: [char; 4] = ['T', 'G', 'M', 'K'];
 
@@ -21,15 +21,7 @@ pub fn draw_it(
 
     for &(ref k, _) in to_display.iter() {
         if base_dirs.contains(k) {
-            display_node(
-                &k,
-                &mut found,
-                &to_display,
-                true,
-                short_paths,
-                depth,
-                "─┬",
-            );
+            display_node(&k, &mut found, &to_display, true, short_paths, depth, "─┬");
         }
     }
 }
@@ -69,16 +61,22 @@ fn display_node<S: Into<String>>(
             print_this_node(node_to_print, size, is_biggest, short_paths, is.as_ref());
             let new_indent = clean_indentation_string(is);
 
-            let ntp_with_slash = ensure_end_slash(node_to_print);
-            let num_slashes = ntp_with_slash.matches('/').count();
+            let ntp_with_slash = strip_end_slash(node_to_print);
+
+            // Annoying edge case for when run on root directory
+            let num_slashes = if ntp_with_slash == "/" {
+                1
+            } else {
+                ntp_with_slash.matches('/').count() + 1
+            };
             let mut num_siblings = count_siblings(to_display, num_slashes - 1, node_to_print);
 
             let mut is_biggest = true;
             for &(ref k, _) in to_display.iter() {
-                if k.starts_with(ntp_with_slash.as_str()) && k.matches('/').count() == num_slashes {
+                let temp = String::from(ensure_end_slash(node_to_print));
+                if k.starts_with(temp.as_str()) && k.matches('/').count() == num_slashes {
                     num_siblings -= 1;
                     let has_children = has_children(to_display, new_depth, k, num_slashes);
-                    //println!("{:?} {:?} {:?}", k ,num_siblings, has_children);
                     display_node(
                         k,
                         found,
diff --git a/src/tests.rs b/src/tests.rs
index a5fa9f7..0e72534 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -45,13 +45,7 @@ fn main_output(short_paths: bool) -> String {
 {}
 {}",
         format_string("src/test_dir", true, short_paths, " 4.0K", "─┬"),
-        format_string(
-            "src/test_dir/many",
-            true,
-            short_paths,
-            " 4.0K",
-            " └─┬",
-        ),
+        format_string("src/test_dir/many", true, short_paths, " 4.0K", " └─┬",),
         format_string(
             "src/test_dir/many/hello_file",
             true,
@@ -77,13 +71,7 @@ fn main_output(short_paths: bool) -> String {
 {}
 {}",
         format_string("src/test_dir", true, short_paths, "  12K", "─┬"),
-        format_string(
-            "src/test_dir/many",
-            true,
-            short_paths,
-            " 8.0K",
-            " └─┬",
-        ),
+        format_string("src/test_dir/many", true, short_paths, " 8.0K", " └─┬",),
         format_string(
             "src/test_dir/many/hello_file",
             true,
@@ -121,6 +109,16 @@ pub fn test_apparent_size() {
         .unwrap();
 }
 
+#[test]
+pub fn test_d_flag_works() {
+    // We should see the top level directory but not the sub dirs / files:
+    assert_cli::Assert::main_binary()
+        .with_args(&["-d", "1", "-s", "src/test_dir"])
+        .stdout()
+        .doesnt_contain("hello_file")
+        .unwrap();
+}
+
 fn build_temp_file(dir: &TempDir) -> (PathBuf) {
     let file_path = dir.path().join("notes.txt");
     let mut file = File::create(&file_path).unwrap();
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 00e8e0d..a8f628a 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -56,7 +56,6 @@ pub fn ensure_end_slash(s: &str) -> String {
     new_name + "/"
 }
 
-// TODO fairly sure we shouldn't need this func
 pub fn strip_end_slash(s: &str) -> String {
     let mut new_name = String::from(s);
     while (new_name.ends_with('/') || new_name.ends_with("/.")) && new_name.len() > 1 {
@@ -134,11 +133,11 @@ pub fn trim_deep_ones(
     let mut result: Vec<(String, u64)> = vec![];
 
     for name in top_level_names {
-        let my_max_depth = name.matches('/').count() - 1 + max_depth as usize;
+        let my_max_depth = name.matches('/').count() + max_depth as usize;
         let name_ref: &str = name.as_ref();
 
         for &(ref k, ref v) in input.iter() {
-            if k.starts_with(name_ref) && k.matches('/').count() - 1 <= my_max_depth {
+            if k.starts_with(name_ref) && k.matches('/').count() <= my_max_depth {
                 result.push((k.clone(), *v));
             }
         }