diff --git a/src/display.rs b/src/display.rs index ac9a6ab..0280e77 100644 --- a/src/display.rs +++ b/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>( 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>( 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, 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 { diff --git a/src/main.rs b/src/main.rs index 61e613d..2ac4d4b 100644 --- a/src/main.rs +++ b/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)] diff --git a/src/utils/mod.rs b/src/utils/mod.rs index a276b11..dd1a512 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -35,9 +35,9 @@ pub fn simplify_dir_names(filenames: Vec<&str>) -> HashSet { } pub fn get_dir_tree( - top_level_names: HashSet, + top_level_names: &HashSet, apparent_size: bool, -) -> (bool, HashMap, HashSet) { +) -> (bool, HashMap) { let mut permissions = 0; let mut inodes: HashSet<(u64, u64)> = HashSet::new(); let mut data: HashMap = 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) -> Vec<(String, u64)> { +pub fn sort(data: HashMap) -> 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, +) -> 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::*;