branchless/core/
repo_ext.rs

1//! Helper functions on [`Repo`].
2
3use std::collections::{HashMap, HashSet};
4
5use color_eyre::Help;
6use eyre::Context;
7use tracing::instrument;
8
9use crate::git::{
10    Branch, BranchType, CategorizedReferenceName, ConfigRead, NonZeroOid, ReferenceName, Repo,
11};
12
13use super::config::get_main_branch_name;
14
15/// A snapshot of all the positions of references we care about in the repository.
16#[derive(Debug)]
17pub struct RepoReferencesSnapshot {
18    /// The location of the `HEAD` reference. This may be `None` if `HEAD` is unborn.
19    pub head_oid: Option<NonZeroOid>,
20
21    /// The location of the main branch.
22    pub main_branch_oid: NonZeroOid,
23
24    /// A mapping from commit OID to the branches which point to that commit.
25    pub branch_oid_to_names: HashMap<NonZeroOid, HashSet<ReferenceName>>,
26}
27
28/// Helper functions on [`Repo`].
29pub trait RepoExt {
30    /// Get the `Branch` for the main branch for the repository.
31    fn get_main_branch(&self) -> eyre::Result<Branch>;
32
33    /// Get the OID corresponding to the main branch.
34    fn get_main_branch_oid(&self) -> eyre::Result<NonZeroOid>;
35
36    /// Get a mapping from OID to the names of branches which point to that OID.
37    ///
38    /// The returned branch names include the `refs/heads/` prefix, so it must
39    /// be stripped if desired.
40    fn get_branch_oid_to_names(&self) -> eyre::Result<HashMap<NonZeroOid, HashSet<ReferenceName>>>;
41
42    /// Get the positions of references in the repository.
43    fn get_references_snapshot(&self) -> eyre::Result<RepoReferencesSnapshot>;
44
45    /// Get the default remote to push to for new branches in this repository.
46    fn get_default_push_remote(&self) -> eyre::Result<Option<String>>;
47}
48
49impl RepoExt for Repo {
50    fn get_main_branch(&self) -> eyre::Result<Branch> {
51        let main_branch_name = get_main_branch_name(self)?;
52        match self.find_branch(&main_branch_name, BranchType::Local)? {
53            Some(branch) => Ok(branch),
54            None => {
55                let suggestion = format!(
56                    r"
57The main branch {:?} could not be found in your repository
58at path: {:?}.
59These branches exist: {:?}
60Either create it, or update the main branch setting by running:
61
62    git branchless init --main-branch <branch>
63
64Note that remote main branches are no longer supported as of v0.6.0. See
65https://github.com/arxanas/git-branchless/discussions/595 for more details.",
66                    get_main_branch_name(self)?,
67                    self.get_path(),
68                    self.get_all_local_branches()?
69                        .into_iter()
70                        .map(|branch| {
71                            branch
72                                .into_reference()
73                                .get_name()
74                                .map(|s| format!("{s:?}"))
75                                .wrap_err("converting branch to reference")
76                        })
77                        .collect::<eyre::Result<Vec<String>>>()?,
78                );
79                Err(eyre::eyre!("Could not find repository main branch")
80                    .with_suggestion(|| suggestion))
81            }
82        }
83    }
84
85    #[instrument]
86    fn get_main_branch_oid(&self) -> eyre::Result<NonZeroOid> {
87        let main_branch = self.get_main_branch()?;
88        let main_branch_oid = main_branch.get_oid()?;
89        match main_branch_oid {
90            Some(main_branch_oid) => Ok(main_branch_oid),
91            None => eyre::bail!(
92                "Could not find commit pointed to by main branch: {:?}",
93                main_branch.get_name()?,
94            ),
95        }
96    }
97
98    #[instrument]
99    fn get_branch_oid_to_names(&self) -> eyre::Result<HashMap<NonZeroOid, HashSet<ReferenceName>>> {
100        let mut result: HashMap<NonZeroOid, HashSet<ReferenceName>> = HashMap::new();
101        for branch in self.get_all_local_branches()? {
102            let reference = branch.into_reference();
103            let reference_name = reference.get_name()?;
104            let reference_info = self.resolve_reference(&reference)?;
105            if let Some(reference_oid) = reference_info.oid {
106                result
107                    .entry(reference_oid)
108                    .or_default()
109                    .insert(reference_name);
110            }
111        }
112
113        Ok(result)
114    }
115
116    fn get_references_snapshot(&self) -> eyre::Result<RepoReferencesSnapshot> {
117        let head_oid = self.get_head_info()?.oid;
118        let main_branch_oid = self.get_main_branch_oid()?;
119        let branch_oid_to_names = self.get_branch_oid_to_names()?;
120
121        Ok(RepoReferencesSnapshot {
122            head_oid,
123            main_branch_oid,
124            branch_oid_to_names,
125        })
126    }
127
128    fn get_default_push_remote(&self) -> eyre::Result<Option<String>> {
129        let main_branch_name = self.get_main_branch()?.get_reference_name()?;
130        match CategorizedReferenceName::new(&main_branch_name) {
131            name @ CategorizedReferenceName::LocalBranch { .. } => {
132                if let Some(main_branch) =
133                    self.find_branch(&name.render_suffix(), BranchType::Local)?
134                {
135                    if let Some(remote_name) = main_branch.get_push_remote_name()? {
136                        return Ok(Some(remote_name));
137                    }
138                }
139            }
140
141            name @ CategorizedReferenceName::RemoteBranch { .. } => {
142                let name = name.render_suffix();
143                if let Some((remote_name, _reference_name)) = name.split_once('/') {
144                    return Ok(Some(remote_name.to_owned()));
145                }
146            }
147
148            CategorizedReferenceName::OtherRef { .. } => {
149                // Do nothing.
150            }
151        }
152
153        let push_default_remote_opt = self.get_readonly_config()?.get("remote.pushDefault")?;
154        Ok(push_default_remote_opt)
155    }
156}