use crate::{
Block, Color, Environment, Icon, IconMode, Pretty, Style, WithStyle as _,
ext::WithPdeathsig as _, file,
};
use git2::{Branch, DescribeFormatOptions, DescribeOptions, ObjectType, Reference, Repository};
use rustix::process::Signal;
use std::{
borrow::Cow,
fs::File,
io::{BufRead as _, BufReader},
path::Path,
process::Command,
};
#[derive(Debug)]
enum Head {
Branch(String),
Detached(String),
Unknown,
}
impl Icon for Head {
fn icon(&self, mode: IconMode) -> &'static str {
use IconMode::*;
match self {
Self::Branch(_) => match mode {
Text => "on",
Icons | MinimalIcons => "",
},
Self::Detached(_) => match mode {
Text => "at",
Icons | MinimalIcons => "",
},
Self::Unknown => "<unk>",
}
}
}
impl Pretty for Head {
fn pretty(&self, f: &mut std::fmt::Formatter<'_>, mode: IconMode) -> std::fmt::Result {
let icon = self.icon(mode);
match &self {
Head::Branch(name) | Self::Detached(name) => write!(f, "{icon} {name}"),
Head::Unknown => write!(f, "{icon}"),
}
}
}
impl Head {
fn new(head: &Reference) -> Self {
if head.is_branch() {
Self::Branch(head.shorthand().unwrap_or_default().into())
} else if let Some(abbrev) = abbrev_commit(head) {
Head::Detached(abbrev)
} else {
Head::Unknown
}
}
fn ref_name(&self) -> Cow<'_, str> {
match self {
Self::Branch(name) => format!("refs/heads/{name}").into(),
Self::Detached(id) => id.into(),
Self::Unknown => "<head>".into(),
}
}
}
fn abbrev_commit(r: &Reference) -> Option<String> {
if let Ok(object) = r.peel(ObjectType::Commit)
&& let Ok(desc) = object.describe(
DescribeOptions::new()
.show_commit_oid_as_fallback(true)
.describe_all(),
)
&& let Ok(abbrev) = desc.format(Some(DescribeFormatOptions::new().abbreviated_size(4)))
{
return Some(abbrev);
}
None
}
enum State {
Merging { head: String },
Rebasing { done: usize, todo: usize },
Applying { next: usize, last: usize },
CherryPicking { head: String },
Reverting { head: String },
Bisecting,
}
impl State {
fn discover(repo: &Repository, root: &Path) -> Option<State> {
let rebase_merge = root.join("rebase-merge");
let rebase_apply = root.join("rebase-apply");
let abbrev_head = |head| {
std::fs::read_to_string(head)
.ok()
.and_then(|id| repo.find_reference(&id).ok())
.and_then(|r| abbrev_commit(&r))
};
Some(if file::exists(root.join("BISECT_LOG")) {
State::Bisecting
} else if let Some(head) = abbrev_head(root.join("REVERT_HEAD")) {
State::Reverting { head }
} else if let Some(head) = abbrev_head(root.join("CHERRY_PICK_HEAD")) {
State::CherryPicking { head }
} else if file::exists(&rebase_merge) {
let todo = match File::open(rebase_merge.join("git-rebase-todo")) {
Ok(file) => BufReader::new(file)
.lines()
.map_while(Result::ok)
.filter(|line| !line.starts_with('#'))
.count(),
Err(_) => 0,
};
let done = match File::open(rebase_merge.join("done")) {
Ok(file) => BufReader::new(file).lines().count(),
Err(_) => 0,
};
State::Rebasing { todo, done }
} else if file::exists(&rebase_apply) {
let read_single_usize = |filename| {
std::fs::read_to_string(rebase_apply.join(filename))
.ok()?
.trim()
.parse::<usize>()
.ok()
};
let next = read_single_usize("next").unwrap_or_default();
let last = read_single_usize("last").unwrap_or_default();
State::Applying { next, last }
} else if let Some(head) = abbrev_head(root.join("MERGE_HEAD")) {
State::Merging { head }
} else {
None?
})
}
}
impl Icon for State {
fn icon(&self, mode: IconMode) -> &'static str {
use IconMode::*;
match self {
Self::Bisecting => match mode {
Text => "bi",
Icons | MinimalIcons => " ",
},
Self::Reverting { .. } => match mode {
Text => "rv",
Icons | MinimalIcons => "",
},
Self::CherryPicking { .. } => match mode {
Text => "cp",
Icons | MinimalIcons => "",
},
Self::Merging { .. } => match mode {
Text => "me",
Icons | MinimalIcons => "",
},
Self::Rebasing { .. } => match mode {
Text => "rb",
Icons | MinimalIcons => "",
},
Self::Applying { .. } => match mode {
Text => "am",
Icons | MinimalIcons => "",
},
}
}
}
impl Pretty for State {
fn pretty(&self, f: &mut std::fmt::Formatter<'_>, mode: IconMode) -> std::fmt::Result {
let icon = self.icon(mode);
match self {
State::Bisecting => write!(f, "{icon}"),
State::CherryPicking { head } | State::Reverting { head } | State::Merging { head } => {
write!(f, "{icon} {head}")
}
State::Rebasing { done, todo } => write!(f, "{icon} {}/{}", done, done + todo),
State::Applying { next, last } => write!(f, "{icon} {next}/{last}"),
}
}
}
fn get_ahead_behind(remote: Option<&str>) -> Option<(usize, usize)> {
let output = Command::new("git")
.arg("rev-list")
.arg("--count")
.arg("--left-right")
.arg(format!("...{}", remote?))
.with_pdeathsig(Some(Signal::TERM))
.output()
.ok()?;
let mut iter = output
.stdout
.trim_ascii_end()
.split(|&c| c == b'\t')
.flat_map(std::str::from_utf8)
.flat_map(str::parse::<usize>);
let ahead = iter.next();
let behind = iter.next();
ahead.zip(behind)
}
pub struct GitRepo {
head: Head,
remote: Option<String>,
stashes: usize,
state: Option<State>,
behind: usize,
ahead: usize,
}
super::register_block!(GitRepo);
impl Block for GitRepo {
fn new(environ: &Environment) -> Option<Self> {
let tree = environ.git_tree.as_ref()?.clone();
let dotgit = tree.join(".git");
let root = if dotgit.is_file() {
tree.join(
std::fs::read_to_string(&dotgit)
.ok()?
.strip_prefix("gitdir: ")?
.trim_end_matches(['\r', '\n']),
)
} else {
dotgit
};
let mut repo = Repository::open_from_env().ok()?;
let mut stashes = 0;
repo.stash_foreach(|_, _, _| {
stashes += 1;
true
})
.ok()?;
let state = State::discover(&repo, &root);
let git_head = repo.head().ok()?;
let head = Head::new(&git_head);
let remote = Branch::wrap(git_head)
.upstream()
.ok()
.map(|br| br.into_reference())
.and_then(|rem| rem.name().ok().map(|s| s.into()));
Some(GitRepo {
head,
remote,
stashes,
state,
behind: 0,
ahead: 0,
})
}
fn extend(&mut self) {
(self.ahead, self.behind) = get_ahead_behind(self.remote.as_deref()).unwrap_or((0, 0));
}
}
pub struct GitTree {
unmerged: usize,
staged: usize,
dirty: usize,
untracked: usize,
}
super::register_block!(GitTree);
impl Block for GitTree {
fn new(environ: &Environment) -> Option<Self> {
environ.git_tree.as_ref().map(|_| GitTree {
unmerged: 0,
staged: 0,
dirty: 0,
untracked: 0,
})
}
fn extend(&mut self) {
let out = Command::new("git")
.arg("status")
.arg("--porcelain=2")
.with_pdeathsig(Some(Signal::TERM))
.output();
let Ok(out) = out else { return };
let lines = out.stdout.split(|&c| c == b'\n');
let mut unmerged = 0;
let mut staged = 0;
let mut dirty = 0;
let mut untracked = 0;
for line in lines {
let words: Vec<_> = line.split(|&c| c == b' ').take(2).collect();
if words.len() != 2 {
continue;
}
let (id, pat) = (words[0], words[1]);
match (id, pat) {
(b"?", _) => {
untracked += 1;
}
(b"u", _) => {
unmerged += 1;
}
(_, pat) if pat.len() == 2 => {
if pat[0] != b'.' {
staged += 1;
}
if pat[1] != b'.' {
dirty += 1;
}
}
_ => {}
}
}
self.unmerged = unmerged;
self.staged = staged;
self.dirty = dirty;
self.untracked = untracked;
}
}
impl Pretty for GitRepo {
fn pretty(&self, f: &mut std::fmt::Formatter<'_>, mode: IconMode) -> std::fmt::Result {
f.with_style(Color::of(&self.head.ref_name()), Style::BOLD, |f| {
write!(f, "[")?;
if let Some(state) = &self.state {
write!(f, "{}|", crate::icon::display(state, mode))?;
}
write!(f, "{}", crate::icon::display(&self.head, mode))?;
if let Head::Branch(local) = &self.head
&& let Some(remote_name) = &self.remote
&& remote_name
.strip_suffix(local)
.is_none_or(|remote| remote.as_bytes().last() != Some(&b'/'))
{
match remote_name.strip_prefix("refs/") {
Some(from_refs) => match from_refs.strip_prefix("remotes/") {
Some(from_remotes) => write!(f, ":{from_remotes}")?,
None => write!(f, ":{from_refs}")?,
},
None => write!(f, ":{remote_name}")?,
}
}
for (icon, val) in [
(GitIcon::Stashes, self.stashes),
(GitIcon::Behind, self.behind),
(GitIcon::Ahead, self.ahead),
] {
if val != 0 {
write!(f, " {}{}", icon.icon(mode), val)?;
}
}
write!(f, "]")
})
}
}
impl Pretty for GitTree {
fn pretty(&self, f: &mut std::fmt::Formatter<'_>, mode: IconMode) -> std::fmt::Result {
if self.unmerged == 0 && self.staged == 0 && self.dirty == 0 && self.untracked == 0 {
return Ok(());
}
f.with_style(Color::PINK, Style::empty(), |f| {
write!(f, "[")?;
let first = &mut true;
for (icon, val) in [
(GitIcon::Conflict, self.unmerged),
(GitIcon::Staged, self.staged),
(GitIcon::Dirty, self.dirty),
(GitIcon::Untracked, self.untracked),
] {
if val != 0 {
if !core::mem::take(first) {
write!(f, " ")?;
}
write!(f, "{}{}", icon.icon(mode), val)?;
}
}
write!(f, "]")
})
}
}
#[derive(PartialEq)]
enum GitIcon {
Ahead,
Behind,
Stashes,
Conflict,
Staged,
Dirty,
Untracked,
}
impl Icon for GitIcon {
fn icon(&self, mode: IconMode) -> &'static str {
use IconMode::*;
match &self {
Self::Ahead => match mode {
Text => "^",
Icons | MinimalIcons => " ",
},
Self::Behind => match mode {
Text => "v",
Icons | MinimalIcons => " ",
},
Self::Stashes => match mode {
Text => "*",
Icons | MinimalIcons => " ",
},
Self::Conflict => match mode {
Text => "=",
Icons => " ",
MinimalIcons => " ",
},
Self::Staged => match mode {
Text => "+",
Icons | MinimalIcons => " ",
},
Self::Dirty => match mode {
Text => "!",
Icons | MinimalIcons => " ",
},
Self::Untracked => match mode {
Text => "?",
Icons => " ",
MinimalIcons => " ",
},
}
}
}