1 files changed, 51 insertions, 14 deletions

diff --git a/src/main.rs b/src/main.rs
index 35417ad..00ac301 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,17 +1,21 @@
 extern crate humansize;
 extern crate ignore;
 extern crate num_cpus;
+#[macro_use]
+extern crate clap;
 
 use std::collections::HashSet;
 use std::os::unix::fs::MetadataExt;
+use std::path::Path;
 use std::sync::mpsc::channel;
 use std::thread;
 
+use clap::{App, AppSettings, Arg};
 use humansize::{file_size_opts, FileSize};
 use ignore::WalkBuilder;
 
-fn main() {
-    let mut builder = WalkBuilder::new("./");
+fn get_size<P: AsRef<Path>>(p: P, num_threads: usize) -> u64 {
+    let mut builder = WalkBuilder::new(p);
     builder.hidden(false);
     builder.parents(false);
     builder.ignore(false);
@@ -20,12 +24,7 @@ fn main() {
     builder.git_exclude(false);
     builder.follow_links(false);
 
-    // Setting the number of threads to 3x the number of cores is a good tradeoff between
-    // cold-cache and warm-cache runs. For a cold disk cache, we are limited by disk IO and
-    // therefore want the number of threads to be rather large in order for the IO scheduler to
-    // plan ahead. On the other hand, the number of threads shouldn't be too high for warm disk
-    // caches where we would otherwise pay a higher synchronization overhead.
-    builder.threads(3 * num_cpus::get());
+    builder.threads(num_threads);
 
     let walker = builder.build_parallel();
 
@@ -44,11 +43,8 @@ fn main() {
                 total += size;
             }
         }
-        println!(
-            "{} ({} bytes)",
-            total.file_size(file_size_opts::DECIMAL).unwrap(),
-            total
-        );
+
+        total
     });
 
     walker.run(|| {
@@ -86,5 +82,46 @@ fn main() {
     });
 
     drop(tx);
-    receiver_thread.join().ok();
+    receiver_thread.join().unwrap()
+}
+
+fn print_result(size: u64) {
+    println!(
+        "{} ({} bytes)",
+        size.file_size(file_size_opts::DECIMAL).unwrap(),
+        size
+    );
+}
+
+fn main() {
+    let app = App::new(crate_name!())
+        .setting(AppSettings::ColorAuto)
+        .setting(AppSettings::ColoredHelp)
+        .setting(AppSettings::DeriveDisplayOrder)
+        .setting(AppSettings::UnifiedHelpMessage)
+        .version(crate_version!())
+        .about("Compute disk usage for the current directory")
+        .arg(
+            Arg::with_name("threads")
+                .long("threads")
+                .short("j")
+                .value_name("N")
+                .takes_value(true)
+                .help("Set the number of threads (default: 3 x num cores)"),
+        );
+
+    let matches = app.get_matches();
+
+    // Setting the number of threads to 3x the number of cores is a good tradeoff between
+    // cold-cache and warm-cache runs. For a cold disk cache, we are limited by disk IO and
+    // therefore want the number of threads to be rather large in order for the IO scheduler to
+    // plan ahead. On the other hand, the number of threads shouldn't be too high for warm disk
+    // caches where we would otherwise pay a higher synchronization overhead.
+    let num_threads = matches
+        .value_of("threads")
+        .and_then(|t| t.parse().ok())
+        .unwrap_or(3 * num_cpus::get());
+
+    let size = get_size(".", num_threads);
+    print_result(size);
 }