working_copy: add a reset() function for Git-like reset

We already have two usecases that can be modeled as updating the
`TreeState` without touching the working copy:

 1. `jj untrack` can be implemented as removing paths from the tree
    object and then doing a reset of the working copy state.

 2. Importing Git HEAD when sharing the working copy with a Git repo.

This patch adds that functionality to `TreeState`.

lib/src/working_copy.rs

@@ -41,7 +41,7 @@ use crate::lock::FileLock;
 use crate::matchers::{EverythingMatcher, Matcher};
 use crate::repo_path::{RepoPath, RepoPathComponent, RepoPathJoin};
 use crate::store::Store;
-use crate::tree::Diff;
+use crate::tree::{Diff, Tree};
 use crate::tree_builder::TreeBuilder;
 
 #[derive(Debug, PartialEq, Eq, Clone)]
@@ -189,6 +189,16 @@ pub enum CheckoutError {
     InternalBackendError(BackendError),
 }
 
+#[derive(Debug, Error, PartialEq, Eq)]
+pub enum ResetError {
+    // The current checkout was deleted, maybe by an overly aggressive GC that happened while
+    // the current process was running.
+    #[error("Current checkout not found")]
+    SourceNotFound,
+    #[error("Internal error: {0:?}")]
+    InternalBackendError(BackendError),
+}
+
 #[derive(Debug, Error, PartialEq, Eq)]
 pub enum UntrackError {
     // The current checkout was deleted, maybe by an overly aggressive GC that happened while
@@ -672,6 +682,46 @@ impl TreeState {
         Ok(stats)
     }
 
+    pub fn reset(&mut self, new_tree: &Tree) -> Result<(), ResetError> {
+        let old_tree = self
+            .store
+            .get_tree(&RepoPath::root(), &self.tree_id)
+            .map_err(|err| match err {
+                BackendError::NotFound => ResetError::SourceNotFound,
+                other => ResetError::InternalBackendError(other),
+            })?;
+
+        for (path, diff) in old_tree.diff(new_tree, &EverythingMatcher) {
+            match diff {
+                Diff::Removed(_before) => {
+                    self.file_states.remove(&path);
+                }
+                Diff::Added(after) | Diff::Modified(_, after) => {
+                    let file_type = match after {
+                        TreeValue::Normal { id: _, executable } => FileType::Normal { executable },
+                        TreeValue::Symlink(_id) => FileType::Symlink,
+                        TreeValue::Conflict(id) => FileType::Conflict { id },
+                        TreeValue::GitSubmodule(_id) => {
+                            println!("ignoring git submodule at {:?}", path);
+                            continue;
+                        }
+                        TreeValue::Tree(_id) => {
+                            panic!("unexpected tree entry in diff at {:?}", path);
+                        }
+                    };
+                    let file_state = FileState {
+                        file_type,
+                        mtime: MillisSinceEpoch(0),
+                        size: 0,
+                    };
+                    self.file_states.insert(path.clone(), file_state);
+                }
+            }
+        }
+        self.tree_id = new_tree.id().clone();
+        Ok(())
+    }
+
     pub fn untrack(&mut self, matcher: &dyn Matcher) -> Result<TreeId, UntrackError> {
         let tree = self
             .store
@@ -867,6 +917,10 @@ impl LockedWorkingCopy<'_> {
         self.wc.tree_state().as_mut().unwrap().write_tree()
     }
 
+    pub fn reset(&mut self, new_tree: &Tree) -> Result<(), ResetError> {
+        self.wc.tree_state().as_mut().unwrap().reset(new_tree)
+    }
+
     pub fn untrack(&mut self, matcher: &dyn Matcher) -> Result<TreeId, UntrackError> {
         self.wc.tree_state().as_mut().unwrap().untrack(matcher)
     }

lib/tests/test_working_copy.rs

@@ -291,6 +291,86 @@ fn test_checkout_file_transitions(use_git: bool) {
     }
 }
 
+#[test]
+fn test_reset() {
+    let settings = testutils::user_settings();
+    let mut test_workspace = testutils::init_repo(&settings, false);
+    let repo = &test_workspace.repo;
+    let workspace_root = test_workspace.workspace.workspace_root().clone();
+
+    let ignored_path = RepoPath::from_internal_string("ignored");
+    let gitignore_path = RepoPath::from_internal_string(".gitignore");
+
+    let tree_without_file = testutils::create_tree(repo, &[(&gitignore_path, "ignored\n")]);
+    let tree_with_file = testutils::create_tree(
+        repo,
+        &[(&gitignore_path, "ignored\n"), (&ignored_path, "code")],
+    );
+    let mut tx = repo.start_transaction("test");
+    let store = repo.store();
+    let root_commit = store.root_commit_id();
+    let commit_without_file = CommitBuilder::for_open_commit(
+        &settings,
+        store,
+        root_commit.clone(),
+        tree_without_file.id().clone(),
+    )
+    .write_to_repo(tx.mut_repo());
+    let commit_with_file = CommitBuilder::for_open_commit(
+        &settings,
+        store,
+        root_commit.clone(),
+        tree_with_file.id().clone(),
+    )
+    .write_to_repo(tx.mut_repo());
+    test_workspace.repo = tx.commit();
+
+    let wc = test_workspace.workspace.working_copy_mut();
+    wc.check_out(commit_with_file.clone()).unwrap();
+
+    // Test the setup: the file should exist on disk and in the tree state.
+    assert!(ignored_path.to_fs_path(&workspace_root).is_file());
+    assert!(wc.file_states().contains_key(&ignored_path));
+
+    // After we reset to the commit without the file, it should still exist on disk,
+    // but it should not be in the tree state, and it should not get added when we
+    // commit the working copy (because it's ignored).
+    let mut locked_wc = wc.start_mutation();
+    locked_wc.reset(&tree_without_file).unwrap();
+    locked_wc.finish(commit_without_file.id().clone());
+    assert!(ignored_path.to_fs_path(&workspace_root).is_file());
+    assert!(!wc.file_states().contains_key(&ignored_path));
+    let mut locked_wc = wc.start_mutation();
+    let new_tree_id = locked_wc.write_tree();
+    assert_eq!(new_tree_id, *tree_without_file.id());
+    locked_wc.discard();
+
+    // After we reset to the commit without the file, it should still exist on disk,
+    // but it should not be in the tree state, and it should not get added when we
+    // commit the working copy (because it's ignored).
+    let mut locked_wc = wc.start_mutation();
+    locked_wc.reset(&tree_without_file).unwrap();
+    locked_wc.finish(commit_without_file.id().clone());
+    assert!(ignored_path.to_fs_path(&workspace_root).is_file());
+    assert!(!wc.file_states().contains_key(&ignored_path));
+    let mut locked_wc = wc.start_mutation();
+    let new_tree_id = locked_wc.write_tree();
+    assert_eq!(new_tree_id, *tree_without_file.id());
+    locked_wc.discard();
+
+    // Now test the opposite direction: resetting to a commit where the file is
+    // tracked. The file should become tracked (even though it's ignored).
+    let mut locked_wc = wc.start_mutation();
+    locked_wc.reset(&tree_with_file).unwrap();
+    locked_wc.finish(commit_with_file.id().clone());
+    assert!(ignored_path.to_fs_path(&workspace_root).is_file());
+    assert!(wc.file_states().contains_key(&ignored_path));
+    let mut locked_wc = wc.start_mutation();
+    let new_tree_id = locked_wc.write_tree();
+    assert_eq!(new_tree_id, *tree_with_file.id());
+    locked_wc.discard();
+}
+
 #[test]
 fn test_untrack() {
     let settings = testutils::user_settings();