diff --git a/lib/src/working_copy.rs b/lib/src/working_copy.rs
index 0e07dbbe1..df087bd1c 100644
--- a/lib/src/working_copy.rs
+++ b/lib/src/working_copy.rs
@@ -176,8 +176,6 @@ pub struct CheckoutStats {
 
 #[derive(Debug, Error, PartialEq, Eq)]
 pub enum CheckoutError {
-    #[error("Update target not found")]
-    TargetNotFound,
     // The current checkout was deleted, maybe by an overly aggressive GC that happened while
     // the current process was running.
     #[error("Current checkout not found")]
@@ -574,7 +572,7 @@ impl TreeState {
         }
     }
 
-    pub fn check_out(&mut self, tree_id: TreeId) -> Result<CheckoutStats, CheckoutError> {
+    pub fn check_out(&mut self, new_tree: &Tree) -> Result<CheckoutStats, CheckoutError> {
         let old_tree = self
             .store
             .get_tree(&RepoPath::root(), &self.tree_id)
@@ -582,13 +580,6 @@ impl TreeState {
                 BackendError::NotFound => CheckoutError::SourceNotFound,
                 other => CheckoutError::InternalBackendError(other),
             })?;
-        let new_tree =
-            self.store
-                .get_tree(&RepoPath::root(), &tree_id)
-                .map_err(|err| match err {
-                    BackendError::NotFound => CheckoutError::TargetNotFound,
-                    other => CheckoutError::InternalBackendError(other),
-                })?;
 
         let mut stats = CheckoutStats {
             updated_files: 0,
@@ -596,7 +587,7 @@ impl TreeState {
             removed_files: 0,
         };
 
-        for (path, diff) in old_tree.diff(&new_tree, &EverythingMatcher) {
+        for (path, diff) in old_tree.diff(new_tree, &EverythingMatcher) {
             let disk_path = path.to_fs_path(&self.working_copy_path);
 
             // TODO: Check that the file has not changed before overwriting/removing it.
@@ -668,7 +659,7 @@ impl TreeState {
                 }
             }
         }
-        self.tree_id = tree_id;
+        self.tree_id = new_tree.id().clone();
         Ok(stats)
     }
 
@@ -874,7 +865,7 @@ impl WorkingCopy {
                 return Err(CheckoutError::ConcurrentCheckout);
             }
         }
-        let stats = locked_wc.check_out(new_commit.tree_id().clone())?;
+        let stats = locked_wc.check_out(&new_commit.tree())?;
         locked_wc.finish(operation_id);
         Ok(stats)
     }
@@ -906,15 +897,10 @@ impl LockedWorkingCopy<'_> {
         self.wc.tree_state().as_mut().unwrap().write_tree()
     }
 
-    pub fn check_out(&mut self, new_tree_id: TreeId) -> Result<CheckoutStats, CheckoutError> {
+    pub fn check_out(&mut self, new_tree: &Tree) -> Result<CheckoutStats, CheckoutError> {
         // TODO: Write a "pending_checkout" file with the old and new TreeIds so we can
         // continue an interrupted checkout if we find such a file.
-        let stats = self
-            .wc
-            .tree_state()
-            .as_mut()
-            .unwrap()
-            .check_out(new_tree_id)?;
+        let stats = self.wc.tree_state().as_mut().unwrap().check_out(new_tree)?;
         Ok(stats)
     }
 
diff --git a/lib/tests/test_working_copy.rs b/lib/tests/test_working_copy.rs
index a213d0023..79cbffc5f 100644
--- a/lib/tests/test_working_copy.rs
+++ b/lib/tests/test_working_copy.rs
@@ -401,7 +401,7 @@ fn test_checkout_discard() {
 
     // Start a checkout
     let mut locked_wc = wc.start_mutation();
-    locked_wc.check_out(tree2.id().clone()).unwrap();
+    locked_wc.check_out(&tree2).unwrap();
     // The change should be reflected in the working copy but not saved
     assert!(!file1_path.to_fs_path(&workspace_root).is_file());
     assert!(file2_path.to_fs_path(&workspace_root).is_file());
diff --git a/src/commands.rs b/src/commands.rs
index 68bf475c7..a40e778fc 100644
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -535,9 +535,7 @@ impl WorkspaceCommandHelper {
                         wc_operation.id().hex(),
                         repo_operation.id().hex()
                     )?;
-                    locked_wc
-                        .check_out(checkout_commit.tree_id().clone())
-                        .unwrap();
+                    locked_wc.check_out(&checkout_commit.tree()).unwrap();
                 } else {
                     return Err(CommandError::InternalError(format!(
                         "The repo was loaded at operation {}, which seems to be a sibling of the \
diff --git a/src/diff_edit.rs b/src/diff_edit.rs
index ea957e44f..efc025149 100644
--- a/src/diff_edit.rs
+++ b/src/diff_edit.rs
@@ -55,12 +55,12 @@ fn check_out(
     store: Arc<Store>,
     wc_dir: PathBuf,
     state_dir: PathBuf,
-    tree_id: TreeId,
+    tree: &Tree,
 ) -> Result<TreeState, DiffEditError> {
     std::fs::create_dir(&wc_dir).unwrap();
     std::fs::create_dir(&state_dir).unwrap();
     let mut tree_state = TreeState::init(store, wc_dir, state_dir);
-    tree_state.check_out(tree_id)?;
+    tree_state.check_out(tree)?;
     Ok(tree_state)
 }
 
@@ -97,6 +97,7 @@ pub fn edit_diff(
     }
     let left_partial_tree_id = left_tree_builder.write_tree();
     let right_partial_tree_id = right_tree_builder.write_tree();
+    let left_partial_tree = store.get_tree(&RepoPath::root(), &left_partial_tree_id)?;
     let right_partial_tree = store.get_tree(&RepoPath::root(), &right_partial_tree_id)?;
 
     // Check out the two partial trees in temporary directories.
@@ -109,14 +110,14 @@ pub fn edit_diff(
         store.clone(),
         left_wc_dir.clone(),
         left_state_dir,
-        left_partial_tree_id,
+        &left_partial_tree,
     )?;
     set_readonly_recursively(&left_wc_dir);
     let mut right_tree_state = check_out(
         store.clone(),
         right_wc_dir.clone(),
         right_state_dir,
-        right_partial_tree_id,
+        &right_partial_tree,
     )?;
     let instructions_path = right_wc_dir.join("JJ-INSTRUCTIONS");
     // In the unlikely event that the file already exists, then the user will simply