refactor(git): move RepoReferencesSnapshot functions to RepoExt

git-branchless-lib/benches/benches.rs

@@ -5,6 +5,7 @@ use branchless::core::dag::Dag;
 use branchless::core::effects::Effects;
 use branchless::core::eventlog::{EventLogDb, EventReplayer};
 use branchless::core::formatting::Glyphs;
+use branchless::core::repo_ext::RepoExt;
 use branchless::core::rewrite::{BuildRebasePlanOptions, RebasePlanBuilder, RepoResource};
 use branchless::git::{CherryPickFastOptions, Commit, Diff, Repo};
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};

git-branchless-lib/src/core/dag.rs

@@ -15,7 +15,9 @@ use tracing::{instrument, trace, warn};
 
 use crate::core::effects::{Effects, OperationType};
 use crate::core::eventlog::{CommitActivityStatus, EventCursor, EventReplayer};
-use crate::git::{Commit, MaybeZeroOid, NonZeroOid, Repo, RepoReferencesSnapshot};
+use crate::git::{Commit, MaybeZeroOid, NonZeroOid, Repo};
+
+use super::repo_ext::RepoReferencesSnapshot;
 
 impl From<NonZeroOid> for eden_dag::VertexName {
     fn from(oid: NonZeroOid) -> Self {

git-branchless-lib/src/core/eventlog.rs

@@ -17,9 +17,10 @@ use eyre::Context;
 use tracing::{error, instrument};
 
 use crate::core::effects::{Effects, OperationType};
-use crate::git::{
-    CategorizedReferenceName, MaybeZeroOid, NonZeroOid, Repo, RepoReferencesSnapshot,
-};
+use crate::core::repo_ext::RepoExt;
+use crate::git::{CategorizedReferenceName, MaybeZeroOid, NonZeroOid, Repo};
+
+use super::repo_ext::RepoReferencesSnapshot;
 
 /// When this environment variable is set, we reuse the ID for the transaction
 /// which the caller has already started.

git-branchless-lib/src/core/node_descriptors.rs

@@ -22,13 +22,11 @@ use crate::core::config::{
     get_commit_descriptors_branches, get_commit_descriptors_differential_revision,
     get_commit_descriptors_relative_time,
 };
-use crate::git::{
-    CategorizedReferenceName, Commit, NonZeroOid, Repo, RepoReferencesSnapshot,
-    ResolvedReferenceInfo,
-};
+use crate::git::{CategorizedReferenceName, Commit, NonZeroOid, Repo, ResolvedReferenceInfo};
 
 use super::eventlog::{Event, EventCursor, EventReplayer};
 use super::formatting::{Glyphs, StyledStringBuilder};
+use super::repo_ext::RepoReferencesSnapshot;
 use super::rewrite::find_rewrite_target;
 
 /// An object which can be rendered in the smartlog.

git-branchless-lib/src/core/repo_ext.rs

@@ -1,7 +1,134 @@
 //! Helper functions on [`Repo`].
 
-use crate::git::Repo;
+use std::collections::{HashMap, HashSet};
+use std::ffi::OsString;
 
-trait RepoExt {}
+use color_eyre::Help;
+use tracing::instrument;
 
-impl RepoExt for Repo {}
+use crate::git::{NonZeroOid, Reference, Repo};
+
+use super::config::get_main_branch_name;
+
+/// A snapshot of all the positions of references we care about in the repository.
+#[derive(Debug)]
+pub struct RepoReferencesSnapshot {
+    /// The location of the `HEAD` reference. This may be `None` if `HEAD` is unborn.
+    pub head_oid: Option<NonZeroOid>,
+
+    /// The location of the main branch.
+    pub main_branch_oid: NonZeroOid,
+
+    /// A mapping from commit OID to the branches which point to that commit.
+    pub branch_oid_to_names: HashMap<NonZeroOid, HashSet<OsString>>,
+}
+
+/// Helper functions on [`Repo`].
+pub trait RepoExt {
+    /// Get the `Reference` for the main branch for the repository.
+    fn get_main_branch_reference(&self) -> eyre::Result<Reference>;
+
+    /// Get the OID corresponding to the main branch.
+    fn get_main_branch_oid(&self) -> eyre::Result<NonZeroOid>;
+
+    /// Get a mapping from OID to the names of branches which point to that OID.
+    ///
+    /// The returned branch names include the `refs/heads/` prefix, so it must
+    /// be stripped if desired.
+    fn get_branch_oid_to_names(&self) -> eyre::Result<HashMap<NonZeroOid, HashSet<OsString>>>;
+
+    /// Get the positions of references in the repository.
+    fn get_references_snapshot(&self) -> eyre::Result<RepoReferencesSnapshot>;
+}
+
+impl RepoExt for Repo {
+    fn get_main_branch_reference(&self) -> eyre::Result<Reference> {
+        let main_branch_name = get_main_branch_name(self)?;
+        match self.find_branch(&main_branch_name, git2::BranchType::Local)? {
+            Some(branch) => match branch.get_upstream_branch()? {
+                Some(upstream_branch) => Ok(upstream_branch.into_reference()),
+                None => Ok(branch.into_reference()),
+            },
+            None => match self.find_branch(&main_branch_name, git2::BranchType::Remote)? {
+                Some(branch) => Ok(branch.into_reference()),
+                None => {
+                    let suggestion = format!(
+                        r"
+The main branch {:?} could not be found in your repository
+at path: {:?}.
+These branches exist: {:?}
+Either create it, or update the main branch setting by running:
+
+    git config branchless.core.mainBranch <branch>
+",
+                        get_main_branch_name(self)?,
+                        self.get_path(),
+                        self.get_all_local_branches()?
+                            .into_iter()
+                            .map(|branch| {
+                                branch
+                                    .into_reference()
+                                    .get_name()
+                                    .map(|s| format!("{:?}", s))
+                            })
+                            .collect::<eyre::Result<Vec<String>>>()?,
+                    );
+                    Err(eyre::eyre!("Could not find repository main branch")
+                        .with_suggestion(|| suggestion))
+                }
+            },
+        }
+    }
+
+    #[instrument]
+    fn get_main_branch_oid(&self) -> eyre::Result<NonZeroOid> {
+        let main_branch_reference = self.get_main_branch_reference()?;
+        let commit = main_branch_reference.peel_to_commit()?;
+        match commit {
+            Some(commit) => Ok(commit.get_oid()),
+            None => eyre::bail!(
+                "Could not find commit pointed to by main branch: {:?}",
+                main_branch_reference.get_name()?
+            ),
+        }
+    }
+
+    #[instrument]
+    fn get_branch_oid_to_names(&self) -> eyre::Result<HashMap<NonZeroOid, HashSet<OsString>>> {
+        let mut result: HashMap<NonZeroOid, HashSet<OsString>> = HashMap::new();
+        for branch in self.get_all_local_branches()? {
+            let reference = branch.into_reference();
+            let reference_name = reference.get_name()?;
+            let reference_info = self.resolve_reference(&reference)?;
+            if let Some(reference_oid) = reference_info.oid {
+                result
+                    .entry(reference_oid)
+                    .or_insert_with(HashSet::new)
+                    .insert(reference_name);
+            }
+        }
+
+        // The main branch may be a remote branch, in which case it won't be
+        // returned in the iteration above.
+        let main_branch_name = self.get_main_branch_reference()?.get_name()?;
+        let main_branch_oid = self.get_main_branch_oid()?;
+        result
+            .entry(main_branch_oid)
+            .or_insert_with(HashSet::new)
+            .insert(main_branch_name);
+
+        Ok(result)
+    }
+
+    fn get_references_snapshot(&self) -> eyre::Result<RepoReferencesSnapshot> {
+        let head_oid = self.get_head_info()?.oid;
+        let main_branch_oid = self.get_main_branch_oid()?;
+        let branch_oid_to_names = self.get_branch_oid_to_names()?;
+
+        Ok(RepoReferencesSnapshot {
+            head_oid,
+            main_branch_oid,
+            branch_oid_to_names,
+        })
+    }
+}

git-branchless-lib/src/core/rewrite/execute.rs

@@ -14,6 +14,7 @@ use crate::core::check_out::{check_out_commit, CheckOutCommitOptions};
 use crate::core::effects::Effects;
 use crate::core::eventlog::EventTransactionId;
 use crate::core::formatting::{printable_styled_string, Pluralize};
+use crate::core::repo_ext::RepoExt;
 use crate::git::{GitRunInfo, MaybeZeroOid, NonZeroOid, Repo, ResolvedReferenceInfo};
 
 use super::plan::RebasePlan;

git-branchless-lib/src/core/rewrite/plan.rs

@@ -945,6 +945,7 @@ mod tests {
 
     use crate::core::eventlog::{EventLogDb, EventReplayer};
     use crate::core::formatting::Glyphs;
+    use crate::core::repo_ext::RepoExt;
     use crate::testing::make_git;
 
     use super::*;

git-branchless-lib/src/core/rewrite/rewrite_hooks.rs

@@ -24,6 +24,7 @@ use crate::core::dag::Dag;
 use crate::core::effects::Effects;
 use crate::core::eventlog::{Event, EventLogDb, EventReplayer};
 use crate::core::formatting::{printable_styled_string, Pluralize};
+use crate::core::repo_ext::RepoExt;
 use crate::git::{
     CategorizedReferenceName, GitRunInfo, MaybeZeroOid, NonZeroOid, Repo, ResolvedReferenceInfo,
 };

git-branchless-lib/src/git/mod.rs

@@ -13,7 +13,7 @@ pub use oid::{MaybeZeroOid, NonZeroOid};
 pub use repo::{
     message_prettify, AmendFastOptions, Branch, CategorizedReferenceName, CherryPickFastError,
     CherryPickFastOptions, Commit, Diff, GitVersion, PatchId, Reference, ReferenceTarget, Repo,
-    RepoReferencesSnapshot, ResolvedReferenceInfo,
+    ResolvedReferenceInfo,
 };
 pub use run::GitRunInfo;
 pub use status::{FileStatus, StatusEntry};

git-branchless-lib/src/git/repo.rs

@@ -18,15 +18,13 @@ use std::str::FromStr;
 use std::time::SystemTime;
 
 use chrono::{DateTime, Local, TimeZone, Utc};
-use color_eyre::Help;
 use cursive::theme::BaseColor;
 use cursive::utils::markup::StyledString;
-use eyre::{eyre, Context};
+use eyre::Context;
 use itertools::Itertools;
 use os_str_bytes::{OsStrBytes, OsStringBytes};
 use tracing::{instrument, warn};
 
-use crate::core::config::get_main_branch_name;
 use crate::core::effects::{Effects, OperationType};
 use crate::core::eventlog::EventTransactionId;
 use crate::core::formatting::{Glyphs, StyledStringBuilder};
@@ -177,19 +175,6 @@ impl AmendFastOptions {
     }
 }
 
-/// A snapshot of all the positions of references we care about in the repository.
-#[derive(Debug)]
-pub struct RepoReferencesSnapshot {
-    /// The location of the `HEAD` reference. This may be `None` if `HEAD` is unborn.
-    pub head_oid: Option<NonZeroOid>,
-
-    /// The location of the main branch.
-    pub main_branch_oid: NonZeroOid,
-
-    /// A mapping from commit OID to the branches which point to that commit.
-    pub branch_oid_to_names: HashMap<NonZeroOid, HashSet<OsString>>,
-}
-
 /// Wrapper around `git2::Repository`.
 pub struct Repo {
     pub(super) inner: git2::Repository,
@@ -367,107 +352,6 @@ impl Repo {
         }
     }
 
-    /// Get the `Reference` for the main branch for the repository.
-    pub fn get_main_branch_reference(&self) -> eyre::Result<Reference> {
-        let main_branch_name = get_main_branch_name(self)?;
-        match self.find_branch(&main_branch_name, git2::BranchType::Local)? {
-            Some(branch) => {
-                let upstream_branch = branch
-                    .inner
-                    .upstream()
-                    .map(|branch| Branch { inner: branch })
-                    .unwrap_or_else(|_| branch);
-                Ok(upstream_branch.into_reference())
-            }
-            None => match self.find_branch(&main_branch_name, git2::BranchType::Remote)? {
-                Some(branch) => Ok(branch.into_reference()),
-                None => {
-                    let suggestion = format!(
-                        r"
-The main branch {:?} could not be found in your repository
-at path: {:?}.
-These branches exist: {:?}
-Either create it, or update the main branch setting by running:
-
-    git config branchless.core.mainBranch <branch>
-",
-                        get_main_branch_name(self)?,
-                        self.get_path(),
-                        self.get_all_local_branches()?
-                            .into_iter()
-                            .map(|branch| {
-                                branch
-                                    .into_reference()
-                                    .get_name()
-                                    .map(|s| format!("{:?}", s))
-                            })
-                            .collect::<eyre::Result<Vec<String>>>()?,
-                    );
-                    Err(eyre!("Could not find repository main branch")
-                        .with_suggestion(|| suggestion))
-                }
-            },
-        }
-    }
-
-    /// Get the OID corresponding to the main branch.
-    #[instrument]
-    pub fn get_main_branch_oid(&self) -> eyre::Result<NonZeroOid> {
-        let main_branch_reference = self.get_main_branch_reference()?;
-        let commit = main_branch_reference.peel_to_commit()?;
-        match commit {
-            Some(commit) => Ok(commit.get_oid()),
-            None => eyre::bail!(
-                "Could not find commit pointed to by main branch: {:?}",
-                main_branch_reference.get_name()?
-            ),
-        }
-    }
-
-    /// Get a mapping from OID to the names of branches which point to that OID.
-    ///
-    /// The returned branch names include the `refs/heads/` prefix, so it must
-    /// be stripped if desired.
-    #[instrument]
-    pub fn get_branch_oid_to_names(&self) -> eyre::Result<HashMap<NonZeroOid, HashSet<OsString>>> {
-        let mut result: HashMap<NonZeroOid, HashSet<OsString>> = HashMap::new();
-        for branch in self.get_all_local_branches()? {
-            let reference = branch.into_reference();
-            let reference_name = reference.get_name()?;
-            let reference_info = self.resolve_reference(&reference)?;
-            if let Some(reference_oid) = reference_info.oid {
-                result
-                    .entry(reference_oid)
-                    .or_insert_with(HashSet::new)
-                    .insert(reference_name);
-            }
-        }
-
-        // The main branch may be a remote branch, in which case it won't be
-        // returned in the iteration above.
-        let main_branch_name = self.get_main_branch_reference()?.get_name()?;
-        let main_branch_oid = self.get_main_branch_oid()?;
-        result
-            .entry(main_branch_oid)
-            .or_insert_with(HashSet::new)
-            .insert(main_branch_name);
-
-        Ok(result)
-    }
-
-    /// Get the positions of references in the repository.
-    pub fn get_references_snapshot(&self) -> eyre::Result<RepoReferencesSnapshot> {
-        let head_oid = self.get_head_info()?.oid;
-        let main_branch_oid = self.get_main_branch_oid()?;
-        let branch_oid_to_names = self.get_branch_oid_to_names()?;
-
-        Ok(RepoReferencesSnapshot {
-            head_oid,
-            main_branch_oid,
-            branch_oid_to_names,
-        })
-    }
-
     /// Detect if an interactive rebase has started but not completed.
     ///
     /// Git will send us spurious `post-rewrite` events marked as `amend` during an
@@ -1726,6 +1610,15 @@ impl<'repo> Branch<'repo> {
         Ok(self.inner.get().target().map(make_non_zero_oid))
     }
 
+    /// If this branch tracks a remote ("upstream") branch, return that branch.
+    pub fn get_upstream_branch(&self) -> eyre::Result<Option<Branch<'repo>>> {
+        match self.inner.upstream() {
+            Ok(upstream) => Ok(Some(Branch { inner: upstream })),
+            Err(err) if err.code() == git2::ErrorCode::NotFound => Ok(None),
+            Err(err) => Err(err.into()),
+        }
+    }
+
     /// Convert the branch into its underlying `Reference`.
     pub fn into_reference(self) -> Reference<'repo> {
         Reference {

git-branchless/src/commands/bug_report.rs

@@ -8,6 +8,7 @@ use bugreport::collector::{CollectionError, Collector};
 use bugreport::format::Markdown;
 use bugreport::report::ReportEntry;
 use itertools::Itertools;
+use lib::core::repo_ext::{RepoExt, RepoReferencesSnapshot};
 
 use crate::commands::smartlog::{make_smartlog_graph, render_graph};
 use lib::core::dag::Dag;
@@ -19,7 +20,7 @@ use lib::core::node_descriptors::{
     DifferentialRevisionDescriptor, ObsolescenceExplanationDescriptor, Redactor,
     RelativeTimeDescriptor,
 };
-use lib::git::{GitRunInfo, Repo, RepoReferencesSnapshot, ResolvedReferenceInfo};
+use lib::git::{GitRunInfo, Repo, ResolvedReferenceInfo};
 
 fn redact_event(redactor: &Redactor, event: &Event) -> String {
     let event = match event.clone() {