diff --git a/src/main.rs b/src/main.rs index a7d9a44..a75edaf 100644 --- a/src/main.rs +++ b/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, + #[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::(); 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(); } + 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); - println!( - "Computed size sum: {} ({} bytes)", - humansize::format_size(size, DECIMAL), - size - ); + let size = if args.bytes { + size.to_string() + } else { + humansize::format_size(size, DECIMAL) + }; + println!("Total: {}", size); }