haze/src/git.rs

use crate::config::HazeConfig;
use crate::Result;
use camino::{Utf8Path, Utf8PathBuf};
use git2::build::CheckoutBuilder;
use git2::{Branch, BranchType, Repository, RepositoryState};
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use miette::{Context, IntoDiagnostic, Report};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use rayon::ThreadPoolBuilder;
use std::fs::read_dir;
use std::iter::once;
use std::path::PathBuf;
use std::process::Command;
use std::time::Duration;

fn find_app_repos(config: &HazeConfig) -> Result<impl Iterator<Item = PathBuf>> {
    let apps_dirs = once(config.sources_root.as_path().join("apps"))
        .chain(config.app_directories.iter().cloned());
    let dir_handles = apps_dirs
        .map(|dir| read_dir(dir).into_diagnostic())
        .collect::<Result<Vec<_>>>()?;
    Ok(dir_handles
        .into_iter()
        .flatten()
        .flatten()
        .filter(|app| app.path().join(".git").is_dir())
        .map(|app| app.path()))
}

fn longest_app_branch(config: &HazeConfig) -> Result<(usize, usize)> {
    Ok(find_app_repos(config)?
        .filter_map(|app_dir| {
            let app_name = app_dir.file_name()?.to_str()?;
            let repo = Repository::init(&app_dir).ok()?;
            let branch_name = current_branch_name(&repo)?;
            Some((app_name.len(), branch_name.len()))
        })
        .max()
        .unwrap_or_default())
}

pub fn checkout_all(config: &HazeConfig, mut name: &str, verbose: bool) -> Result<()> {
    // "main" and "master" are interchangeable
    if name == "main" {
        name = "master";
    }

    for app_dir in find_app_repos(config)? {
        let repo = Repository::init(&app_dir)
            .into_diagnostic()
            .wrap_err_with(|| format!("Failed to open repository {}", app_dir.display()))?;
        let app_name = app_dir.file_name().unwrap().to_string_lossy();
        if let Some(branch) = get_branch(&repo, name)? {
            if !branch.is_head() {
                let is_remote = branch.get().is_remote();

                print!("{app_name}");
                if let Err(e) = checkout(&repo, &branch, is_remote.then_some(name)) {
                    println!(": {:#} ❌", e);
                } else {
                    println!(" ✓");
                }
            } else if verbose {
                println!("{app_name} -");
            }
        } else if verbose {
            println!("{app_name} 🛇 Branch not found");
        };
    }
    Ok(())
}

pub fn find_app_repo(config: &HazeConfig, app: &str) -> Option<Utf8PathBuf> {
    find_app_repos(config)
        .ok()?
        .filter_map(|dir| Utf8PathBuf::try_from(dir).ok())
        .find(|dir| dir.file_name().is_some_and(|name| name == app))
}

fn git(repo: &Utf8Path, args: &[&str]) -> Result<String> {
    let output = Command::new(GIT_BINARY)
        .args(args)
        .current_dir(repo)
        .output()
        .into_diagnostic()
        .wrap_err_with(|| format!("Failed to run git {} for {}", args.join(" "), repo))?;
    if !output.status.success() {
        let err = String::from_utf8_lossy(&output.stderr);
        return Err(Report::msg(format!(
            "git {} failed in {}: {}",
            args.join(" "),
            repo,
            err.trim()
        )));
    }
    Ok(String::from_utf8_lossy(&output.stdout).trim().into())
}

fn has_ref(repo: &Utf8Path, git_ref: &str) -> bool {
    git(
        repo,
        &[
            "rev-parse",
            "--verify",
            "--quiet",
            &format!("{git_ref}^{{commit}}"),
        ],
    )
    .is_ok()
}

pub fn create_worktree(repo: &Utf8Path, dest: &Utf8Path, git_ref: &str) -> Result<()> {
    if !has_ref(repo, git_ref) {
        // The reference might be a tag that hasn't been fetched yet.
        let _ = git(repo, &["fetch", "origin", "tag", git_ref]);
        if !has_ref(repo, git_ref) {
            return Err(Report::msg(format!("{git_ref} not found in {repo}")));
        }
    }

    git(
        repo,
        &["worktree", "add", "--detach", dest.as_str(), git_ref], // Use a detached head so the worktree doesn't conflict with another one.
    )?;
    Ok(())
}

pub fn remove_worktree(repo: &Utf8Path) -> Result<()> {
    git(repo, &["worktree", "remove", "--force", repo.as_str()])?;
    Ok(())
}

pub fn init_submodules(worktree: &Utf8Path) -> Result<()> {
    git(worktree, &["submodule", "update", "--init", "--recursive"])?;
    Ok(())
}

pub fn prune_worktrees(config: &HazeConfig) -> Result<()> {
    let app_repos = find_app_repos(config)?
        .filter_map(|repo| Utf8PathBuf::try_from(repo).ok())
        .collect::<Vec<_>>();
    for repo in once(config.sources_root.clone()).chain(app_repos) {
        if let Err(e) = git(&repo, &["worktree", "prune"]) {
            eprintln!("Failed to prune worktrees of {repo}: {e:#}");
        }
    }
    Ok(())
}

fn current_branch_name(repo: &Repository) -> Option<String> {
    let (branch, _) = repo
        .branches(None)
        .ok()?
        .flatten()
        .find(|(b, _)| b.is_head())?;

    branch.name().ok().flatten().map(String::from)
}

const GIT_BINARY: &str = match option_env!("GIT_BINARY") {
    Some(git) => git,
    None => "git",
};

pub fn pull_all(config: &HazeConfig) -> Result<()> {
    let (max_app, max_branch) = longest_app_branch(config)?;

    let progress = MultiProgress::new();
    let pull_style = ProgressStyle::with_template("{spinner:.green} {msg}").unwrap();

    let pool = ThreadPoolBuilder::new()
        .num_threads(8)
        .build()
        .into_diagnostic()?;
    let repos = find_app_repos(config)?.collect::<Vec<_>>();

    pool.install(|| {
        repos.par_iter().for_each(|app_dir| {
            let app_name = app_dir.file_name().unwrap().to_string_lossy();
            let Ok(repo) = Repository::init(app_dir) else {
                return;
            };
            let branch_name = current_branch_name(&repo).unwrap_or("unknown".into());

            let bar = ProgressBar::new_spinner().with_style(pull_style.clone());
            bar.enable_steady_tick(Duration::from_millis(100));
            let bar = progress.add(bar);

            let msg = |state: &str| {
                format!(
                    "{app_name:<app_width$} - {branch_name:<branch_width$}{state}",
                    app_width = max_app,
                    branch_width = max_branch
                )
            };

            if repo.state() != RepositoryState::Clean {
                bar.set_message(msg(" ⨯ repository not clean"));
                bar.finish();
                return;
            }

            bar.set_message(msg(""));

            let output = match Command::new(GIT_BINARY)
                .arg("pull")
                .current_dir(app_dir)
                .output()
            {
                Ok(output) => output,
                Err(error) => {
                    bar.set_message(msg(&format!(" ⨯ {error}")));
                    return;
                }
            };

            if output.status.success() {
                bar.set_message(msg(" ✓"));
            } else {
                let err = String::from_utf8_lossy(&output.stderr);
                let err_line = err.lines().next().unwrap();
                bar.set_message(msg(&format!(" ⨯ {err_line}")));
            }
            bar.finish();
        });
    });

    Ok(())
}

fn get_branch<'repo>(repo: &'repo Repository, name: &str) -> Result<Option<Branch<'repo>>> {
    let local_branches = repo.branches(Some(BranchType::Local)).into_diagnostic()?;
    if let Some(local_branch) = local_branches.flatten().find_map(|(branch, _)| {
        match (branch.name(), name) {
            (Ok(Some("main")), "master") => Some(branch), // make "main" synonymous with "master"
            (Ok(Some(branch_name)), _) if branch_name == name => Some(branch),
            _ => None,
        }
    }) {
        return Ok(Some(local_branch));
    };

    let remote_branches = repo.branches(Some(BranchType::Remote)).into_diagnostic()?;
    let remote_name = format!("origin/{name}");
    Ok(remote_branches
        .flatten()
        .find_map(|(branch, _)| match branch.name() {
            Ok(Some(branch_name)) if branch_name == remote_name => Some(branch),
            _ => None,
        }))
}

fn checkout(repo: &Repository, branch: &Branch, target_name: Option<&str>) -> Result<()> {
    let commit = branch.get().peel_to_commit().into_diagnostic()?;
    let mut name = branch
        .name()
        .into_diagnostic()?
        .expect("we already know the name if utf8");
    let mut checkout = CheckoutBuilder::default();
    checkout.update_index(true);

    repo.checkout_tree(commit.as_object(), Some(&mut checkout))
        .into_diagnostic()
        .wrap_err("Failed to checkout tree")?;

    if let Some(target_name) = target_name {
        let mut new_branch = repo
            .branch(target_name, &commit, false)
            .into_diagnostic()
            .wrap_err("Failed to create local branch")?;
        new_branch
            .set_upstream(Some(name))
            .into_diagnostic()
            .wrap_err("Failed to set local branch upstream")?;
        name = target_name;
    }

    repo.set_head(&format!("refs/heads/{name}"))
        .into_diagnostic()
        .wrap_err("Failed to set HEAD")?;

    Ok(())
}