git-heatmap/src/lib.rs

#![feature(let_chains)]

use std::{
	cmp::Reverse,
	fmt::Display,
	path::{self, PathBuf},
	sync::OnceLock,
};

use anyhow::anyhow;
use chrono::{DateTime, Duration, Local, NaiveDate, NaiveTime, TimeZone};
use clap::Parser;
use cli::CliArgs;
use gix::{
	ObjectId, ThreadSafeRepository, bstr::ByteSlice, revision::walk::Sorting,
	traverse::commit::simple::CommitTimeOrder,
};
use heatmap::{ColorLogic, HeatmapColors};
use itertools::Itertools;
use mailmap::Mailmap;
use rayon::prelude::*;
use rgb::Rgb;

pub mod cli;
pub mod heatmap;
pub mod mailmap;
pub mod rgb;

pub const ESCAPE: &str = "\x1B";
pub const RESET: &str = "\x1B[0m";
pub const DAYS: [&str; 7] = ["Mo", "Tu", "We", "Th", "Fr", "Sa", "Su"];
pub static CHAR: OnceLock<char> = OnceLock::new();
pub static COLOR_LOGIC: OnceLock<ColorLogic> = OnceLock::new();
pub static COLOR_MAP: OnceLock<Vec<String>> = OnceLock::new();

pub const GREEN_COLOR_MAP: [Rgb; 5] = [
	Rgb(0, 0, 0),
	Rgb(14, 68, 41),
	Rgb(0, 109, 50),
	Rgb(38, 166, 65),
	Rgb(25, 255, 64),
];

pub const RED_COLOR_MAP: [Rgb; 5] = [
	Rgb(0, 0, 0),
	Rgb(208, 169, 35),
	Rgb(208, 128, 35),
	Rgb(208, 78, 35),
	Rgb(255, 0, 0),
];

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Commit {
	pub id: ObjectId,
	pub title: String,
	pub author: Author,
	pub repo: String,
	pub time: DateTime<Local>,
}

impl Commit {
	pub fn new(
		id: ObjectId,
		title: String,
		author: String,
		email: String,
		repo: String,
		time: DateTime<Local>,
	) -> Self {
		Self {
			id,
			title,
			author: Author {
				name: author,
				email,
			},
			repo,
			time,
		}
	}
}

#[derive(Debug, PartialEq, Eq, Hash, Clone)]
pub struct Author {
	pub name: String,
	pub email: String,
}

impl Display for Author {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		f.write_fmt(format_args!("{} <{}>", self.name, self.email))
	}
}

pub fn args() -> CliArgs {
	let args = CliArgs::parse();

	CHAR.set(args.char).unwrap();
	COLOR_LOGIC.set(args.counting).unwrap();
	let color_map = match args.color_scheme {
		HeatmapColors::Green => GREEN_COLOR_MAP,
		HeatmapColors::Red => RED_COLOR_MAP,
	};
	let color_map = color_map
		.into_iter()
		.map(|c| c.to_ansi())
		.collect::<Vec<_>>();
	COLOR_MAP.set(color_map).unwrap();

	args
}

pub fn get_commits(
	args: CliArgs,
	start_date: NaiveDate,
	end_date: NaiveDate,
) -> anyhow::Result<(usize, usize, Vec<Commit>)> {
	let mut commits: Vec<Commit> = vec![];

	let ignored_repos = args.ignored_repos.as_ref().unwrap_or(&vec![]).to_owned();

	let (repos, branches) = match &args.root_dir {
		Some(roots) => {
			let mut repos: Vec<PathBuf> = vec![];
			for root in roots {
				find_git_repos(root, &mut repos, &ignored_repos);
			}
			let branches = vec!["".to_string(); repos.len()];
			(repos, branches)
		}
		None => {
			let mut repos = match args.repos {
				Some(r) => r,
				None => vec![PathBuf::from(".")],
			};

			// Turn any potential relative paths (such as `.`) into an absolute path
			repos = repos
				.iter_mut()
				.filter_map(|p| std::fs::canonicalize(p).ok())
				.collect_vec();

			let branches = args.branches.unwrap_or_else(|| vec!["".to_string()]);

			if repos.len() > 1 && repos.len() != branches.len() {
				return Err(anyhow!(
					"Number of repos ({}) needs to match the number of branch lists ({})!",
					repos.len(),
					branches.len()
				));
			}
			(repos, branches)
		}
	};

	let midnight = NaiveTime::from_hms_opt(0, 0, 0).unwrap();
	let start_date = start_date.and_time(midnight);
	let start_date = Local.from_local_datetime(&start_date).unwrap();

	let current_time = Local::now().time();
	let end_date = end_date.and_time(current_time);
	let end_date = Local.from_local_datetime(&end_date).unwrap();

	let authors = args.authors.unwrap_or_default();

	let mut repos_count: usize = 0;
	let mut branches_count: usize = 0;
	let mut cached_commits: Vec<Commit> = vec![];

	for (i, repo_path) in repos.iter().enumerate() {
		let repo = ThreadSafeRepository::open(repo_path).unwrap();

		let mut repo_name = repo_path.file_name();
		if repo_path.ends_with(gix::discover::DOT_GIT_DIR) {
			if let Some(parent) = repo_path.parent() {
				repo_name = parent.file_name();
			};
		}

		let repo_name = repo_name
			.iter()
			.filter_map(|r| r.to_str())
			.next()
			.unwrap_or("unknown");

		let branch_names = &*branches[i];
		let branches = get_repo_branches(&repo, branch_names).unwrap();

		let mailmap = Mailmap::new(repo_path);

		let branch_commits: Vec<_> = branches
			.par_iter()
			.filter_map(|branch| get_commit_ids(&repo, branch, start_date))
			.reduce(Vec::new, |mut c, n| {
				c.extend(n);
				c
			});

		let repo = repo.to_thread_local();

		let branch_commits = branch_commits
			.into_iter()
			.unique()
			.filter_map(|c| repo.find_commit(c).ok())
			.filter_map(|c| {
				let title = c
					.message()
					.ok()?
					.title
					.trim_ascii()
					.to_str()
					.ok()?
					.to_string();

				if args.no_merges {
					let is_merge = c.parent_ids().count() > 1;
					if is_merge {
						return None;
					}
				}

				let author = c.author().ok()?;

				// Ignores commits with different commit / author times
				// Usually due to rebases or cherry picking, however other edge cases may apply too
				if args.no_diff {
					let commit_info = c.committer().unwrap();
					if commit_info.time != author.time {
						return None;
					}
				}

				let email = author.email.to_string();
				let name = author.name.to_string();

				let time = if args.use_author_time {
					author.time().ok()?
				}
				else {
					c.time().ok()?
				};
				let time =
					DateTime::from_timestamp_millis(time.seconds * 1000)?.with_timezone(&Local);
				if time < start_date || time > end_date {
					return None;
				}

				let author = Author { name, email };
				let author = mailmap.resolve(author);

				if !authors.is_empty() && !authors.contains(&author.name) {
					return None;
				}

				let commit = Commit {
					id: c.id,
					title,
					author,
					repo: repo_name.to_string(),
					time,
				};

				if args.use_author_time {
					for other in &cached_commits {
						if other.author == commit.author
							&& other.title == commit.title
							&& other.time == commit.time
						{
							return None;
						}
					}

					cached_commits.push(commit.clone());
				}

				Some(commit)
			})
			.collect_vec();

		if !branch_commits.is_empty() {
			repos_count += 1;
			branches_count += branches.len();
		}

		commits.extend(branch_commits);
	}

	commits.par_sort_by_cached_key(|a| Reverse(a.time));

	Ok((repos_count, branches_count, commits))
}

fn get_repo_branches(repo: &ThreadSafeRepository, branch_names: &str) -> Option<Vec<String>> {
	if branch_names.is_empty() {
		let repo = repo.to_thread_local();
		let Ok(refs) = repo.references()
		else {
			return None;
		};

		let Ok(prefix) = refs.prefixed("refs/heads")
		else {
			return None;
		};

		let branches = prefix
			.filter_map(Result::ok)
			.filter_map(|b| {
				b.inner
					.name
					.to_string()
					.strip_prefix("refs/heads/")
					.map(|s| s.to_string())
			})
			.collect();

		Some(branches)
	}
	else {
		Some(branch_names.split(' ').map(|s| s.to_string()).collect())
	}
}

fn get_commit_ids(
	repo: &ThreadSafeRepository,
	branch: &str,
	start_date: DateTime<Local>,
) -> Option<Vec<ObjectId>> {
	let repo = repo.to_thread_local();

	// When passing the default @ (HEAD) branch this might actually not exist at all
	// locally so we're skipping it
	let rev = repo.rev_parse(branch).ok()?;

	let branch_commits = rev
		.single()
		.unwrap()
		.ancestors()
		.sorting(Sorting::ByCommitTimeCutoff {
			order: CommitTimeOrder::NewestFirst,
			seconds: start_date.timestamp(),
		})
		.all()
		.ok()?;

	let commits = branch_commits
		.filter_map(|c| c.ok())
		.map(|c| c.id)
		.collect();

	Some(commits)
}

fn find_git_repos(scan_path: &path::Path, repos: &mut Vec<PathBuf>, ignored_repos: &Vec<String>) {
	if let Some(path) = walk_dir(scan_path, ignored_repos) {
		repos.extend(path)
	}
}

pub fn walk_dir(scan_path: &path::Path, ignored_repos: &Vec<String>) -> Option<Vec<PathBuf>> {
	let Ok(dirs) = scan_path.read_dir()
	else {
		return None;
	};

	let dirs: Vec<PathBuf> = dirs
		.par_bridge()
		.filter_map(|d| d.ok())
		.filter(|d| {
			let dir_name = d.file_name().to_string_lossy().to_string();
			!ignored_repos.contains(&dir_name)
		})
		.filter(|d| d.file_type().is_ok_and(|t| t.is_dir()))
		.filter_map(|d| {
			let dir = d.path();
			let filename = dir.file_name().unwrap_or_default().to_string_lossy();

			match filename.as_ref() {
				".git" => Some(vec![dir]),
				_ => walk_dir(&dir, ignored_repos),
			}
		})
		.reduce(Vec::new, |mut c, n| {
			c.extend(n);
			c
		});

	Some(dirs)
}

pub fn get_default_until(since: NaiveDate) -> String {
	let mut until = Local::now().date_naive();
	if since + Duration::days(365) < until {
		until = since + Duration::days(365);
	}
	until.format("%Y-%m-%d").to_string()
}

fn get_color(val: i32, high: i32) -> usize {
	match COLOR_LOGIC.get() {
		Some(logic) => match logic {
			ColorLogic::ByAmount => match val {
				0 => 0,
				x if x < 2 => 1,
				x if x < 4 => 2,
				x if x < 6 => 3,
				x if x >= 6 => 4,
				_ => 0,
			},
			ColorLogic::ByWeight => {
				let color = val as f32 / high as f32;
				match color {
					0.0 => 0,
					x if x <= 0.2 => 1,
					x if x <= 0.4 => 2,
					x if x <= 0.8 => 3,
					x if x > 0.8 => 4,
					_ => 0,
				}
			}
		},
		None => 0,
	}
}

fn get_char() -> char {
	*CHAR.get_or_init(|| '▩')
}

fn get_color_map() -> Vec<String> {
	COLOR_MAP
		.get_or_init(|| {
			GREEN_COLOR_MAP
				.into_iter()
				.map(|c| c.to_ansi())
				.collect::<Vec<_>>()
		})
		.to_vec()
}