use clap::Parser;
use dashmap::DashSet;
use indicatif::{ParallelProgressIterator, ProgressBar, ProgressStyle};
use rayon::prelude::*;
use rloc::cli::Cli;
use rloc::diff;
use rloc::output::{self, render, OutputFormat};
use rloc::strip::{self, StripMode};
use std::fs::File;
use std::io::{self, BufWriter, Write};
use std::process::ExitCode;
use std::time::Instant;
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => {
eprintln!("error: {}", e);
ExitCode::FAILURE
}
}
}
fn run() -> Result<(), Box<dyn std::error::Error>> {
let cli = Cli::parse();
if cli.show_lang {
rloc::cli::show_languages();
return Ok(());
}
if cli.show_ext {
rloc::cli::show_extensions();
return Ok(());
}
if let Some(ref path) = cli.read_lang_def {
rloc::custom_langs::CustomLanguages::load(path)?;
}
if !cli.sum_reports.is_empty() {
return sum_reports(&cli);
}
if cli.strip_comments.is_some() || cli.strip_code.is_some() {
return run_strip(&cli);
}
if let Some(ref diff_path) = cli.diff {
return run_diff(&cli, diff_path);
}
if cli.threads > 0 {
rayon::ThreadPoolBuilder::new()
.num_threads(cli.threads)
.build_global()
.ok();
}
let mut walker_config = cli.to_walker_config()?;
let output_config = cli.to_output_config();
let start = Instant::now();
let temp_dir = if cli.extract_archives {
let temp = std::env::temp_dir().join(format!("rloc-{}", std::process::id()));
std::fs::create_dir_all(&temp)?;
let mut extra_paths = Vec::new();
for path in &walker_config.paths {
if path.is_file() && rloc::archive::is_archive(path) {
let archive_dest = temp.join(path.file_stem().unwrap_or_default());
std::fs::create_dir_all(&archive_dest)?;
if rloc::archive::extract_archive(path, &archive_dest).is_ok() {
extra_paths.push(archive_dest);
}
}
}
walker_config.paths.extend(extra_paths);
Some(temp)
} else {
None
};
let files = rloc::walker::walk_files(&walker_config);
if files.is_empty() {
if !cli.quiet {
eprintln!("No source files found.");
}
return Ok(());
}
let file_count = files.len();
let skip_uniqueness = walker_config.skip_uniqueness;
let seen_hashes: DashSet<u64> = DashSet::new();
let progress = if cli.quiet || output_config.format != OutputFormat::Table {
ProgressBar::hidden()
} else {
let pb = ProgressBar::new(file_count as u64);
pb.set_style(
ProgressStyle::default_bar()
.template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({per_sec})")
.unwrap()
.progress_chars("=>-"),
);
pb
};
let file_stats: Vec<_> = files
.into_par_iter()
.progress_with(progress.clone())
.filter_map(|entry| {
if !skip_uniqueness {
if let Ok(hash) = rloc::counter::compute_file_hash(&entry.path) {
if !seen_hashes.insert(hash) {
return None;
}
}
}
match rloc::counter::count_lines(&entry.path, entry.language) {
Ok(stats) if stats.total() > 0 => Some(stats),
Ok(_) => None,
Err(e) => {
if cli.verbose > 0 {
eprintln!("warning: {}: {}", entry.path.display(), e);
}
None
}
}
})
.collect();
progress.finish_and_clear();
let elapsed = start.elapsed();
let summary = rloc::stats::Summary::from_file_stats(file_stats).with_elapsed(elapsed);
if let Some(output_path) = cli.output_path() {
let file = File::create(output_path)?;
let mut writer = BufWriter::new(file);
render_to_writer(&summary, &output_config, &mut writer)?;
writer.flush()?;
} else {
render(&summary, &output_config)?;
}
if let Some(temp) = temp_dir {
let _ = std::fs::remove_dir_all(temp);
}
Ok(())
}
fn render_to_writer(
summary: &rloc::stats::Summary,
config: &output::OutputConfig,
out: &mut impl Write,
) -> io::Result<()> {
match config.format {
OutputFormat::Table => {
if !config.hide_rate {
if let Some(elapsed) = summary.elapsed {
writeln!(out)?;
write!(
out,
"{} files processed in {:.3}s",
summary.total_files,
elapsed.as_secs_f64()
)?;
if let (Some(fps), Some(lps)) =
(summary.files_per_second(), summary.lines_per_second())
{
write!(out, " ({:.0} files/s, {:.0} lines/s)", fps, lps)?;
}
writeln!(out)?;
}
}
writeln!(out)?;
writeln!(out, "Language Files Blank Comment Code")?;
writeln!(out, "─────────────────────────────────────────────────")?;
for lang in &summary.languages {
writeln!(
out,
"{:<14} {:>5} {:>8} {:>8} {:>8}",
lang.name, lang.files, lang.blanks, lang.comments, lang.code
)?;
}
writeln!(out, "─────────────────────────────────────────────────")?;
writeln!(
out,
"{:<14} {:>5} {:>8} {:>8} {:>8}",
"SUM",
summary.total_files,
summary.total_blanks,
summary.total_comments,
summary.total_code
)?;
Ok(())
}
OutputFormat::Json => {
let output = rloc::stats::JsonOutput::from(summary);
let json = serde_json::to_string_pretty(&output).map_err(io::Error::other)?;
writeln!(out, "{}", json)
}
OutputFormat::Csv => {
let mut writer = csv::Writer::from_writer(out);
writer.write_record(["Language", "Files", "Blank", "Comment", "Code"])?;
for lang in &summary.languages {
writer.write_record([
&lang.name,
&lang.files.to_string(),
&lang.blanks.to_string(),
&lang.comments.to_string(),
&lang.code.to_string(),
])?;
}
writer.write_record([
"SUM",
&summary.total_files.to_string(),
&summary.total_blanks.to_string(),
&summary.total_comments.to_string(),
&summary.total_code.to_string(),
])?;
writer.flush()?;
Ok(())
}
OutputFormat::Yaml => {
let output = rloc::stats::JsonOutput::from(summary);
let yaml = serde_yaml::to_string(&output).map_err(io::Error::other)?;
write!(out, "{}", yaml)
}
OutputFormat::Markdown => {
writeln!(out, "| Language | Files | Blank | Comment | Code |")?;
writeln!(out, "|----------|------:|------:|--------:|-----:|")?;
for lang in &summary.languages {
writeln!(
out,
"| {} | {} | {} | {} | {} |",
lang.name, lang.files, lang.blanks, lang.comments, lang.code
)?;
}
writeln!(
out,
"| **SUM** | **{}** | **{}** | **{}** | **{}** |",
summary.total_files,
summary.total_blanks,
summary.total_comments,
summary.total_code
)
}
OutputFormat::Sql => {
writeln!(out, "CREATE TABLE t (")?;
writeln!(out, " Language TEXT,")?;
writeln!(out, " nFiles INTEGER,")?;
writeln!(out, " nBlank INTEGER,")?;
writeln!(out, " nComment INTEGER,")?;
writeln!(out, " nCode INTEGER")?;
writeln!(out, ");")?;
writeln!(out)?;
for lang in &summary.languages {
writeln!(
out,
"INSERT INTO t VALUES ('{}', {}, {}, {}, {});",
lang.name.replace('\'', "''"),
lang.files,
lang.blanks,
lang.comments,
lang.code
)?;
}
writeln!(
out,
"INSERT INTO t VALUES ('SUM', {}, {}, {}, {});",
summary.total_files,
summary.total_blanks,
summary.total_comments,
summary.total_code
)
}
OutputFormat::Xml => {
writeln!(out, r#"<?xml version="1.0" encoding="UTF-8"?>"#)?;
writeln!(out, "<results>")?;
if let Some(elapsed) = summary.elapsed {
writeln!(out, " <header>")?;
writeln!(out, " <n_files>{}</n_files>", summary.total_files)?;
writeln!(out, " <n_lines>{}</n_lines>", summary.total_lines())?;
writeln!(
out,
" <elapsed_seconds>{:.3}</elapsed_seconds>",
elapsed.as_secs_f64()
)?;
writeln!(out, " </header>")?;
}
writeln!(out, " <languages>")?;
for lang in &summary.languages {
let escaped_name = lang
.name
.replace('&', "&")
.replace('<', "<")
.replace('>', ">")
.replace('"', """);
writeln!(out, " <language name=\"{}\">", escaped_name)?;
writeln!(out, " <files>{}</files>", lang.files)?;
writeln!(out, " <blank>{}</blank>", lang.blanks)?;
writeln!(out, " <comment>{}</comment>", lang.comments)?;
writeln!(out, " <code>{}</code>", lang.code)?;
writeln!(out, " </language>")?;
}
writeln!(out, " </languages>")?;
writeln!(out, " <total>")?;
writeln!(out, " <files>{}</files>", summary.total_files)?;
writeln!(out, " <blank>{}</blank>", summary.total_blanks)?;
writeln!(out, " <comment>{}</comment>", summary.total_comments)?;
writeln!(out, " <code>{}</code>", summary.total_code)?;
writeln!(out, " </total>")?;
writeln!(out, "</results>")
}
}
}
fn sum_reports(cli: &Cli) -> Result<(), Box<dyn std::error::Error>> {
use rloc::stats::JsonOutput;
let mut reports = Vec::new();
for path in &cli.sum_reports {
let content = std::fs::read_to_string(path)
.map_err(|e| format!("Failed to read {}: {}", path.display(), e))?;
let report: JsonOutput = serde_json::from_str(&content)
.map_err(|e| format!("Failed to parse {}: {}", path.display(), e))?;
reports.push(report);
}
let combined = JsonOutput::sum_reports(reports);
let json = serde_json::to_string_pretty(&combined)?;
println!("{}", json);
Ok(())
}
fn run_strip(cli: &Cli) -> Result<(), Box<dyn std::error::Error>> {
let walker_config = cli.to_walker_config()?;
let files = rloc::walker::walk_files(&walker_config);
let (mode, ext) = if let Some(ref ext) = cli.strip_comments {
(StripMode::Comments, ext.as_str())
} else if let Some(ref ext) = cli.strip_code {
(StripMode::Code, ext.as_str())
} else {
return Err("No strip mode specified".into());
};
let mut processed = 0;
let mut errors = 0;
for entry in files {
match strip::strip_file(
&entry.path,
entry.language,
match mode {
StripMode::Comments => StripMode::Comments,
StripMode::Code => StripMode::Code,
},
ext,
) {
Ok(()) => {
processed += 1;
if cli.verbose > 0 {
eprintln!("Stripped: {}", entry.path.display());
}
}
Err(e) => {
errors += 1;
if cli.verbose > 0 {
eprintln!("Error stripping {}: {}", entry.path.display(), e);
}
}
}
}
if !cli.quiet {
eprintln!("Processed {} files ({} errors)", processed, errors);
}
Ok(())
}
fn run_diff(cli: &Cli, diff_path: &std::path::Path) -> Result<(), Box<dyn std::error::Error>> {
let config1 = cli.to_walker_config()?;
let mut config2 = config1.clone();
config2.paths = vec![diff_path.to_path_buf()];
let result = diff::compute_diff(&config1, &config2, cli.verbose > 0);
diff::render_diff(&result);
Ok(())
}