From 4919138d241af247f7c619f130f7af449a7495f5 Mon Sep 17 00:00:00 2001
From: Cormac Relf <web@cormacrelf.net>
Date: Sun, 6 Oct 2024 01:04:51 +0900
Subject: [PATCH] git: Remove git worktrees when forgetting workspaces

---
 cli/src/commands/workspace/forget.rs | 30 ++++++++++++++++++++++++++--
 cli/tests/test_git_colocated.rs      |  8 --------
 lib/src/repo.rs                      |  5 +++--
 lib/src/view.rs                      |  5 +++--
 4 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/cli/src/commands/workspace/forget.rs b/cli/src/commands/workspace/forget.rs
index 8b5f416bd8..a41053d3c9 100644
--- a/cli/src/commands/workspace/forget.rs
+++ b/cli/src/commands/workspace/forget.rs
@@ -86,12 +86,33 @@ pub fn cmd_workspace_forget(
         }
     }
 
+    let git_repo_path = workspace_command
+        .git_backend()
+        .map(|git_backend| git_backend.git_repo_path().to_owned());
+
     // bundle every workspace forget into a single transaction, so that e.g.
     // undo correctly restores all of them at once.
     let mut tx = workspace_command.start_transaction();
+
+    let mut worktrees_to_remove = vec![];
     wss.iter().try_for_each(|ws| {
-        tx.repo_mut().remove_git_head_target(ws);
-        tx.repo_mut().remove_wc_commit(ws)
+        if let Some(git_repo_path) = git_repo_path.as_deref() {
+            tx.repo_mut().remove_git_head_target(ws);
+            match jj_lib::git::git_worktree_validate(git_repo_path, ws.as_str()) {
+                Ok(stat) => worktrees_to_remove.push(stat),
+                Err(jj_lib::git::WorktreeValidationError::NonexistentWorktree(_named)) => {
+                    tracing::debug!("No git worktree found named {_named}");
+                }
+                Err(e) => {
+                    return Err(user_error(format!(
+                        "Invalid worktree for workspace {ws}: {e}"
+                    )))
+                }
+            }
+        }
+        tx.repo_mut()
+            .remove_wc_commit(ws)
+            .map_err(CommandError::from)
     })?;
 
     let description = if let [ws] = wss.as_slice() {
@@ -104,5 +125,10 @@ pub fn cmd_workspace_forget(
     };
 
     tx.finish(ui, description)?;
+
+    for validated in worktrees_to_remove {
+        jj_lib::git::git_worktree_remove(validated).map_err(user_error)?;
+    }
+
     Ok(())
 }
diff --git a/cli/tests/test_git_colocated.rs b/cli/tests/test_git_colocated.rs
index c474bbfb05..b79cf6cfc8 100644
--- a/cli/tests/test_git_colocated.rs
+++ b/cli/tests/test_git_colocated.rs
@@ -1106,13 +1106,5 @@ fn test_colocated_workspace_add_forget_add() {
 
     std::fs::remove_dir_all(&second_path).unwrap();
 
-    // HACK: prune the git worktree manually. We should probably implement `git worktree remove`
-    // and do that when forgetting workspaces.
-    Command::new("git")
-        .args(["worktree", "prune"])
-        .current_dir(&repo_path)
-        .assert()
-        .success();
-
     let _ = test_env.jj_cmd_ok(&repo_path, &["workspace", "add", "--colocate", "../second"]);
 }
diff --git a/lib/src/repo.rs b/lib/src/repo.rs
index 943e0a0b19..ab2f168169 100644
--- a/lib/src/repo.rs
+++ b/lib/src/repo.rs
@@ -1346,8 +1346,9 @@ impl MutableRepo {
         Ok(())
     }
 
-    pub fn remove_git_head_target(&mut self, workspace_id: &WorkspaceId) {
-        self.view_mut().remove_git_head_target(workspace_id);
+    /// Returns true if did remove the HEAD
+    pub fn remove_git_head_target(&mut self, workspace_id: &WorkspaceId) -> bool {
+        self.view_mut().remove_git_head_target(workspace_id)
     }
 
     pub fn remove_wc_commit(&mut self, workspace_id: &WorkspaceId) -> Result<(), EditCommitError> {
diff --git a/lib/src/view.rs b/lib/src/view.rs
index 50e18bd0e2..797ea588e0 100644
--- a/lib/src/view.rs
+++ b/lib/src/view.rs
@@ -324,8 +324,9 @@ impl View {
         self.data.git_heads = git_heads;
     }
 
-    pub(crate) fn remove_git_head_target(&mut self, workspace_id: &WorkspaceId) {
-        self.data.git_heads.remove(workspace_id);
+    /// Returns true if did remove the HEAD
+    pub(crate) fn remove_git_head_target(&mut self, workspace_id: &WorkspaceId) -> bool {
+        self.data.git_heads.remove(workspace_id).is_some()
     }
 
     /// Iterates all commit ids referenced by this view.