use std::collections::HashMap;
use std::path::{Path, PathBuf};
use crate::error::RustlocError;
use crate::query::options::{Aggregation, LineTypes};
use crate::source::filter::{discover_files, discover_files_in_dirs, FilterConfig};
use crate::source::workspace::{CrateInfo, WorkspaceInfo};
use crate::Result;
use super::stats::{CrateStats, FileStats, Locs, ModuleStats};
use super::visitor::gather_stats_for_path;
#[derive(Debug, Clone)]
pub struct CountOptions {
pub crate_filter: Vec<String>,
pub file_filter: FilterConfig,
pub aggregation: Aggregation,
pub line_types: LineTypes,
}
impl Default for CountOptions {
fn default() -> Self {
Self {
crate_filter: Vec::new(),
file_filter: FilterConfig::new(),
aggregation: Aggregation::Total,
line_types: LineTypes::default(),
}
}
}
impl CountOptions {
pub fn new() -> Self {
Self::default()
}
pub fn crates(mut self, names: Vec<String>) -> Self {
self.crate_filter = names;
self
}
pub fn filter(mut self, filter: FilterConfig) -> Self {
self.file_filter = filter;
self
}
pub fn aggregation(mut self, level: Aggregation) -> Self {
self.aggregation = level;
self
}
pub fn line_types(mut self, types: LineTypes) -> Self {
self.line_types = types;
self
}
}
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct CountResult {
pub root: PathBuf,
pub file_count: usize,
pub total: Locs,
pub crates: Vec<CrateStats>,
pub files: Vec<FileStats>,
pub modules: Vec<ModuleStats>,
}
impl CountResult {
pub fn new() -> Self {
Self::default()
}
pub fn filter(&self, types: LineTypes) -> Self {
Self {
root: self.root.clone(),
file_count: self.file_count,
total: self.total.filter(types),
crates: self.crates.iter().map(|c| c.filter(types)).collect(),
files: self.files.iter().map(|f| f.filter(types)).collect(),
modules: self.modules.iter().map(|m| m.filter(types)).collect(),
}
}
}
pub fn count_workspace(path: impl AsRef<Path>, options: CountOptions) -> Result<CountResult> {
let workspace = WorkspaceInfo::discover(path)?;
let crates: Vec<&CrateInfo> = if options.crate_filter.is_empty() {
workspace.crates.iter().collect()
} else {
let names: Vec<&str> = options.crate_filter.iter().map(|s| s.as_str()).collect();
workspace
.crates
.iter()
.filter(|c| names.contains(&c.name.as_str()))
.collect()
};
let mut result = CountResult::new();
result.root = workspace.root.clone();
let include_files = matches!(options.aggregation, Aggregation::ByFile);
let include_modules = matches!(options.aggregation, Aggregation::ByModule);
let include_crates = matches!(
options.aggregation,
Aggregation::ByCrate | Aggregation::ByModule | Aggregation::ByFile
);
for crate_info in &crates {
let crate_stats = count_crate(crate_info, &options)?;
result.total += crate_stats.stats;
result.file_count += crate_stats.files.len();
if include_files {
result.files.extend(crate_stats.files.clone());
}
if include_modules {
let crate_modules = aggregate_modules(&crate_stats.files, &crate_info.name, crate_info);
result.modules.extend(crate_modules);
}
if include_crates {
result.crates.push(crate_stats);
}
}
if include_modules {
result.modules.sort_by(|a, b| a.name.cmp(&b.name));
}
Ok(result.filter(options.line_types))
}
pub fn compute_module_name(file_path: &Path, src_root: &Path) -> String {
let relative = file_path.strip_prefix(src_root).unwrap_or(file_path);
let mut components: Vec<&str> = relative
.components()
.filter_map(|c| c.as_os_str().to_str())
.collect();
if components.is_empty() {
return String::new();
}
let filename = components.pop().unwrap_or("");
let stem = filename.strip_suffix(".rs").unwrap_or(filename);
if components.is_empty() && (stem == "lib" || stem == "main" || stem == "mod") {
return String::new();
}
if stem == "mod" {
return components.join("::");
}
if !components.is_empty() {
components.join("::")
} else {
stem.to_string()
}
}
fn aggregate_modules(
files: &[FileStats],
crate_name: &str,
crate_info: &CrateInfo,
) -> Vec<ModuleStats> {
let mut module_map: HashMap<String, ModuleStats> = HashMap::new();
for file in files {
let src_root = crate_info
.src_dirs
.iter()
.find(|dir| file.path.starts_with(dir))
.map(|p| p.as_path())
.unwrap_or(&crate_info.root);
let local_module = compute_module_name(&file.path, src_root);
let full_module_name = if local_module.is_empty() {
crate_name.to_string()
} else {
format!("{}::{}", crate_name, local_module)
};
let module = module_map
.entry(full_module_name.clone())
.or_insert_with(|| ModuleStats::new(full_module_name));
module.add_file(file.path.clone(), file.stats);
}
module_map.into_values().collect()
}
fn count_crate(crate_info: &CrateInfo, options: &CountOptions) -> Result<CrateStats> {
let dirs: Vec<&Path> = crate_info.all_dirs();
let files = discover_files_in_dirs(&dirs, &options.file_filter)?;
let mut crate_stats = CrateStats::new(crate_info.name.clone(), crate_info.root.clone());
for file_path in files {
let stats = gather_stats_for_path(&file_path)?;
let file_stats = FileStats::new(file_path, stats);
crate_stats.add_file(file_stats);
}
Ok(crate_stats)
}
pub fn count_directory(path: impl AsRef<Path>, filter: &FilterConfig) -> Result<CountResult> {
let path = path.as_ref();
if !path.exists() {
return Err(RustlocError::PathNotFound(path.to_path_buf()));
}
let files = discover_files(path, filter)?;
let mut result = CountResult::new();
result.root = path.to_path_buf();
for file_path in files {
let stats = gather_stats_for_path(&file_path)?;
result.total += stats;
result.file_count += 1;
result.files.push(FileStats::new(file_path, stats));
}
Ok(result)
}
pub fn count_file(path: impl AsRef<Path>) -> Result<Locs> {
gather_stats_for_path(path)
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
use tempfile::tempdir;
fn create_rust_file(path: &Path, content: &str) {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).unwrap();
}
fs::write(path, content).unwrap();
}
fn create_simple_project(root: &Path) {
fs::write(
root.join("Cargo.toml"),
r#"[package]
name = "test-project"
version = "0.1.0"
edition = "2021"
"#,
)
.unwrap();
create_rust_file(
&root.join("src/main.rs"),
r#"fn main() {
println!("Hello");
}
#[cfg(test)]
mod tests {
#[test]
fn test_main() {
assert!(true);
}
}
"#,
);
create_rust_file(
&root.join("src/lib.rs"),
r#"//! Library documentation
/// A public function
pub fn hello() {
println!("Hello from lib");
}
"#,
);
}
fn create_workspace(root: &Path) {
fs::write(
root.join("Cargo.toml"),
r#"[workspace]
members = ["crate-a", "crate-b"]
"#,
)
.unwrap();
fs::create_dir_all(root.join("crate-a/src")).unwrap();
fs::write(
root.join("crate-a/Cargo.toml"),
r#"[package]
name = "crate-a"
version = "0.1.0"
edition = "2021"
"#,
)
.unwrap();
create_rust_file(
&root.join("crate-a/src/lib.rs"),
r#"pub fn a() {
println!("A");
}
"#,
);
fs::create_dir_all(root.join("crate-b/src")).unwrap();
fs::write(
root.join("crate-b/Cargo.toml"),
r#"[package]
name = "crate-b"
version = "0.1.0"
edition = "2021"
"#,
)
.unwrap();
create_rust_file(
&root.join("crate-b/src/lib.rs"),
r#"pub fn b() {
println!("B");
}
// A comment
"#,
);
}
#[test]
fn test_count_directory() {
let temp = tempdir().unwrap();
let src = temp.path().join("src");
fs::create_dir_all(&src).unwrap();
create_rust_file(
&src.join("main.rs"),
r#"fn main() {
println!("Hello");
}
"#,
);
let filter = FilterConfig::new();
let result = count_directory(&src, &filter).unwrap();
assert_eq!(result.files.len(), 1);
assert_eq!(result.total.code, 3);
}
#[test]
fn test_count_file() {
let temp = tempdir().unwrap();
let file = temp.path().join("test.rs");
create_rust_file(
&file,
r#"/// Doc comment
fn foo() {
// Regular comment
let x = 1;
}
"#,
);
let stats = count_file(&file).unwrap();
assert_eq!(stats.docs, 1);
assert_eq!(stats.code, 3); assert_eq!(stats.comments, 1);
}
#[test]
fn test_count_workspace() {
let temp = tempdir().unwrap();
create_workspace(temp.path());
let result = count_workspace(
temp.path(),
CountOptions::new().aggregation(Aggregation::ByCrate),
)
.unwrap();
assert_eq!(result.crates.len(), 2);
}
#[test]
fn test_count_workspace_filtered() {
let temp = tempdir().unwrap();
create_workspace(temp.path());
let options = CountOptions::new()
.crates(vec!["crate-a".to_string()])
.aggregation(Aggregation::ByCrate);
let result = count_workspace(temp.path(), options).unwrap();
assert_eq!(result.crates.len(), 1);
assert_eq!(result.crates[0].name, "crate-a");
}
#[test]
fn test_count_workspace_with_file_stats() {
let temp = tempdir().unwrap();
create_workspace(temp.path());
let options = CountOptions::new().aggregation(Aggregation::ByFile);
let result = count_workspace(temp.path(), options).unwrap();
assert_eq!(result.files.len(), 2);
}
#[test]
fn test_count_mixed_code_and_tests() {
let temp = tempdir().unwrap();
create_simple_project(temp.path());
let result = count_workspace(temp.path(), CountOptions::new()).unwrap();
assert!(result.total.code > 0);
assert!(result.total.tests > 0);
assert!(result.total.docs > 0);
}
#[test]
fn test_compute_module_name_root_files() {
let src = Path::new("/project/src");
assert_eq!(
compute_module_name(Path::new("/project/src/lib.rs"), src),
""
);
assert_eq!(
compute_module_name(Path::new("/project/src/main.rs"), src),
""
);
assert_eq!(
compute_module_name(Path::new("/project/src/error.rs"), src),
"error"
);
}
#[test]
fn test_compute_module_name_directory_aggregation() {
let src = Path::new("/project/src");
assert_eq!(
compute_module_name(Path::new("/project/src/data/mod.rs"), src),
"data"
);
assert_eq!(
compute_module_name(Path::new("/project/src/data/counter.rs"), src),
"data"
);
assert_eq!(
compute_module_name(Path::new("/project/src/data/stats.rs"), src),
"data"
);
}
#[test]
fn test_compute_module_name_nested_directories() {
let src = Path::new("/project/src");
assert_eq!(
compute_module_name(Path::new("/project/src/data/sub/foo.rs"), src),
"data::sub"
);
assert_eq!(
compute_module_name(Path::new("/project/src/data/sub/mod.rs"), src),
"data::sub"
);
}
#[test]
fn test_compute_module_name_new_style_module() {
let src = Path::new("/project/src");
assert_eq!(
compute_module_name(Path::new("/project/src/data.rs"), src),
"data"
);
assert_eq!(
compute_module_name(Path::new("/project/src/data/counter.rs"), src),
"data"
);
}
#[test]
fn test_module_aggregation_groups_files_by_directory() {
let temp = tempdir().unwrap();
let root = temp.path();
fs::write(
root.join("Cargo.toml"),
"[package]\nname = \"test-proj\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
)
.unwrap();
create_rust_file(&root.join("src/lib.rs"), "pub mod data;\npub mod utils;\n");
create_rust_file(
&root.join("src/data/mod.rs"),
"pub mod counter;\npub mod stats;\n",
);
create_rust_file(&root.join("src/data/counter.rs"), "pub fn count() {}\n");
create_rust_file(&root.join("src/data/stats.rs"), "pub fn stats() {}\n");
create_rust_file(&root.join("src/utils.rs"), "pub fn helper() {}\n");
let options = CountOptions::new().aggregation(Aggregation::ByModule);
let result = count_workspace(root, options).unwrap();
let module_names: Vec<&str> = result.modules.iter().map(|m| m.name.as_str()).collect();
assert!(module_names.contains(&"test-proj::data"));
assert!(module_names.contains(&"test-proj::utils"));
assert!(module_names.contains(&"test-proj"));
assert!(!module_names.contains(&"test-proj::data::counter"));
assert!(!module_names.contains(&"test-proj::data::stats"));
}
}