cli_util: move snapshotting of the stale working copy into cli_util

This centralizes the logic so cli helpers can use it.

cli/src/cli_util.rs

@@ -283,6 +283,11 @@ struct CommandHelperData {
     working_copy_factories: WorkingCopyFactories,
 }
 
+pub enum StaleWorkingCopy {
+    Recovered(WorkspaceCommandHelper),
+    Snapshotted((Arc<ReadonlyRepo>, Commit)),
+}
+
 impl CommandHelper {
     pub fn app(&self) -> &Command {
         &self.data.app
@@ -413,6 +418,39 @@ impl CommandHelper {
             })
     }
 
+    pub fn load_stale_working_copy_commit(
+        &self,
+        ui: &Ui,
+    ) -> Result<StaleWorkingCopy, CommandError> {
+        let workspace = self.load_workspace()?;
+        let op_id = workspace.working_copy().operation_id();
+
+        match workspace.repo_loader().load_operation(op_id) {
+            Ok(op) => {
+                let repo = workspace.repo_loader().load_at(&op)?;
+                let mut workspace_command = self.for_workable_repo(ui, workspace, repo)?;
+                workspace_command.maybe_snapshot(ui)?;
+
+                let wc_commit_id = workspace_command.get_wc_commit_id().unwrap();
+                let repo = workspace_command.repo().clone();
+                let wc_commit = repo.store().get_commit(wc_commit_id)?;
+                Ok(StaleWorkingCopy::Snapshotted((repo, wc_commit)))
+            }
+            Err(e @ OpStoreError::ObjectNotFound { .. }) => {
+                writeln!(
+                    ui.status(),
+                    "Failed to read working copy's current operation; attempting recovery. Error \
+                     message from read attempt: {e}"
+                )?;
+
+                let mut workspace_command = self.workspace_helper_no_snapshot(ui)?;
+                workspace_command.create_and_check_out_recovery_commit(ui)?;
+                Ok(StaleWorkingCopy::Recovered(workspace_command))
+            }
+            Err(e) => Err(e.into()),
+        }
+    }
+
     /// Loads command environment for the given `workspace`.
     pub fn workspace_environment(
         &self,
@@ -1061,7 +1099,8 @@ to the current parents may contain changes from multiple commits.
         )?;
         locked_ws.finish(repo.op_id().clone())?;
 
-        self.user_repo.repo = repo;
+        self.user_repo = ReadonlyUserRepo::new(repo);
+        self.maybe_snapshot(ui)?;
         Ok(())
     }
 

cli/src/commands/workspace/update_stale.rs

@@ -12,14 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use jj_lib::op_store::OpStoreError;
-use jj_lib::repo::Repo;
 use jj_lib::working_copy::WorkingCopyFreshness;
 use tracing::instrument;
 
 use crate::cli_util::update_stale_working_copy;
 use crate::cli_util::CommandHelper;
-use crate::cli_util::WorkspaceCommandHelper;
+use crate::cli_util::StaleWorkingCopy;
 use crate::command_error::CommandError;
 use crate::ui::Ui;
 
@@ -40,19 +38,14 @@ pub fn cmd_workspace_update_stale(
     // operation, then merge the divergent operations. The wc_commit_id of the
     // merged repo wouldn't change because the old one wins, but it's probably
     // fine if we picked the new wc_commit_id.
-    let known_wc_commit = {
-        let (mut workspace_command, recovered) = for_stale_working_copy(ui, command)?;
-        workspace_command.maybe_snapshot(ui)?;
-
-        if recovered {
+    let known_wc_commit = match command.load_stale_working_copy_commit(ui)? {
+        StaleWorkingCopy::Recovered(_) => {
             // We have already recovered from the situation that prompted the user to run
             // this command, and it is known that the workspace is not stale
             // (since we just updated it), so we can return early.
             return Ok(());
         }
-
-        let wc_commit_id = workspace_command.get_wc_commit_id().unwrap();
-        workspace_command.repo().store().get_commit(wc_commit_id)?
+        StaleWorkingCopy::Snapshotted((_workspace_command, commit)) => commit,
     };
     let mut workspace_command = command.workspace_helper_no_snapshot(ui)?;
 
@@ -79,30 +72,3 @@ pub fn cmd_workspace_update_stale(
     }
     Ok(())
 }
-
-/// Loads workspace that will diverge from the last working-copy operation.
-fn for_stale_working_copy(
-    ui: &mut Ui,
-    command: &CommandHelper,
-) -> Result<(WorkspaceCommandHelper, bool), CommandError> {
-    let workspace = command.load_workspace()?;
-    let (repo, recovered) = {
-        let op_id = workspace.working_copy().operation_id();
-        match workspace.repo_loader().load_operation(op_id) {
-            Ok(op) => (workspace.repo_loader().load_at(&op)?, false),
-            Err(e @ OpStoreError::ObjectNotFound { .. }) => {
-                writeln!(
-                    ui.status(),
-                    "Failed to read working copy's current operation; attempting recovery. Error \
-                     message from read attempt: {e}"
-                )?;
-
-                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()),
-        }
-    };
-    Ok((command.for_workable_repo(ui, workspace, repo)?, recovered))
-}