cli: add a function for restoring part of a tree from another tree

We had similar code in two places for restoring paths from one tree to
another. Let's reuse it instead.

I put the new function in the `rewrite` module. I'm not sure if that's
right place. Maybe it belongs in `tree`?

cli/src/cli_util.rs

@@ -41,8 +41,8 @@ use jj_lib::git_backend::GitBackend;
 use jj_lib::gitignore::GitIgnoreFile;
 use jj_lib::hex_util::to_reverse_hex;
 use jj_lib::id_prefix::IdPrefixContext;
-use jj_lib::matchers::{EverythingMatcher, Matcher, PrefixMatcher, Visit};
-use jj_lib::merged_tree::{MergedTree, MergedTreeBuilder};
+use jj_lib::matchers::{EverythingMatcher, Matcher, PrefixMatcher};
+use jj_lib::merged_tree::MergedTree;
 use jj_lib::op_heads_store::{self, OpHeadResolutionError, OpHeadsStore};
 use jj_lib::op_store::{OpStore, OpStoreError, OperationId, WorkspaceId};
 use jj_lib::operation::Operation;
@@ -56,6 +56,7 @@ use jj_lib::revset::{
     RevsetExpression, RevsetIteratorExt, RevsetParseContext, RevsetParseError,
     RevsetParseErrorKind, RevsetResolutionError, RevsetWorkspaceContext,
 };
+use jj_lib::rewrite::restore_tree;
 use jj_lib::settings::{ConfigResultExt as _, UserSettings};
 use jj_lib::str_util::{StringPattern, StringPatternParseError};
 use jj_lib::transaction::Transaction;
@@ -1570,16 +1571,9 @@ impl WorkspaceCommandTransaction<'_> {
     ) -> Result<MergedTreeId, CommandError> {
         if interactive {
             self.edit_diff(ui, left_tree, right_tree, matcher, instructions)
-        } else if matcher.visit(&RepoPath::root()) == Visit::AllRecursively {
-            // Optimization for a common case
-            Ok(right_tree.id().clone())
         } else {
-            let mut tree_builder = MergedTreeBuilder::new(left_tree.id().clone());
-            for (repo_path, diff) in left_tree.diff(right_tree, matcher) {
-                let (_left, right) = diff?;
-                tree_builder.set_or_remove(repo_path, right);
-            }
-            Ok(tree_builder.write_tree(self.repo().store())?)
+            let new_tree_id = restore_tree(right_tree, left_tree, matcher)?;
+            Ok(new_tree_id)
         }
     }
 

cli/src/commands/mod.rs

@@ -60,7 +60,7 @@ use jj_lib::merged_tree::{MergedTree, MergedTreeBuilder};
 use jj_lib::op_store::WorkspaceId;
 use jj_lib::repo::{ReadonlyRepo, Repo};
 use jj_lib::repo_path::RepoPath;
-use jj_lib::rewrite::{merge_commit_trees, DescendantRebaser};
+use jj_lib::rewrite::{merge_commit_trees, restore_tree, DescendantRebaser};
 use jj_lib::settings::UserSettings;
 use jj_lib::working_copy::SnapshotOptions;
 use jj_lib::workspace::{default_working_copy_initializer, Workspace};
@@ -1057,18 +1057,9 @@ fn cmd_restore(
     }
     workspace_command.check_rewritable([&to_commit])?;
 
-    let new_tree_id = if args.paths.is_empty() {
-        from_tree.id().clone()
-    } else {
-        let matcher = workspace_command.matcher_from_values(&args.paths)?;
-        let mut tree_builder = MergedTreeBuilder::new(to_commit.tree_id().clone());
-        let to_tree = to_commit.tree()?;
-        for (repo_path, diff) in from_tree.diff(&to_tree, matcher.as_ref()) {
-            let (before, _after) = diff?;
-            tree_builder.set_or_remove(repo_path, before);
-        }
-        tree_builder.write_tree(workspace_command.repo().store())?
-    };
+    let matcher = workspace_command.matcher_from_values(&args.paths)?;
+    let to_tree = to_commit.tree()?;
+    let new_tree_id = restore_tree(&from_tree, &to_tree, matcher.as_ref())?;
     if &new_tree_id == to_commit.tree_id() {
         writeln!(ui.stderr(), "Nothing changed.")?;
     } else {

lib/src/rewrite.rs

@@ -20,13 +20,15 @@ use std::sync::Arc;
 use itertools::Itertools;
 use tracing::instrument;
 
-use crate::backend::{BackendError, CommitId, ObjectId};
+use crate::backend::{BackendError, BackendResult, CommitId, MergedTreeId, ObjectId};
 use crate::commit::Commit;
 use crate::dag_walk;
 use crate::index::Index;
-use crate::merged_tree::MergedTree;
+use crate::matchers::{Matcher, Visit};
+use crate::merged_tree::{MergedTree, MergedTreeBuilder};
 use crate::op_store::RefTarget;
 use crate::repo::{MutableRepo, Repo};
+use crate::repo_path::RepoPath;
 use crate::revset::{RevsetExpression, RevsetIteratorExt};
 use crate::settings::UserSettings;
 use crate::store::Store;
@@ -68,6 +70,27 @@ pub fn merge_commit_trees_without_repo(
     }
 }
 
+/// Restore matching paths from the source into the destination.
+pub fn restore_tree(
+    source: &MergedTree,
+    destination: &MergedTree,
+    matcher: &dyn Matcher,
+) -> BackendResult<MergedTreeId> {
+    if matcher.visit(&RepoPath::root()) == Visit::AllRecursively {
+        // Optimization for a common case
+        Ok(source.id())
+    } else {
+        // TODO: We should be able to not traverse deeper in the diff if the matcher
+        // matches an entire subtree.
+        let mut tree_builder = MergedTreeBuilder::new(destination.id().clone());
+        for (repo_path, diff) in source.diff(destination, matcher) {
+            let (source_value, _destination_value) = diff?;
+            tree_builder.set_or_remove(repo_path, source_value);
+        }
+        tree_builder.write_tree(destination.store())
+    }
+}
+
 pub fn rebase_commit(
     settings: &UserSettings,
     mut_repo: &mut MutableRepo,

lib/tests/test_rewrite.rs

@@ -13,16 +13,59 @@
 // limitations under the License.
 
 use itertools::Itertools as _;
+use jj_lib::matchers::{EverythingMatcher, FilesMatcher};
 use jj_lib::op_store::{RefTarget, RemoteRef, RemoteRefState, WorkspaceId};
 use jj_lib::repo::Repo;
 use jj_lib::repo_path::RepoPath;
-use jj_lib::rewrite::DescendantRebaser;
+use jj_lib::rewrite::{restore_tree, DescendantRebaser};
 use maplit::{hashmap, hashset};
 use testutils::{
     assert_rebased, create_random_commit, create_tree, write_random_commit, CommitGraphBuilder,
     TestRepo,
 };
 
+#[test]
+fn test_restore_tree() {
+    let test_repo = TestRepo::init();
+    let repo = &test_repo.repo;
+
+    let path1 = RepoPath::from_internal_string("file1");
+    let path2 = RepoPath::from_internal_string("dir1/file2");
+    let path3 = RepoPath::from_internal_string("dir1/file3");
+    let path4 = RepoPath::from_internal_string("dir2/file4");
+    let left = create_tree(
+        repo,
+        &[(&path2, "left"), (&path3, "left"), (&path4, "left")],
+    );
+    let right = create_tree(
+        repo,
+        &[(&path1, "right"), (&path2, "right"), (&path3, "right")],
+    );
+
+    // Restore everything using EverythingMatcher
+    let restored = restore_tree(&left, &right, &EverythingMatcher).unwrap();
+    assert_eq!(restored, left.id());
+
+    // Restore everything using FilesMatcher
+    let restored = restore_tree(
+        &left,
+        &right,
+        &FilesMatcher::new(&[path1.clone(), path2.clone(), path3.clone(), path4.clone()]),
+    )
+    .unwrap();
+    assert_eq!(restored, left.id());
+
+    // Restore some files
+    let restored = restore_tree(
+        &left,
+        &right,
+        &FilesMatcher::new(&[path1.clone(), path2.clone()]),
+    )
+    .unwrap();
+    let expected = create_tree(repo, &[(&path2, "left"), (&path3, "right")]);
+    assert_eq!(restored, expected.id());
+}
+
 #[test]
 fn test_rebase_descendants_sideways() {
     let settings = testutils::user_settings();