Fix running on root dir /

Fixes: https://github.com/bootandy/dust/issues/22

Allows code to run on the root directory

src/display.rs

@@ -3,6 +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;
 
 static UNITS: [char; 4] = ['T', 'G', 'M', 'K'];
 
@@ -64,21 +65,20 @@ fn display_node<S: Into<String>>(
     match get_size(to_display, node_to_print) {
         None => println!("Can not find path: {}", node_to_print),
         Some(size) => {
-            let ntp: &str = node_to_print;
-            let num_slashes = node_to_print.matches('/').count();
-
             let is = indentation_str.into();
-            print_this_node(ntp, size, is_biggest, short_paths, is.as_ref());
+            print_this_node(node_to_print, size, is_biggest, short_paths, is.as_ref());
             let new_indent = clean_indentation_string(is);
 
-            let mut num_siblings = count_siblings(to_display, num_slashes, ntp);
+            let ntp_with_slash = ensure_end_slash(node_to_print);
+            let num_slashes = ntp_with_slash.matches('/').count();
+            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() {
-                let ntp_with_slash= String::from(ntp.to_owned() + "/");
+                if k.starts_with(ntp_with_slash.as_str()) && k.matches('/').count() == num_slashes {
-                if k.starts_with(ntp_with_slash.as_str()) && k.matches('/').count() == num_slashes + 1 {
                     num_siblings -= 1;
-                    let has_children = has_children(to_display, new_depth, k, num_slashes + 1);
+                    let has_children = has_children(to_display, new_depth, k, num_slashes);
+                    //println!("{:?} {:?} {:?}", k ,num_siblings, has_children);
                     display_node(
                         k,
                         found,
@@ -107,7 +107,7 @@ fn clean_indentation_string<S: Into<String>>(s: S) -> String {
 
 fn count_siblings(to_display: &[(String, u64)], num_slashes: usize, ntp: &str) -> u64 {
     to_display.iter().fold(0, |a, b| {
-        if b.0.starts_with(ntp) && b.0.matches('/').count() == num_slashes + 1 {
+        if b.0.starts_with(ntp) && b.0.as_str().matches('/').count() == num_slashes + 1 {
             a + 1
         } else {
             a
@@ -123,8 +123,9 @@ fn has_children(
 ) -> bool {
     if new_depth.is_none() || new_depth.unwrap() != 1 {
         for &(ref k2, _) in to_display.iter() {
-            let ntp_with_slash= String::from(ntp.to_owned() + "/");
+            let ntp_with_slash = String::from(ntp.to_owned() + "/");
-            if k2.starts_with(ntp_with_slash.as_str()) && k2.matches('/').count() == num_slashes + 1 {
+            if k2.starts_with(ntp_with_slash.as_str()) && k2.matches('/').count() == num_slashes + 1
+            {
                 return true;
             }
         }

src/utils/mod.rs

@@ -11,7 +11,7 @@ pub fn simplify_dir_names(filenames: Vec<&str>) -> HashSet<String> {
     let mut top_level_names: HashSet<String> = HashSet::new();
 
     for t in filenames {
-        let top_level_name = strip_end_slashes(t);
+        let top_level_name = ensure_end_slash(t);
         let mut can_add = true;
         let mut to_remove: Vec<String> = Vec::new();
 
@@ -27,7 +27,7 @@ pub fn simplify_dir_names(filenames: Vec<&str>) -> HashSet<String> {
             top_level_names.remove(&tr);
         }
         if can_add {
-            top_level_names.insert(top_level_name);
+            top_level_names.insert(strip_end_slash(t));
         }
     }
 
@@ -48,9 +48,18 @@ pub fn get_dir_tree(
     (permissions == 0, data)
 }
 
-fn strip_end_slashes(s: &str) -> String {
+pub fn ensure_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 {
+    while new_name.ends_with('/') || new_name.ends_with("/.") {
+        new_name.pop();
+    }
+    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 {
         new_name.pop();
     }
     new_name
@@ -125,11 +134,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() + max_depth as usize;
+        let my_max_depth = name.matches('/').count() - 1 + 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() <= my_max_depth {
+            if k.starts_with(name_ref) && k.matches('/').count() - 1 <= my_max_depth {
                 result.push((k.clone(), *v));
             }
         }
@@ -165,9 +174,9 @@ mod tests {
     #[test]
     fn test_simplify_dir_rm_subdir_and_not_substrings() {
         let mut correct = HashSet::new();
-        correct.insert("a/b".to_string());
-        correct.insert("c/a/b".to_string());
         correct.insert("b".to_string());
+        correct.insert("c/a/b".to_string());
+        correct.insert("a/b".to_string());
         assert_eq!(simplify_dir_names(vec!["a/b", "c/a/b/", "b"]), correct);
     }