use std::{ fs, path::{Path, PathBuf}, thread, }; use ccusage_core::{LoadedEntry, cli::SharedArgs, date_within_range}; pub mod jsonl; pub mod report; pub use report::print_table_for_agent; pub fn collect_usage_files(dir: &Path, files: &mut Vec) { collect_files_with_extension(dir, "jsonl", files); } pub fn collect_files_with_extension(dir: &Path, extension: &str, files: &mut Vec) { let Ok(entries) = fs::read_dir(dir) else { return; }; for entry in entries.filter_map(std::result::Result::ok) { let Ok(file_type) = entry.file_type() else { continue; }; let path = entry.path(); if file_type.is_file() && path.extension().is_some_and(|ext| ext == extension) { files.push(path); } else if file_type.is_dir() { collect_files_with_extension(&path, extension, files); } } } pub fn filter_loaded_entries_by_date(entries: &mut Vec, shared: &SharedArgs) { if shared.since.is_none() && shared.until.is_none() { return; } entries.retain(|entry| { date_within_range( &entry.date, shared.since.as_deref(), shared.until.as_deref(), ) }); } pub fn chunk_file_indexes_by_size(files: &[PathBuf], chunk_count: usize) -> Vec> { // Callers derive the count from available_parallelism, which can be 0 for a // caller that clamps against an empty file list; one chunk still returns // every index rather than indexing an empty vector. let chunk_count = chunk_count.max(1); let mut weighted_indexes = Vec::with_capacity(files.len()); for (index, file) in files.iter().enumerate() { let size = fs::metadata(file).map_or(0, |metadata| metadata.len()); weighted_indexes.push((index, size)); } weighted_indexes.sort_unstable_by(|a, b| match b.1.cmp(&a.1) { std::cmp::Ordering::Equal => a.0.cmp(&b.0), order => order, }); let mut chunks = vec![Vec::new(); chunk_count]; let mut chunk_sizes = vec![0_u64; chunk_count]; for (index, size) in weighted_indexes { let mut target = 0; for candidate in 1..chunk_sizes.len() { if chunk_sizes[candidate] < chunk_sizes[target] { target = candidate; } } chunks[target].push(index); chunk_sizes[target] = chunk_sizes[target].saturating_add(size); } chunks .into_iter() .filter(|chunk| !chunk.is_empty()) .collect() } /// Reads `files` by applying `read` to each path and returns results in file order. pub fn read_files_parallel(files: &[PathBuf], single_thread: bool, read: F) -> Vec where T: Send, F: Fn(&Path) -> T + Sync, { let worker_count = if single_thread { 1 } else { thread::available_parallelism() .map(usize::from) .unwrap_or(1) .min(files.len()) }; if worker_count <= 1 { return files.iter().map(|file| read(file.as_path())).collect(); } let chunks = chunk_file_indexes_by_size(files, worker_count); let read = &read; thread::scope(|scope| { let mut handles = Vec::with_capacity(chunks.len()); for chunk in chunks { handles.push(scope.spawn(move || { chunk .into_iter() .map(|index| (index, read(files[index].as_path()))) .collect::>() })); } let mut results: Vec> = Vec::with_capacity(files.len()); results.resize_with(files.len(), || None); for (index, value) in handles .into_iter() .flat_map(|handle| handle.join().expect("file read worker panicked")) { results[index] = Some(value); } results .into_iter() .map(|value| value.expect("file read worker returned every file")) .collect() }) } #[cfg(test)] mod tests { use super::{chunk_file_indexes_by_size, read_files_parallel}; use ccusage_test_support::Fixture; #[test] fn preserves_file_order_and_matches_single_thread() { let fixture = Fixture::new(); let files = (0..256) .map(|index| { let body = "x".repeat((index % 17) * 64 + 1); fixture.write_file(format!("file-{index:03}.txt"), format!("{index}:{body}")) }) .collect::>(); let read = |path: &std::path::Path| { let content = std::fs::read_to_string(path).unwrap(); content.split(':').next().unwrap().to_string() }; let single = read_files_parallel(&files, true, read); let multi = read_files_parallel(&files, false, read); let expected = (0..256).map(|index| index.to_string()).collect::>(); assert_eq!(single, expected); assert_eq!(multi, expected); } #[test] fn treats_a_zero_chunk_count_as_one_chunk() { let fixture = Fixture::new(); let files = vec![fixture.write_file("only.txt", "body")]; assert_eq!(chunk_file_indexes_by_size(&files, 0), vec![vec![0]]); } #[test] fn handles_empty_input() { let empty: Vec = Vec::new(); assert!(read_files_parallel(&empty, false, |_| 0_u8).is_empty()); } }