Vectrum
Sign In
494 lines · 19.3 KB · Rust
//! The git plumbing behind the project panel and the per-turn change
//! overview. Everything here blocks, so call it from a background task.

use std::{
    path::{Path, PathBuf},
    process::{Command, Stdio},
};

use serde::{Deserialize, Serialize};

/// Runs git in `path`, returning stdout, or stderr as the error.
pub fn git(path: &Path, args: &[&str]) -> Result<String, String> {
    run(Command::new("git").arg("-C").arg(path).args(args))
}

fn run(command: &mut Command) -> Result<String, String> {
    let output = command.stdin(Stdio::null()).output().map_err(|err| err.to_string())?;
    if output.status.success() {
        Ok(String::from_utf8_lossy(&output.stdout).into_owned())
    } else {
        Err(String::from_utf8_lossy(&output.stderr).trim().to_string())
    }
}

/// The repository a folder is in, named by its common directory (shared by
/// all of its worktrees), and the folder's path inside the checkout.
pub fn repository_of(path: &Path) -> Option<(PathBuf, PathBuf)> {
    let common = git(path, &["rev-parse", "--path-format=absolute", "--git-common-dir"]).ok()?;
    let top = git(path, &["rev-parse", "--show-toplevel"]).ok()?;
    let top = PathBuf::from(top.trim());
    let inside = path.strip_prefix(&top).map(Path::to_path_buf).unwrap_or_default();
    Some((PathBuf::from(common.trim()), inside))
}

/// What the project panel shows about a project's git state.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct GitInfo {
    pub branch: Option<String>,
    pub branches: Vec<String>,
    pub added: u64,
    pub removed: u64,
}

/// Reads the branch and the size of the working-tree changes, counting the
/// lines of untracked text files as added. `None` outside a repository.
pub fn read_info(path: &Path) -> Option<GitInfo> {
    git(path, &["rev-parse", "--is-inside-work-tree"]).ok()?;
    let branch = git(path, &["symbolic-ref", "--short", "-q", "HEAD"])
        .or_else(|_| git(path, &["rev-parse", "--short", "HEAD"]))
        .ok()
        .map(|branch| branch.trim().to_string())
        .filter(|branch| !branch.is_empty());
    let branches = git(path, &["for-each-ref", "--format=%(refname:short)", "refs/heads"])
        .map(|out| out.lines().map(str::to_string).collect())
        .unwrap_or_default();

    let mut numstat = String::new();
    if git(path, &["rev-parse", "--verify", "-q", "HEAD"]).is_ok() {
        numstat.push_str(&git(path, &["diff", "--numstat", "HEAD"]).unwrap_or_default());
    } else {
        // No commits yet: staged against nothing, plus unstaged.
        numstat.push_str(&git(path, &["diff", "--numstat", "--cached"]).unwrap_or_default());
        numstat.push_str(&git(path, &["diff", "--numstat"]).unwrap_or_default());
    }
    let (mut added, mut removed) = (0, 0);
    for line in numstat.lines() {
        let mut fields = line.split('\t');
        // Binary files report `-`.
        if let (Some(Ok(plus)), Some(Ok(minus))) =
            (fields.next().map(str::parse::<u64>), fields.next().map(str::parse::<u64>))
        {
            added += plus;
            removed += minus;
        }
    }
    let untracked = git(path, &["ls-files", "--others", "--exclude-standard", "-z"]).unwrap_or_default();
    for file in untracked.split('\0').filter(|file| !file.is_empty()).take(5_000) {
        let file = path.join(file);
        if file.metadata().is_ok_and(|meta| meta.len() <= 1_048_576)
            && let Ok(bytes) = std::fs::read(&file)
            && !bytes.contains(&0)
        {
            let lines = bytes.iter().filter(|&&byte| byte == b'\n').count() as u64;
            added += lines + u64::from(bytes.last().is_some_and(|&byte| byte != b'\n'));
        }
    }
    Some(GitInfo { branch, branches, added, removed })
}

/// A branch the branch picker lists.
#[derive(Clone, Debug, PartialEq)]
pub struct GitRef {
    /// `main`, or `origin/main` for a remote branch.
    pub name: String,
    pub remote: bool,
    /// Checked out in the folder the refs were read from.
    pub current: bool,
    /// The repository's default branch (origin's HEAD).
    pub default: bool,
    /// The worktree it's checked out in, when that's another folder.
    pub worktree: Option<PathBuf>,
}

impl GitRef {
    /// The badge the branch picker shows next to it.
    pub fn badge(&self) -> Option<&'static str> {
        if self.current {
            Some("current")
        } else if self.worktree.is_some() {
            Some("worktree")
        } else if self.remote {
            Some("remote")
        } else if self.default {
            Some("default")
        } else {
            None
        }
    }
}

/// The branches of the repository at `path`, ordered like T3 Code's branch
/// picker: the current one, the default one, then local before remote
/// branches, each newest commit first. A remote branch is left out when a
/// local one has its name.
pub fn list_refs(path: &Path) -> Vec<GitRef> {
    let Ok(out) = git(
        path,
        &[
            "for-each-ref",
            "--format=%(refname)%00%(committerdate:unix)%00%(HEAD)%00%(worktreepath)",
            "refs/heads",
            "refs/remotes",
        ],
    ) else {
        return Vec::new();
    };
    let here = git(path, &["rev-parse", "--show-toplevel"]).map(|top| PathBuf::from(top.trim())).ok();
    let default = git(path, &["symbolic-ref", "--short", "-q", "refs/remotes/origin/HEAD"])
        .ok()
        .map(|name| name.trim().trim_start_matches("origin/").to_string());
    let mut refs: Vec<(i64, GitRef)> = Vec::new();
    for line in out.lines() {
        let mut fields = line.split('\0');
        let (Some(refname), Some(date), Some(head), worktree) = (fields.next(), fields.next(), fields.next(), fields.next())
        else {
            continue;
        };
        let (name, remote) = match (refname.strip_prefix("refs/heads/"), refname.strip_prefix("refs/remotes/")) {
            (Some(name), _) => (name.to_string(), false),
            (None, Some(name)) if !name.ends_with("/HEAD") => (name.to_string(), true),
            _ => continue,
        };
        let worktree = worktree.filter(|path| !path.is_empty()).map(PathBuf::from).filter(|path| path.exists());
        let current = head == "*" || worktree.is_some() && worktree == here;
        let is_default = default.as_deref().is_some_and(|default| {
            if remote { name.strip_prefix("origin/") == Some(default) } else { name == default }
        });
        refs.push((
            date.parse().unwrap_or(0),
            GitRef { name, remote, current, default: is_default, worktree: worktree.filter(|_| !current) },
        ));
    }
    let locals: Vec<String> = refs.iter().filter(|(_, r)| !r.remote).map(|(_, r)| r.name.clone()).collect();
    refs.retain(|(_, r)| !r.remote || !r.name.strip_prefix("origin/").is_some_and(|name| locals.iter().any(|local| local == name)));
    refs.sort_by(|(a_date, a), (b_date, b)| {
        (!a.current, !a.default, a.remote, -a_date, &a.name).cmp(&(!b.current, !b.default, b.remote, -b_date, &b.name))
    });
    refs.into_iter().map(|(_, r)| r).collect()
}

/// Records the whole working tree, untracked files included, as a git tree
/// and returns its id. It stages into `index`, a private index file, so the
/// user's own staging area is never touched; the blobs land in the repository
/// as unreferenced objects that `git gc` cleans up.
pub fn snapshot(cwd: &Path, index: &Path) -> Option<String> {
    let git_dir = PathBuf::from(git(cwd, &["rev-parse", "--absolute-git-dir"]).ok()?.trim());
    if !index.exists() {
        // Seeding from the real index lets git skip rehashing unchanged files.
        let _ = std::fs::copy(git_dir.join("index"), index);
    }
    let with_index = |args: &[&str]| {
        run(Command::new("git").arg("-C").arg(cwd).args(args).env("GIT_INDEX_FILE", index))
    };
    with_index(&["add", "-A", "--", ":/"]).ok()?;
    let tree = with_index(&["write-tree"]).ok()?;
    Some(tree.trim().to_string())
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct FileChange {
    pub path: String,
    /// `None` for binary files.
    pub added: Option<u64>,
    pub removed: Option<u64>,
}

/// The files one turn changed: the difference between two snapshots.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct TurnChanges {
    pub base: String,
    pub head: String,
    pub files: Vec<FileChange>,
}

impl TurnChanges {
    pub fn added(&self) -> u64 {
        self.files.iter().filter_map(|file| file.added).sum()
    }

    pub fn removed(&self) -> u64 {
        self.files.iter().filter_map(|file| file.removed).sum()
    }
}

pub fn changes_between(cwd: &Path, base: &str, head: &str) -> Option<TurnChanges> {
    let out = git(cwd, &["diff-tree", "-r", "-z", "--numstat", "--no-renames", base, head]).ok()?;
    let mut files = Vec::new();
    for record in out.split('\0').filter(|record| !record.is_empty()) {
        let mut fields = record.splitn(3, '\t');
        let (Some(added), Some(removed), Some(path)) = (fields.next(), fields.next(), fields.next()) else {
            continue;
        };
        files.push(FileChange {
            path: path.to_string(),
            added: added.parse().ok(),
            removed: removed.parse().ok(),
        });
    }
    files.sort_by(|a, b| a.path.cmp(&b.path));
    Some(TurnChanges { base: base.into(), head: head.into(), files })
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LineKind {
    Hunk,
    Context,
    Added,
    Removed,
}

#[derive(Clone, Debug)]
pub struct DiffLine {
    pub kind: LineKind,
    pub text: String,
    /// Where the line sits in the old and the new file. Added lines have no
    /// old number, removed ones no new number, hunk rows neither.
    pub old_line: Option<u32>,
    pub new_line: Option<u32>,
}

#[derive(Clone, Debug)]
pub struct FileDiff {
    pub path: String,
    pub lines: Vec<DiffLine>,
    pub binary: bool,
}

/// The unified diff between two snapshots, optionally for one path.
/// `ignore_whitespace` leaves out whitespace-only edits, and with them the
/// files that changed nothing else.
pub fn diff_between(
    cwd: &Path,
    base: &str,
    head: &str,
    path: Option<&str>,
    ignore_whitespace: bool,
) -> Result<Vec<FileDiff>, String> {
    let mut args = vec!["diff", "--no-color", "--no-ext-diff", "--no-renames", "-U3"];
    if ignore_whitespace {
        args.push("--ignore-all-space");
    }
    args.extend([base, head]);
    if let Some(path) = path {
        args.extend(["--", path]);
    }
    let mut files = parse_diff(&git(cwd, &args)?);
    if ignore_whitespace {
        // Newer gits drop such a file; older ones keep its header, hunkless.
        files.retain(|file| file.binary || !file.lines.is_empty());
    }
    Ok(files)
}

/// The first old and new line of a hunk, from `@@ -a,b +c,d @@ context`.
fn parse_hunk(line: &str) -> Option<(u32, u32)> {
    let mut ranges = line.strip_prefix("@@ ")?.split(' ');
    let start = |range: Option<&str>, sign: char| range?.strip_prefix(sign)?.split(',').next()?.parse().ok();
    Some((start(ranges.next(), '-')?, start(ranges.next(), '+')?))
}

fn parse_diff(text: &str) -> Vec<FileDiff> {
    let mut files: Vec<FileDiff> = Vec::new();
    // The numbers the next old and new line will have.
    let (mut old, mut new) = (0, 0);
    for line in text.lines() {
        if let Some(rest) = line.strip_prefix("diff --git ") {
            // `a/path b/path`: take the new side.
            let path = rest.rsplit_once(" b/").map(|(_, path)| path).unwrap_or(rest);
            files.push(FileDiff { path: path.to_string(), lines: Vec::new(), binary: false });
            continue;
        }
        let Some(file) = files.last_mut() else { continue };
        let (kind, text, old_line, new_line) = if line.starts_with("@@") {
            (old, new) = parse_hunk(line).unwrap_or_default();
            (LineKind::Hunk, line, None, None)
        } else if line.starts_with("Binary files") {
            file.binary = true;
            continue;
        } else if file.lines.is_empty() && (line.starts_with("---") || line.starts_with("+++")) {
            continue;
        } else if let Some(text) = line.strip_prefix('+') {
            new += 1;
            (LineKind::Added, text, None, Some(new - 1))
        } else if let Some(text) = line.strip_prefix('-') {
            old += 1;
            (LineKind::Removed, text, Some(old - 1), None)
        } else if let Some(text) = line.strip_prefix(' ') {
            (old, new) = (old + 1, new + 1);
            (LineKind::Context, text, Some(old - 1), Some(new - 1))
        } else {
            // Headers: index, mode and new/deleted file lines.
            continue;
        };
        file.lines.push(DiffLine { kind, text: text.to_string(), old_line, new_line });
    }
    files
}

#[cfg(test)]
mod tests {
    use super::{LineKind, parse_diff, parse_hunk};

    #[::core::prelude::v1::test]
    fn parses_hunk_headers() {
        assert_eq!(parse_hunk("@@ -1,2 +3,4 @@ fn a() {"), Some((1, 3)));
        assert_eq!(parse_hunk("@@ -0,0 +1 @@"), Some((0, 1)));
        assert_eq!(parse_hunk("@@ -7 +7 @@"), Some((7, 7)));
        assert_eq!(parse_hunk("@@@ -1,2 -1,2 +1,3 @@@"), None);
    }

    #[::core::prelude::v1::test]
    fn numbers_lines() {
        let diff = "diff --git a/a b/a\n--- a/a\n+++ b/a\n@@ -10,4 +12,5 @@ fn a() {\n keep\n-gone\n+one\n+two\n keep\n@@ -40 +43 @@\n-x\n+y\n";
        let files = parse_diff(diff);
        let numbers: Vec<_> = files[0].lines.iter().map(|line| (line.old_line, line.new_line)).collect();
        assert_eq!(
            numbers,
            [
                (None, None),
                (Some(10), Some(12)),
                (Some(11), None),
                (None, Some(13)),
                (None, Some(14)),
                (Some(12), Some(15)),
                (None, None),
                (Some(40), None),
                (None, Some(43)),
            ]
        );
        assert_eq!(files[0].lines[0].text, "@@ -10,4 +12,5 @@ fn a() {");
    }

    #[::core::prelude::v1::test]
    fn parses_unified_diff() {
        let diff = "diff --git a/src/a.rs b/src/a.rs\nindex 1..2 100644\n--- a/src/a.rs\n+++ b/src/a.rs\n@@ -1,2 +1,2 @@\n fn a() {}\n-old\n+new\n";
        let files = parse_diff(diff);
        assert_eq!(files.len(), 1);
        assert_eq!(files[0].path, "src/a.rs");
        let kinds: Vec<_> = files[0].lines.iter().map(|line| line.kind).collect();
        assert_eq!(kinds, [LineKind::Hunk, LineKind::Context, LineKind::Removed, LineKind::Added]);
    }
}

/// One changed file, for the commit dialog.
#[derive(Clone, Debug, PartialEq)]
pub struct StatusFile {
    pub path: String,
    pub added: Option<u64>,
    pub removed: Option<u64>,
}

/// What T3 Code's git actions need to know about a checkout.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct GitStatus {
    pub branch: Option<String>,
    pub files: Vec<StatusFile>,
    /// The branch it tracks, like `origin/main`.
    pub upstream: Option<String>,
    pub ahead: u32,
    pub behind: u32,
    pub has_origin: bool,
    /// The repository's default branch, like `main`.
    pub default_branch: Option<String>,
    /// Commits on the branch that the default branch doesn't have.
    pub ahead_of_default: u32,
}

impl GitStatus {
    pub fn detached(&self) -> bool {
        self.branch.is_none()
    }

    pub fn has_changes(&self) -> bool {
        !self.files.is_empty()
    }

    pub fn on_default(&self) -> bool {
        self.branch.is_some() && self.branch == self.default_branch
    }
}

/// Reads a checkout's branch, changed files and how it stands against its
/// upstream and the default branch. `None` outside a repository.
pub fn status(path: &Path) -> Option<GitStatus> {
    git(path, &["rev-parse", "--is-inside-work-tree"]).ok()?;
    let branch = git(path, &["symbolic-ref", "--short", "-q", "HEAD"]).ok().map(|out| out.trim().to_string()).filter(|b| !b.is_empty());
    let mut files: Vec<StatusFile> = Vec::new();
    let porcelain = git(path, &["status", "--porcelain", "-z", "--untracked-files=all"]).unwrap_or_default();
    let mut entries = porcelain.split('\0').filter(|entry| !entry.is_empty());
    while let Some(entry) = entries.next() {
        if entry.len() < 4 {
            continue;
        }
        let code = &entry[..2];
        let file = entry[3..].to_string();
        // A rename lists its old path next.
        if code.starts_with('R') || code.starts_with('C') {
            entries.next();
        }
        files.push(StatusFile { path: file, added: None, removed: None });
    }
    let has_head = git(path, &["rev-parse", "--verify", "-q", "HEAD"]).is_ok();
    if has_head {
        for line in git(path, &["diff", "--numstat", "HEAD"]).unwrap_or_default().lines() {
            let mut fields = line.splitn(3, '\t');
            let (Some(plus), Some(minus), Some(file)) = (fields.next(), fields.next(), fields.next()) else { continue };
            if let Some(entry) = files.iter_mut().find(|entry| entry.path == file) {
                entry.added = plus.parse().ok();
                entry.removed = minus.parse().ok();
            }
        }
    }
    for entry in files.iter_mut().filter(|entry| entry.added.is_none()) {
        let full = path.join(&entry.path);
        if full.metadata().is_ok_and(|meta| meta.is_file() && meta.len() <= 1_048_576)
            && let Ok(bytes) = std::fs::read(&full)
            && !bytes.contains(&0)
        {
            entry.added = Some(bytes.iter().filter(|&&byte| byte == b'\n').count() as u64);
            entry.removed = Some(0);
        }
    }
    let upstream = git(path, &["rev-parse", "--abbrev-ref", "--symbolic-full-name", "@{upstream}"])
        .ok()
        .map(|out| out.trim().to_string())
        .filter(|upstream| !upstream.is_empty());
    let (mut ahead, mut behind) = (0, 0);
    if upstream.is_some()
        && let Ok(counts) = git(path, &["rev-list", "--left-right", "--count", "HEAD...@{upstream}"])
    {
        let mut counts = counts.split_whitespace().map(|count| count.parse().unwrap_or(0));
        ahead = counts.next().unwrap_or(0);
        behind = counts.next().unwrap_or(0);
    }
    let has_origin = git(path, &["remote"]).is_ok_and(|remotes| remotes.lines().any(|remote| remote == "origin"));
    let default_branch = git(path, &["symbolic-ref", "--short", "-q", "refs/remotes/origin/HEAD"])
        .ok()
        .map(|name| name.trim().trim_start_matches("origin/").to_string())
        .filter(|name| !name.is_empty())
        .or_else(|| {
            ["main", "master"]
                .into_iter()
                .find(|name| git(path, &["rev-parse", "--verify", "-q", &format!("refs/heads/{name}")]).is_ok())
                .map(str::to_string)
        });
    let ahead_of_default = match (&default_branch, has_head) {
        (Some(default), true) => {
            let base = if has_origin && git(path, &["rev-parse", "--verify", "-q", &format!("refs/remotes/origin/{default}")]).is_ok() {
                format!("origin/{default}")
            } else {
                default.clone()
            };
            git(path, &["rev-list", "--count", &format!("{base}..HEAD")]).ok().and_then(|out| out.trim().parse().ok()).unwrap_or(0)
        }
        _ => 0,
    };
    Some(GitStatus { branch, files, upstream, ahead, behind, has_origin, default_branch, ahead_of_default })
}

/// What's staged, as the commit message prompt reads it: the file list and
/// the patch.
pub fn staged(path: &Path) -> (String, String) {
    (git(path, &["diff", "--cached", "--stat"]).unwrap_or_default(), git(path, &["diff", "--cached"]).unwrap_or_default())
}