From 8231341fb1b65afe42acc337f112c87424a45eac Mon Sep 17 00:00:00 2001
From: Waleed Khan <me@waleedkhan.name>
Date: Tue, 4 Oct 2022 17:13:30 -0700
Subject: [PATCH] fix(init): support `init` in worktrees

---
 git-branchless-lib/src/git/repo.rs        | 29 ++++++++++++++++++++++
 git-branchless-lib/src/testing.rs         | 22 +++++++++++++++++
 git-branchless/src/commands/init.rs       |  4 ++-
 git-branchless/tests/command/test_init.rs | 30 ++++++++++++++++++++++-
 4 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/git-branchless-lib/src/git/repo.rs b/git-branchless-lib/src/git/repo.rs
index a0136e1b4..a9eb52775 100644
--- a/git-branchless-lib/src/git/repo.rs
+++ b/git-branchless-lib/src/git/repo.rs
@@ -52,6 +52,9 @@ pub enum Error {
     #[error("could not open repository: {0}")]
     OpenRepo(#[source] git2::Error),
 
+    #[error("could not find repository to open for worktree {path:?}")]
+    OpenParentWorktreeRepository { path: PathBuf },
+
     #[error("could not open repository: {0}")]
     UnsupportedExtensionWorktreeConfig(#[source] git2::Error),
 
@@ -520,6 +523,32 @@ impl Repo {
         Ok(Index { inner: index })
     }
 
+    /// If this repository is a worktree for another "parent" repository, return a [`Repo`] object
+    /// corresponding to that repository.
+    #[instrument]
+    pub fn open_worktree_parent_repo(&self) -> Result<Option<Self>> {
+        if !self.inner.is_worktree() {
+            return Ok(None);
+        }
+
+        // `git2` doesn't seem to support a way to directly look up the parent repository for the
+        // worktree.
+        let worktree_info_dir = self.get_path();
+        let parent_repo_path = match worktree_info_dir
+            .parent() // remove `.git`
+            .and_then(|path| path.parent()) // remove worktree name
+            .and_then(|path| path.parent()) // remove `worktrees`
+        {
+            Some(path) => path,
+            None => {
+                return Err(Error::OpenParentWorktreeRepository {
+                    path: worktree_info_dir.to_owned()});
+            },
+        };
+        let parent_repo = Self::from_dir(parent_repo_path)?;
+        Ok(Some(parent_repo))
+    }
+
     /// Get the configuration object for the repository.
     ///
     /// **Warning**: This object should only be used for read operations. Write
diff --git a/git-branchless-lib/src/testing.rs b/git-branchless-lib/src/testing.rs
index 709220adc..dfb09102a 100644
--- a/git-branchless-lib/src/testing.rs
+++ b/git-branchless-lib/src/testing.rs
@@ -604,3 +604,25 @@ pub fn make_git_with_remote_repo() -> eyre::Result<GitWrapperWithRemoteRepo> {
         cloned_repo,
     })
 }
+
+/// Represents a Git worktree for an existing Git repository on disk.
+pub struct GitWorktreeWrapper {
+    /// Guard to clean up the containing temporary directory. Make sure to bind
+    /// this to a local variable not named `_`.
+    pub temp_dir: TempDir,
+
+    /// A wrapper around the worktree.
+    pub worktree: Git,
+}
+
+/// Create a new worktree for the provided repository.
+pub fn make_git_worktree(git: &Git, worktree_name: &str) -> eyre::Result<GitWorktreeWrapper> {
+    let temp_dir = tempfile::tempdir()?;
+    let worktree_path = temp_dir.path().join(worktree_name);
+    git.run(&["worktree", "add", worktree_path.to_str().unwrap()])?;
+    let worktree = Git {
+        repo_path: worktree_path,
+        ..git.clone()
+    };
+    Ok(GitWorktreeWrapper { temp_dir, worktree })
+}
diff --git a/git-branchless/src/commands/init.rs b/git-branchless/src/commands/init.rs
index 6c91fa374..c4c43b50a 100644
--- a/git-branchless/src/commands/init.rs
+++ b/git-branchless/src/commands/init.rs
@@ -556,7 +556,9 @@ pub fn init(
     main_branch_name: Option<&str>,
 ) -> eyre::Result<()> {
     let mut in_ = BufReader::new(stdin());
-    let mut repo = Repo::from_current_dir()?;
+    let repo = Repo::from_current_dir()?;
+    let mut repo = repo.open_worktree_parent_repo()?.unwrap_or(repo);
+
     let default_config = Config::open_default()?;
     let readonly_config = repo.get_readonly_config()?;
     let mut config = create_isolated_config(effects, &repo, readonly_config.into_config())?;
diff --git a/git-branchless/tests/command/test_init.rs b/git-branchless/tests/command/test_init.rs
index a582c99c0..93df1ab54 100644
--- a/git-branchless/tests/command/test_init.rs
+++ b/git-branchless/tests/command/test_init.rs
@@ -6,7 +6,9 @@ use crate::util::trim_lines;
 
 use eyre::Context;
 use lib::git::GitVersion;
-use lib::testing::{make_git, GitInitOptions, GitRunOptions};
+use lib::testing::{
+    make_git, make_git_worktree, GitInitOptions, GitRunOptions, GitWorktreeWrapper,
+};
 
 #[test]
 fn test_hook_installed() -> eyre::Result<()> {
@@ -567,3 +569,29 @@ fn test_init_core_hooks_path_warning() -> eyre::Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_init_worktree() -> eyre::Result<()> {
+    let git = make_git()?;
+    git.init_repo_with_options(&GitInitOptions {
+        run_branchless_init: false,
+        make_initial_commit: true,
+    })?;
+    git.commit_file("test1", 1)?;
+    git.commit_file("test2", 2)?;
+
+    let GitWorktreeWrapper {
+        temp_dir: _temp_dir,
+        worktree,
+    } = make_git_worktree(&git, "new-worktree")?;
+    worktree.run(&["branchless", "init"])?;
+    {
+        let (stdout, _stderr) = worktree.run(&["smartlog"])?;
+        insta::assert_snapshot!(stdout, @r###"
+        :
+        @ 96d1c37 (> new-worktree, master) create test2.txt
+        "###);
+    }
+
+    Ok(())
+}