workspace: move recovery commit logic into lib for sharing

This is to facilitate automatic update-stale in extensions and in the CommandHelper layer.

cli/src/cli_util.rs

@@ -109,6 +109,7 @@ use jj_lib::signing::SignInitError;
 use jj_lib::str_util::StringPattern;
 use jj_lib::transaction::Transaction;
 use jj_lib::view::View;
+use jj_lib::working_copy;
 use jj_lib::working_copy::CheckoutStats;
 use jj_lib::working_copy::SnapshotOptions;
 use jj_lib::working_copy::WorkingCopy;
@@ -1033,6 +1034,37 @@ impl WorkspaceCommandHelper {
         Ok((locked_ws, wc_commit))
     }
 
+    pub fn create_and_check_out_recovery_commit(&mut self, ui: &Ui) -> Result<(), CommandError> {
+        self.check_working_copy_writable()?;
+
+        let workspace_id = self.workspace_id().clone();
+        let mut locked_ws = self.workspace.start_working_copy_mutation()?;
+        let (repo, new_commit) = working_copy::create_and_check_out_recovery_commit(
+            locked_ws.locked_wc(),
+            &self.user_repo.repo,
+            workspace_id,
+            self.env.settings(),
+            "RECOVERY COMMIT FROM `jj workspace update-stale`
+
+This commit contains changes that were written to the working copy by an
+operation that was subsequently lost (or was at least unavailable when you ran
+`jj workspace update-stale`). Because the operation was lost, we don't know
+what the parent commits are supposed to be. That means that the diff compared
+to the current parents may contain changes from multiple commits.
+",
+        )?;
+
+        writeln!(
+            ui.status(),
+            "Created and checked out recovery commit {}",
+            short_commit_hash(new_commit.id())
+        )?;
+        locked_ws.finish(repo.op_id().clone())?;
+
+        self.user_repo.repo = repo;
+        Ok(())
+    }
+
     pub fn workspace_root(&self) -> &Path {
         self.workspace.workspace_root()
     }

cli/src/command_error.rs

@@ -49,6 +49,7 @@ use jj_lib::revset::RevsetResolutionError;
 use jj_lib::signing::SignInitError;
 use jj_lib::str_util::StringPatternParseError;
 use jj_lib::view::RenameWorkspaceError;
+use jj_lib::working_copy::RecoverWorkspaceError;
 use jj_lib::working_copy::ResetError;
 use jj_lib::working_copy::SnapshotError;
 use jj_lib::working_copy::WorkingCopyStateError;
@@ -508,6 +509,18 @@ impl From<FilesetParseError> for CommandError {
     }
 }
 
+impl From<RecoverWorkspaceError> for CommandError {
+    fn from(err: RecoverWorkspaceError) -> Self {
+        match err {
+            RecoverWorkspaceError::Backend(err) => err.into(),
+            RecoverWorkspaceError::OpHeadsStore(err) => err.into(),
+            RecoverWorkspaceError::Reset(err) => err.into(),
+            RecoverWorkspaceError::RewriteRootCommit(err) => err.into(),
+            err @ RecoverWorkspaceError::WorkspaceMissingWorkingCopy(_) => user_error(err),
+        }
+    }
+}
+
 impl From<RevsetParseError> for CommandError {
     fn from(err: RevsetParseError) -> Self {
         let hint = revset_parse_error_hint(&err);

cli/src/commands/workspace/update_stale.rs

@@ -12,17 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::sync::Arc;
-
 use jj_lib::object_id::ObjectId;
 use jj_lib::op_store::OpStoreError;
-use jj_lib::repo::ReadonlyRepo;
 use jj_lib::repo::Repo;
 use jj_lib::working_copy::WorkingCopyFreshness;
 use tracing::instrument;
 
 use crate::cli_util::print_checkout_stats;
-use crate::cli_util::short_commit_hash;
 use crate::cli_util::CommandHelper;
 use crate::cli_util::WorkspaceCommandHelper;
 use crate::command_error::internal_error_with_message;
@@ -105,51 +101,6 @@ pub fn cmd_workspace_update_stale(
     Ok(())
 }
 
-fn create_and_check_out_recovery_commit(
-    ui: &mut Ui,
-    command: &CommandHelper,
-) -> Result<Arc<ReadonlyRepo>, CommandError> {
-    let mut workspace_command = command.workspace_helper_no_snapshot(ui)?;
-    let workspace_id = workspace_command.workspace_id().clone();
-    let mut tx = workspace_command.start_transaction().into_inner();
-
-    let (mut locked_workspace, commit) =
-        workspace_command.unchecked_start_working_copy_mutation()?;
-    let commit_id = commit.id();
-
-    let mut_repo = tx.repo_mut();
-    let new_commit = mut_repo
-        .new_commit(
-            command.settings(),
-            vec![commit_id.clone()],
-            commit.tree_id().clone(),
-        )
-        .set_description(
-            "RECOVERY COMMIT FROM `jj workspace update-stale`
-
-This commit contains changes that were written to the working copy by an
-operation that was subsequently lost (or was at least unavailable when you ran
-`jj workspace update-stale`). Because the operation was lost, we don't know
-what the parent commits are supposed to be. That means that the diff compared
-to the current parents may contain changes from multiple commits.
-",
-        )
-        .write()?;
-    mut_repo.set_wc_commit(workspace_id, new_commit.id().clone())?;
-    let repo = tx.commit("recovery commit")?;
-
-    locked_workspace.locked_wc().recover(&new_commit)?;
-    locked_workspace.finish(repo.op_id().clone())?;
-
-    writeln!(
-        ui.status(),
-        "Created and checked out recovery commit {}",
-        short_commit_hash(new_commit.id())
-    )?;
-
-    Ok(repo)
-}
-
 /// Loads workspace that will diverge from the last working-copy operation.
 fn for_stale_working_copy(
     ui: &mut Ui,
@@ -166,7 +117,10 @@ fn for_stale_working_copy(
                     "Failed to read working copy's current operation; attempting recovery. Error \
                      message from read attempt: {e}"
                 )?;
-                (create_and_check_out_recovery_commit(ui, command)?, true)
+
+                let mut workspace_command = command.workspace_helper_no_snapshot(ui)?;
+                workspace_command.create_and_check_out_recovery_commit(ui)?;
+                (workspace_command.repo().clone(), true)
             }
             Err(e) => return Err(e.into()),
         }

lib/src/working_copy.rs

@@ -33,15 +33,19 @@ use crate::gitignore::GitIgnoreError;
 use crate::gitignore::GitIgnoreFile;
 use crate::matchers::EverythingMatcher;
 use crate::matchers::Matcher;
+use crate::op_heads_store::OpHeadsStoreError;
 use crate::op_store::OpStoreError;
 use crate::op_store::OperationId;
 use crate::op_store::WorkspaceId;
 use crate::operation::Operation;
 use crate::repo::ReadonlyRepo;
+use crate::repo::Repo;
+use crate::repo::RewriteRootCommit;
 use crate::repo_path::InvalidRepoPathError;
 use crate::repo_path::RepoPath;
 use crate::repo_path::RepoPathBuf;
 use crate::settings::HumanByteSize;
+use crate::settings::UserSettings;
 use crate::store::Store;
 
 /// The trait all working-copy implementations must implement.
@@ -369,6 +373,58 @@ impl WorkingCopyFreshness {
     }
 }
 
+/// An error while recovering a stale working copy.
+#[derive(Debug, Error)]
+pub enum RecoverWorkspaceError {
+    /// Backend error.
+    #[error(transparent)]
+    Backend(#[from] BackendError),
+    /// Error during transaction.
+    #[error(transparent)]
+    OpHeadsStore(#[from] OpHeadsStoreError),
+    /// Error during checkout.
+    #[error(transparent)]
+    Reset(#[from] ResetError),
+    /// Checkout attempted to modify the root commit.
+    #[error(transparent)]
+    RewriteRootCommit(#[from] RewriteRootCommit),
+    /// Working copy commit is missing.
+    #[error("\"{0:?}\" doesn't have a working-copy commit")]
+    WorkspaceMissingWorkingCopy(WorkspaceId),
+}
+
+/// Recover this workspace to its last known checkout.
+pub fn create_and_check_out_recovery_commit(
+    locked_wc: &mut dyn LockedWorkingCopy,
+    repo: &Arc<ReadonlyRepo>,
+    workspace_id: WorkspaceId,
+    user_settings: &UserSettings,
+    description: &str,
+) -> Result<(Arc<ReadonlyRepo>, Commit), RecoverWorkspaceError> {
+    let mut tx = repo.start_transaction(user_settings);
+    let repo_mut = tx.repo_mut();
+
+    let commit_id = repo
+        .view()
+        .get_wc_commit_id(&workspace_id)
+        .ok_or_else(|| RecoverWorkspaceError::WorkspaceMissingWorkingCopy(workspace_id.clone()))?;
+    let commit = repo.store().get_commit(commit_id)?;
+    let new_commit = repo_mut
+        .new_commit(
+            user_settings,
+            vec![commit_id.clone()],
+            commit.tree_id().clone(),
+        )
+        .set_description(description)
+        .write()?;
+    repo_mut.set_wc_commit(workspace_id, new_commit.id().clone())?;
+
+    let repo = tx.commit("recovery commit")?;
+    locked_wc.recover(&new_commit)?;
+
+    Ok((repo, new_commit))
+}
+
 /// An error while reading the working copy state.
 #[derive(Debug, Error)]
 #[error("{message}")]