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> { 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::>>()?; 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 { 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 { 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::>(); 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 { 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::>(); 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: 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>> { 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(()) }