From 6c17d9ad3057a64d48bd9bb1ee8a79aa7bfa96ea 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 | 37 ++++++++++++++++++++++--
 cli/tests/test_git_colocated.rs      | 42 +++++++++++++++++++++++++---
 2 files changed, 73 insertions(+), 6 deletions(-)

diff --git a/cli/src/commands/workspace/forget.rs b/cli/src/commands/workspace/forget.rs
index 07f2b165eb..fc3cc4b72d 100644
--- a/cli/src/commands/workspace/forget.rs
+++ b/cli/src/commands/workspace/forget.rs
@@ -18,6 +18,7 @@ use tracing::instrument;
 
 use crate::cli_util::CommandHelper;
 use crate::command_error::user_error;
+use crate::command_error::user_error_with_hint;
 use crate::command_error::CommandError;
 use crate::ui::Ui;
 
@@ -60,12 +61,39 @@ pub fn cmd_workspace_forget(
         }
     }
 
+    let git_backend = workspace_command
+        .git_backend()
+        .map(|backend| backend.git_repo());
+
     // 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) = git_backend.as_ref() {
+            tx.repo_mut().remove_git_head_target(ws);
+            match jj_lib::git::git_worktree_validate_removal(git_repo, ws.as_str()) {
+                Ok(stat) => worktrees_to_remove.push(stat),
+                Err(
+                    error @ jj_lib::git::WorktreeRemovalValidationError::NonexistentWorktree(_),
+                ) => {
+                    // Indistinguishable from a workspace that never had a worktree.
+                    tracing::debug!(%error, "Ignoring non-existent worktree");
+                }
+                Err(e) => {
+                    let err = format!("Could not remove Git worktree for workspace {ws}: {e}");
+                    return if let Some(hint) = e.hint(git_repo) {
+                        Err(user_error_with_hint(err, hint))
+                    } else {
+                        Err(user_error(err))
+                    };
+                }
+            }
+        }
+        tx.repo_mut()
+            .remove_wc_commit(ws)
+            .map_err(CommandError::from)
     })?;
 
     let description = if let [ws] = wss.as_slice() {
@@ -78,5 +106,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 a54fb2307a..850b5cb356 100644
--- a/cli/tests/test_git_colocated.rs
+++ b/cli/tests/test_git_colocated.rs
@@ -1107,13 +1107,47 @@ 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.
+    let _ = test_env.jj_cmd_ok(&repo_path, &["workspace", "add", "--colocate", "../second"]);
+}
+
+#[test]
+fn test_colocated_workspace_forget_locked() {
+    // TODO: Better way to disable the test if git command couldn't be executed
+    if Command::new("git").arg("--version").status().is_err() {
+        eprintln!("Skipping because git command might fail to run");
+        return;
+    }
+
+    let test_env = TestEnvironment::default();
+    let repo_path = test_env.env_root().join("repo");
+
+    let _ = test_env.jj_cmd_ok(test_env.env_root(), &["git", "init", "--colocate", "repo"]);
+    // There are no git heads if you do not commit
+    std::fs::write(repo_path.join("file"), "content").unwrap();
+    let _ = test_env.jj_cmd_ok(&repo_path, &["commit", "-m", "initial commit"]);
+    let _ = test_env.jj_cmd_ok(&repo_path, &["workspace", "add", "--colocate", "../second"]);
+
+    // Lock the new worktree
     Command::new("git")
-        .args(["worktree", "prune"])
+        .args(["worktree", "lock", "second"])
         .current_dir(&repo_path)
         .assert()
         .success();
 
-    let _ = test_env.jj_cmd_ok(&repo_path, &["workspace", "add", "--colocate", "../second"]);
+    insta::assert_snapshot!(
+        test_env.jj_cmd_failure(&repo_path, &["workspace", "forget", "second"]),
+        @r#"
+    Error: Could not remove Git worktree for workspace second: Worktree 'second' was locked, and could not be removed from git.
+    Hint: To unlock, run `git worktree unlock second`.
+    "#,
+    );
+
+    // Unlock it again
+    Command::new("git")
+        .args(["worktree", "unlock", "second"])
+        .current_dir(&repo_path)
+        .assert()
+        .success();
+
+    let _ = test_env.jj_cmd_ok(&repo_path, &["workspace", "forget", "second"]);
 }