Refactor & Optimize use of depth

Code Refactored

Using Depth now prunes the tree below the displayable depth so display
calls finish quicker.

src/display.rs

@@ -19,7 +19,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, "─┬");
         }
     }
 }
@@ -55,30 +55,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 is = indentation_str.into();
             let ntp: &str = node_to_print.as_ref();
-
-            print_this_node(ntp, size, is_biggest, short_paths, is.as_ref());
-            let new_indent_str = clean_indentation_string(is);
-
             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());
+            let new_indent = clean_indentation_string(is);
+
             let mut num_siblings = count_siblings(to_display, num_slashes, ntp);
 
             let mut is_biggest = true;
             for &(ref k, _) in to_display.iter() {
                 if k.starts_with(ntp) && k.matches('/').count() == num_slashes + 1 {
                     num_siblings -= 1;
-
-                    let mut has_children = false;
-                    if new_depth.is_none() || new_depth.unwrap() != 1 {
-                        for &(ref k2, _) in to_display.iter() {
-                            let kk: &str = k.as_ref();
-                            if k2.starts_with(kk) && k2.matches('/').count() == num_slashes + 2 {
-                                has_children = true;
-                            }
-                        }
-                    }
-
+                    let has_children = has_children(to_display, new_depth, k, num_slashes + 1);
                     display_node(
                         k,
                         found,
@@ -86,7 +76,7 @@ fn display_node<S: Into<String>>(
                         is_biggest,
                         short_paths,
                         new_depth,
-                        new_indent_str.to_string() + get_tree_chars(num_siblings, has_children),
+                        new_indent.to_string() + get_tree_chars(num_siblings != 0, has_children),
                     );
                     is_biggest = false;
                 }
@@ -115,8 +105,19 @@ fn count_siblings(to_display: &Vec<(String, u64)>, num_slashes: usize, ntp: &str
     })
 }
 
-fn get_tree_chars(num_siblings: u64, has_children: bool) -> &'static str {
-    if num_siblings == 0 {
+fn has_children(to_display: &Vec<(String, u64)>, new_depth: Option<u64>, ntp: &str, num_slashes: usize) -> bool {
+    if new_depth.is_none() || new_depth.unwrap() != 1 {
+        for &(ref k2, _) in to_display.iter() {
+            if k2.starts_with(ntp) && k2.matches('/').count() == num_slashes + 1 {
+                return true;
+            }
+        }
+    }
+    return false;
+}
+
+fn get_tree_chars(has_smaller_siblings: bool, has_children: bool) -> &'static str {
+    if !has_smaller_siblings {
         if has_children {
             "└─┬"
         } else {

src/main.rs

@@ -5,7 +5,7 @@ extern crate walkdir;
 
 use self::display::draw_it;
 use clap::{App, AppSettings, Arg};
-use utils::{find_big_ones, get_dir_tree, simplify_dir_names, sort};
+use utils::{find_big_ones, get_dir_tree, simplify_dir_names, sort, trim_deep_ones};
 
 mod display;
 mod utils;
@@ -83,22 +83,15 @@ fn main() {
     let use_full_path = options.is_present("display_full_paths");
 
     let simplified_dirs = simplify_dir_names(target_dirs);
-    let (permissions, nodes, top_level_names) = get_dir_tree(simplified_dirs, use_apparent_size);
+    let (permissions, nodes) = get_dir_tree(&simplified_dirs, use_apparent_size);
     let sorted_data = sort(nodes);
     let biggest_ones = {
-        if depth.is_none() {
-            find_big_ones(sorted_data, number_of_lines)
-        } else {
-            sorted_data
+        match depth {
+            None => find_big_ones(sorted_data, number_of_lines + simplified_dirs.len()),
+            Some(d) => trim_deep_ones(sorted_data, d, &simplified_dirs)
         }
     };
-    draw_it(
-        permissions,
-        !use_full_path,
-        depth,
-        top_level_names,
-        biggest_ones,
-    );
+    draw_it(permissions, !use_full_path, depth, simplified_dirs, biggest_ones);
 }
 
 #[cfg(test)]

src/utils/mod.rs

@@ -35,9 +35,9 @@ pub fn simplify_dir_names(filenames: Vec<&str>) -> HashSet<String> {
 }
 
 pub fn get_dir_tree(
-    top_level_names: HashSet<String>,
+    top_level_names: &HashSet<String>,
     apparent_size: bool,
-) -> (bool, HashMap<String, u64>, HashSet<String>) {
+) -> (bool, HashMap<String, u64>) {
     let mut permissions = 0;
     let mut inodes: HashSet<(u64, u64)> = HashSet::new();
     let mut data: HashMap<String, u64> = HashMap::new();
@@ -45,7 +45,7 @@ pub fn get_dir_tree(
     for b in top_level_names.iter() {
         examine_dir(&b, apparent_size, &mut inodes, &mut data, &mut permissions);
     }
-    (permissions == 0, data, top_level_names)
+    (permissions == 0, data)
 }
 
 fn strip_end_slashes(s: &str) -> String {
@@ -105,20 +105,40 @@ pub fn compare_tuple(a: &(String, u64), b: &(String, u64)) -> Ordering {
     }
 }
 
-pub fn sort<'a>(data: HashMap<String, u64>) -> Vec<(String, u64)> {
+pub fn sort(data: HashMap<String, u64>) -> Vec<(String, u64)> {
     let mut new_l: Vec<(String, u64)> = data.iter().map(|(a, b)| (a.clone(), *b)).collect();
     new_l.sort_by(|a, b| compare_tuple(&a, &b));
     new_l
 }
 
-pub fn find_big_ones<'a>(new_l: Vec<(String, u64)>, max_to_show: usize) -> Vec<(String, u64)> {
+pub fn find_big_ones(new_l: Vec<(String, u64)>, max_to_show: usize) -> Vec<(String, u64)> {
     if max_to_show > 0 && new_l.len() > max_to_show {
-        new_l[0..max_to_show + 1].to_vec()
+        new_l[0..max_to_show].to_vec()
     } else {
         new_l
     }
 }
 
+pub fn trim_deep_ones(
+    input: Vec<(String, u64)>,
+    max_depth: u64,
+    top_level_names: &HashSet<String>,
+) -> Vec<(String, u64)> {
+    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 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 {
+                result.push((k.clone(), *v));
+            }
+        }
+    }
+    result
+}
+
 mod tests {
     #[allow(unused_imports)]
     use super::*;