diff --git a/src/main.rs b/src/main.rs index ce4fcfa..4a596fe 100644 --- a/src/main.rs +++ b/src/main.rs @@ -23,6 +23,13 @@ fn main() { .help("Depth to show") .takes_value(true), ) + .arg( + Arg::with_name("threads") + .short("t") + .long("threads") + .help("Number of threads to spawn simultaneously") + .takes_value(true), + ) .arg( Arg::with_name("number_of_lines") .short("n") @@ -67,19 +74,20 @@ fn main() { } }; - let depth = { - if options.is_present("depth") { - match value_t!(options.value_of("depth"), u64) { - Ok(v) => Some(v + 1), - Err(_) => { - eprintln!("Ignoring bad value for depth"); - None - } - } - } else { - None - } - }; + let threads = options.value_of("threads").and_then(|threads| { + threads + .parse::() + .map_err(|_| eprintln!("Ignoring bad value for threads: {:?}", threads)) + .ok() + }); + + let depth = options.value_of("depth").and_then(|depth| { + depth + .parse::() + .map(|v| v + 1) + .map_err(|_| eprintln!("Ignoring bad value for depth")) + .ok() + }); if options.is_present("depth") && number_of_lines != DEFAULT_NUMBER_OF_LINES { eprintln!("Use either -n or -d. Not both"); return; @@ -89,7 +97,7 @@ fn main() { let use_full_path = options.is_present("display_full_paths"); let simplified_dirs = simplify_dir_names(target_dirs); - let (permissions, nodes) = get_dir_tree(&simplified_dirs, use_apparent_size); + let (permissions, nodes) = get_dir_tree(&simplified_dirs, use_apparent_size, threads); let sorted_data = sort(nodes); let biggest_ones = { match depth { diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 55f9487..5b63be8 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -44,13 +44,21 @@ pub fn simplify_dir_names(filenames: Vec<&str>) -> HashSet { pub fn get_dir_tree( top_level_names: &HashSet, apparent_size: bool, + threads: Option, ) -> (bool, HashMap) { let mut permissions = 0; let mut inodes: HashSet<(u64, u64)> = HashSet::new(); let mut data: HashMap = HashMap::new(); for b in top_level_names.iter() { - examine_dir(&b, apparent_size, &mut inodes, &mut data, &mut permissions); + examine_dir( + &b, + apparent_size, + &mut inodes, + &mut data, + &mut permissions, + threads, + ); } (permissions == 0, data) } @@ -76,11 +84,15 @@ fn examine_dir( inodes: &mut HashSet<(u64, u64)>, data: &mut HashMap, file_count_no_permission: &mut u64, + cpus: Option, ) { - for entry in WalkDir::new(top_dir) + let mut iter = WalkDir::new(top_dir) .preload_metadata(true) - .skip_hidden(false) - { + .skip_hidden(false); + if let Some(cpus) = cpus { + iter = iter.num_threads(cpus); + } + for entry in iter { if let Ok(e) = entry { let maybe_size_and_inode = get_metadata(&e, apparent_size);