Sort options

Signed-off-by: Slendi <slendi@socopon.com>

src/main.rs

@@ -1,41 +1,82 @@
 use std::{
+    collections::HashMap,
     path::PathBuf,
     sync::{
-        Arc,
+        Arc, Mutex,
         atomic::{AtomicU64, AtomicUsize, Ordering},
     },
     thread,
 };
 
-use clap::Parser;
+use clap::{Parser, ValueEnum};
 use crossbeam_channel as chan;
 use humansize::DECIMAL;
 
+#[derive(Debug, Clone, ValueEnum)]
+enum SortOrder {
+    Increasing,
+    Decreasing,
+}
+
+#[derive(Debug, Clone, ValueEnum)]
+enum SortMode {
+    Path,
+    Size,
+}
+
 #[derive(Debug, Parser)]
 struct Args {
     path: PathBuf,
     #[arg(short = 'j', long = "jobs")]
     jobs: Option<usize>,
+    #[arg(short = 'd', long = "dir-sizes")]
+    dir_sizes: bool,
+    #[arg(short, long)]
+    bytes: bool,
+    #[arg(short = 'O', long, value_enum, default_value_t = SortOrder::Increasing)]
+    sort_order: SortOrder,
+    #[arg(short = 'M', long, value_enum, default_value_t = SortMode::Path)]
+    sort_mode: SortMode,
 }
 
 fn main() {
     let args = Args::parse();
     let jobs = args.jobs.unwrap_or_else(num_cpus::get).max(1);
+    let root = args.path.clone();
 
     let (tx, rx) = chan::unbounded::<PathBuf>();
 
     let pending = Arc::new(AtomicUsize::new(1));
     let total = Arc::new(AtomicU64::new(0));
+    let dir_sizes = if args.dir_sizes {
+        Some(Arc::new(Mutex::new(HashMap::new())))
+    } else {
+        None
+    };
 
     let mut workers = Vec::with_capacity(jobs);
     for _ in 0..jobs {
         let rx = rx.clone();
         let total = total.clone();
+        let dir_sizes = dir_sizes.clone();
+        let root = root.clone();
         workers.push(thread::spawn(move || {
             while let Ok(path) = rx.recv() {
                 if let Ok(meta) = std::fs::metadata(&path) {
                     if meta.is_file() {
-                        total.fetch_add(meta.len(), Ordering::Relaxed);
+                        let size = meta.len();
+                        total.fetch_add(size, Ordering::Relaxed);
+                        if let Some(dir_sizes) = dir_sizes.as_ref() {
+                            let mut map = dir_sizes.lock().unwrap();
+                            let mut current = path.parent();
+                            while let Some(dir) = current {
+                                if !dir.starts_with(&root) {
+                                    break;
+                                }
+                                *map.entry(dir.to_path_buf()).or_insert(0) += size;
+                                current = dir.parent();
+                            }
+                        }
                     }
                 }
             }
@@ -98,10 +139,40 @@ fn main() {
         let _ = w.join();
     }
 
-    let size = total.load(Ordering::Relaxed);
-    println!(
-        "Computed size sum: {} ({} bytes)",
-        humansize::format_size(size, DECIMAL),
-        size
-    );
+    if let Some(dir_sizes) = dir_sizes {
+        let mut entries: Vec<(PathBuf, u64)> = dir_sizes
+            .lock()
+            .unwrap()
+            .iter()
+            .map(|(path, size)| (path.clone(), *size))
+            .collect();
+
+        entries.sort_by(|(lp, li), (rp, ri)| {
+            let ord = match args.sort_mode {
+                SortMode::Path => lp.cmp(rp),
+                SortMode::Size => li.cmp(&ri),
+            };
+
+            match args.sort_order {
+                SortOrder::Increasing => ord,
+                SortOrder::Decreasing => ord.reverse(),
+            }
+        });
+        for (path, size) in entries {
+            let size = if args.bytes {
+                size.to_string()
+            } else {
+                humansize::format_size(size, DECIMAL)
+            };
+            println!("{}: {}", size, path.display());
+        }
+    }
+
+    let size = total.load(Ordering::Relaxed);
+    let size = if args.bytes {
+        size.to_string()
+    } else {
+        humansize::format_size(size, DECIMAL)
+    };
+    println!("Total: {}", size);
 }