store: cache tree on write and return it

This matches what we do when writing commits.

lib/src/rewrite.rs

@@ -54,8 +54,7 @@ pub fn merge_commit_trees_without_repo(
                 .map(|id| store.get_commit(id))
                 .try_collect()?;
             let ancestor_tree = merge_commit_trees_without_repo(store, index, &ancestors)?;
-            let new_tree_id = merge_trees(&new_tree, &ancestor_tree, &other_commit.tree())?;
+            new_tree = merge_trees(&new_tree, &ancestor_tree, &other_commit.tree())?;
-            new_tree = store.get_tree(&RepoPath::root(), &new_tree_id)?;
         }
         Ok(new_tree)
     }
@@ -83,7 +82,8 @@ pub fn rebase_commit(
         let old_base_tree = merge_commit_trees(mut_repo, &old_parents)?;
         let new_base_tree = merge_commit_trees(mut_repo, new_parents)?;
         // TODO: pass in labels for the merge parts
-        merge_trees(&new_base_tree, &old_base_tree, &old_commit.tree())?
+        let merged_tree = merge_trees(&new_base_tree, &old_base_tree, &old_commit.tree())?;
+        merged_tree.id().clone()
     };
     let new_parent_ids = new_parents
         .iter()
@@ -105,14 +105,14 @@ pub fn back_out_commit(
     let old_base_tree = merge_commit_trees(mut_repo, &old_commit.parents())?;
     let new_base_tree = merge_commit_trees(mut_repo, new_parents)?;
     // TODO: pass in labels for the merge parts
-    let new_tree_id = merge_trees(&new_base_tree, &old_commit.tree(), &old_base_tree).unwrap();
+    let new_tree = merge_trees(&new_base_tree, &old_commit.tree(), &old_base_tree).unwrap();
     let new_parent_ids = new_parents
         .iter()
         .map(|commit| commit.id().clone())
         .collect();
     // TODO: i18n the description based on repo language
     Ok(mut_repo
-        .new_commit(settings, new_parent_ids, new_tree_id)
+        .new_commit(settings, new_parent_ids, new_tree.id().clone())
         .set_description(format!("backout of commit {}", &old_commit.id().hex()))
         .write()?)
 }

lib/src/store.rs

@@ -122,9 +122,19 @@ impl Store {
         Ok(data)
     }
 
-    pub fn write_tree(&self, path: &RepoPath, contents: &backend::Tree) -> BackendResult<TreeId> {
+    pub fn write_tree(
-        // TODO: This should also do caching like write_commit does.
+        self: &Arc<Self>,
-        self.backend.write_tree(path, contents)
+        path: &RepoPath,
+        tree: backend::Tree,
+    ) -> BackendResult<Tree> {
+        let tree_id = self.backend.write_tree(path, &tree)?;
+        let data = Arc::new(tree);
+        {
+            let mut write_locked_cache = self.tree_cache.write().unwrap();
+            write_locked_cache.insert((path.clone(), tree_id.clone()), data.clone());
+        }
+
+        Ok(Tree::new(self.clone(), path.clone(), tree_id, data))
     }
 
     pub fn read_file(&self, path: &RepoPath, id: &FileId) -> BackendResult<Box<dyn Read>> {

lib/src/tree.rs

@@ -14,6 +14,7 @@
 
 use std::cmp::Ordering;
 use std::fmt::{Debug, Error, Formatter};
+use std::hash::{Hash, Hasher};
 use std::io::Read;
 use std::iter::Peekable;
 use std::pin::Pin;
@@ -62,6 +63,21 @@ impl Debug for Tree {
     }
 }
 
+impl PartialEq for Tree {
+    fn eq(&self, other: &Self) -> bool {
+        self.id == other.id && self.dir == other.dir
+    }
+}
+
+impl Eq for Tree {}
+
+impl Hash for Tree {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.dir.hash(state);
+        self.id.hash(state);
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct DiffSummary {
     pub modified: Vec<RepoPath>,
@@ -525,13 +541,13 @@ pub fn merge_trees(
     side1_tree: &Tree,
     base_tree: &Tree,
     side2_tree: &Tree,
-) -> Result<TreeId, TreeMergeError> {
+) -> Result<Tree, TreeMergeError> {
     let store = base_tree.store();
     let dir = base_tree.dir();
     assert_eq!(side1_tree.dir(), dir);
     assert_eq!(side2_tree.dir(), dir);
 
-    if let Some(resolved) = trivial_merge(&[&base_tree.id], &[&side1_tree.id, &side2_tree.id]) {
+    if let Some(resolved) = trivial_merge(&[base_tree], &[side1_tree, side2_tree]) {
         return Ok((*resolved).clone());
     }
 
@@ -559,7 +575,7 @@ pub fn merge_trees(
             }
         }
     }
-    Ok(store.write_tree(dir, &new_tree)?)
+    Ok(store.write_tree(dir, new_tree)?)
 }
 
 /// Returns `Some(TreeId)` if this is a directory or missing. If it's missing,
@@ -599,11 +615,11 @@ fn merge_tree_value(
             let base_tree = store.get_tree(&subdir, base_id)?;
             let side1_tree = store.get_tree(&subdir, side1_id)?;
             let side2_tree = store.get_tree(&subdir, side2_id)?;
-            let merged_tree_id = merge_trees(&side1_tree, &base_tree, &side2_tree)?;
+            let merged_tree = merge_trees(&side1_tree, &base_tree, &side2_tree)?;
-            if merged_tree_id == *empty_tree_id {
+            if merged_tree.id() == empty_tree_id {
                 None
             } else {
-                Some(TreeValue::Tree(merged_tree_id))
+                Some(TreeValue::Tree(merged_tree.id().clone()))
             }
         }
         _ => {

lib/src/tree_builder.rs

@@ -87,7 +87,7 @@ impl TreeBuilder {
 
         // Write trees in reverse lexicographical order, starting with trees without
         // children.
-        let store = self.store.as_ref();
+        let store = &self.store;
         // TODO: trees_to_write.pop_last() can be used, but requires Rust 1.66.0
         let mut dirs_to_write = trees_to_write.keys().cloned().collect_vec();
         while let Some(dir) = dirs_to_write.pop() {
@@ -101,13 +101,13 @@ impl TreeBuilder {
                         // Entry would have been replaced with file (see above)
                     }
                 } else {
-                    let tree_id = store.write_tree(&dir, &tree).unwrap();
+                    let tree = store.write_tree(&dir, tree).unwrap();
-                    parent_tree.set(basename.clone(), TreeValue::Tree(tree_id));
+                    parent_tree.set(basename.clone(), TreeValue::Tree(tree.id().clone()));
                 }
             } else {
                 // We're writing the root tree. Write it even if empty. Return its id.
                 assert!(dirs_to_write.is_empty());
-                return store.write_tree(&dir, &tree).unwrap();
+                return store.write_tree(&dir, tree).unwrap().id().clone();
             }
         }
 

lib/tests/test_merge_trees.rs

@@ -18,8 +18,7 @@ use jujutsu_lib::backend::TreeValue;
 use jujutsu_lib::repo::Repo;
 use jujutsu_lib::repo_path::{RepoPath, RepoPathComponent};
 use jujutsu_lib::rewrite::rebase_commit;
-use jujutsu_lib::tree;
+use jujutsu_lib::tree::{merge_trees, Tree};
-use jujutsu_lib::tree::Tree;
 use test_case::test_case;
 use testutils::TestRepo;
 
@@ -73,8 +72,7 @@ fn test_same_type(use_git: bool) {
     let side2_tree = write_tree(2);
 
     // Create the merged tree
-    let merged_tree_id = tree::merge_trees(&side1_tree, &base_tree, &side2_tree).unwrap();
+    let merged_tree = merge_trees(&side1_tree, &base_tree, &side2_tree).unwrap();
-    let merged_tree = store.get_tree(&RepoPath::root(), &merged_tree_id).unwrap();
 
     // Check that we have exactly the paths we expect in the merged tree
     let names = merged_tree
@@ -227,8 +225,7 @@ fn test_executable(use_git: bool) {
     let side2_tree = write_tree(&contents_in_tree(&files, 2));
 
     // Create the merged tree
-    let merged_tree_id = tree::merge_trees(&side1_tree, &base_tree, &side2_tree).unwrap();
+    let merged_tree = merge_trees(&side1_tree, &base_tree, &side2_tree).unwrap();
-    let merged_tree = store.get_tree(&RepoPath::root(), &merged_tree_id).unwrap();
 
     // Check that the merged tree has the correct executable bits
     let norm = base_tree.value(&RepoPathComponent::from("nnn"));
@@ -282,8 +279,7 @@ fn test_subtrees(use_git: bool) {
         "d1/d1/d1/f2",
     ]);
 
-    let merged_tree_id = tree::merge_trees(&side1_tree, &base_tree, &side2_tree).unwrap();
+    let merged_tree = merge_trees(&side1_tree, &base_tree, &side2_tree).unwrap();
-    let merged_tree = store.get_tree(&RepoPath::root(), &merged_tree_id).unwrap();
     let entries = merged_tree.entries().collect_vec();
 
     let expected_tree = write_tree(vec![
@@ -325,8 +321,7 @@ fn test_subtree_becomes_empty(use_git: bool) {
     let side1_tree = write_tree(vec!["f1", "d1/f1", "d1/d1/d1/f1"]);
     let side2_tree = write_tree(vec!["d1/d1/d1/f2"]);
 
-    let merged_tree_id = tree::merge_trees(&side1_tree, &base_tree, &side2_tree).unwrap();
+    let merged_tree = merge_trees(&side1_tree, &base_tree, &side2_tree).unwrap();
-    let merged_tree = store.get_tree(&RepoPath::root(), &merged_tree_id).unwrap();
     assert_eq!(merged_tree.id(), store.empty_tree_id());
 }
 
@@ -357,22 +352,20 @@ fn test_subtree_one_missing(use_git: bool) {
     let tree3 = write_tree(vec!["d1/f1", "d1/f2"]);
 
     // The two sides add different trees
-    let merged_tree_id = tree::merge_trees(&tree2, &tree1, &tree3).unwrap();
+    let merged_tree = merge_trees(&tree2, &tree1, &tree3).unwrap();
-    let merged_tree = store.get_tree(&RepoPath::root(), &merged_tree_id).unwrap();
     let expected_entries = write_tree(vec!["d1/f1", "d1/f2"]).entries().collect_vec();
     assert_eq!(merged_tree.entries().collect_vec(), expected_entries);
     // Same tree other way
-    let merged_tree_id = tree::merge_trees(&tree3, &tree1, &tree2).unwrap();
+    let reverse_merged_tree = merge_trees(&tree3, &tree1, &tree2).unwrap();
-    assert_eq!(merged_tree_id, *merged_tree.id());
+    assert_eq!(reverse_merged_tree.id(), merged_tree.id());
 
     // One side removes, the other side modifies
-    let merged_tree_id = tree::merge_trees(&tree1, &tree2, &tree3).unwrap();
+    let merged_tree = merge_trees(&tree1, &tree2, &tree3).unwrap();
-    let merged_tree = store.get_tree(&RepoPath::root(), &merged_tree_id).unwrap();
     let expected_entries = write_tree(vec!["d1/f2"]).entries().collect_vec();
     assert_eq!(merged_tree.entries().collect_vec(), expected_entries);
     // Same tree other way
-    let merged_tree_id = tree::merge_trees(&tree3, &tree2, &tree1).unwrap();
+    let reverse_merged_tree = merge_trees(&tree3, &tree2, &tree1).unwrap();
-    assert_eq!(merged_tree_id, *merged_tree.id());
+    assert_eq!(reverse_merged_tree.id(), merged_tree.id());
 }
 
 #[test_case(false ; "local backend")]
@@ -426,8 +419,7 @@ fn test_types(use_git: bool) {
     let side2_tree = store.get_tree(&RepoPath::root(), &side2_tree_id).unwrap();
 
     // Created the merged tree
-    let merged_tree_id = tree::merge_trees(&side1_tree, &base_tree, &side2_tree).unwrap();
+    let merged_tree = merge_trees(&side1_tree, &base_tree, &side2_tree).unwrap();
-    let merged_tree = store.get_tree(&RepoPath::root(), &merged_tree_id).unwrap();
 
     // Check the conflicting cases
     let component = RepoPathComponent::from("normal_executable_symlink");
@@ -492,13 +484,8 @@ fn test_simplify_conflict(use_git: bool) {
     let upstream1_tree = write_tree("upstream1 contents");
     let upstream2_tree = write_tree("upstream2 contents");
 
-    let merge_trees = |side1: &Tree, base: &Tree, side2: &Tree| -> Tree {
-        let tree_id = tree::merge_trees(side1, base, side2).unwrap();
-        store.get_tree(&RepoPath::root(), &tree_id).unwrap()
-    };
-
     // Rebase the branch tree to the first upstream tree
-    let rebased1_tree = merge_trees(&branch_tree, &base_tree, &upstream1_tree);
+    let rebased1_tree = merge_trees(&branch_tree, &base_tree, &upstream1_tree).unwrap();
     // Make sure we have a conflict (testing the test setup)
     match rebased1_tree.value(&component).unwrap() {
         TreeValue::Conflict(_) => {
@@ -509,12 +496,12 @@ fn test_simplify_conflict(use_git: bool) {
 
     // Rebase the rebased tree back to the base. The conflict should be gone. Try
     // both directions.
-    let rebased_back_tree = merge_trees(&rebased1_tree, &upstream1_tree, &base_tree);
+    let rebased_back_tree = merge_trees(&rebased1_tree, &upstream1_tree, &base_tree).unwrap();
     assert_eq!(
         rebased_back_tree.value(&component),
         branch_tree.value(&component)
     );
-    let rebased_back_tree = merge_trees(&base_tree, &upstream1_tree, &rebased1_tree);
+    let rebased_back_tree = merge_trees(&base_tree, &upstream1_tree, &rebased1_tree).unwrap();
     assert_eq!(
         rebased_back_tree.value(&component),
         branch_tree.value(&component)
@@ -522,7 +509,8 @@ fn test_simplify_conflict(use_git: bool) {
 
     // Rebase the rebased tree further upstream. The conflict should be simplified
     // to not mention the contents from the first rebase.
-    let further_rebased_tree = merge_trees(&rebased1_tree, &upstream1_tree, &upstream2_tree);
+    let further_rebased_tree =
+        merge_trees(&rebased1_tree, &upstream1_tree, &upstream2_tree).unwrap();
     match further_rebased_tree.value(&component).unwrap() {
         TreeValue::Conflict(id) => {
             let conflict = store
@@ -542,7 +530,8 @@ fn test_simplify_conflict(use_git: bool) {
         }
         _ => panic!("unexpected value"),
     };
-    let further_rebased_tree = merge_trees(&upstream2_tree, &upstream1_tree, &rebased1_tree);
+    let further_rebased_tree =
+        merge_trees(&upstream2_tree, &upstream1_tree, &rebased1_tree).unwrap();
     match further_rebased_tree.value(&component).unwrap() {
         TreeValue::Conflict(id) => {
             let conflict = store.read_conflict(&path, id).unwrap();

lib/tests/test_revset.rs

@@ -2306,11 +2306,7 @@ fn test_evaluate_expression_conflict(use_git: bool) {
     let tree1 = testutils::create_tree(repo, &[(&file_path1, "1"), (&file_path2, "1")]);
     let tree2 = testutils::create_tree(repo, &[(&file_path1, "2"), (&file_path2, "2")]);
     let tree3 = testutils::create_tree(repo, &[(&file_path1, "3"), (&file_path2, "1")]);
-    let tree_id4 = merge_trees(&tree2, &tree1, &tree3).unwrap();
+    let tree4 = merge_trees(&tree2, &tree1, &tree3).unwrap();
-    let tree4 = mut_repo
-        .store()
-        .get_tree(&RepoPath::root(), &tree_id4)
-        .unwrap();
 
     let mut create_commit = |parent_ids, tree_id| {
         mut_repo

src/commands/mod.rs

@@ -1832,16 +1832,12 @@ fn rebase_to_dest_parent(
             merge_commit_trees(workspace_command.repo().as_ref(), &destination.parents())?;
         let source_parent_tree =
             merge_commit_trees(workspace_command.repo().as_ref(), &source.parents())?;
-        let rebased_tree_id = merge_trees(
+        let rebased_tree = merge_trees(
             &destination_parent_tree,
             &source_parent_tree,
             &source.tree(),
         )?;
-        let tree = workspace_command
+        Ok(rebased_tree)
-            .repo()
-            .store()
-            .get_tree(&RepoPath::root(), &rebased_tree_id)?;
-        Ok(tree)
     }
 }
 
@@ -2378,14 +2374,14 @@ from the source will be moved into the destination.
         .store()
         .get_tree(&RepoPath::root(), &new_parent_tree_id)?;
     // Apply the reverse of the selected changes onto the source
-    let new_source_tree_id = merge_trees(&source_tree, &new_parent_tree, &parent_tree)?;
+    let new_source_tree = merge_trees(&source_tree, &new_parent_tree, &parent_tree)?;
-    let abandon_source = new_source_tree_id == *parent_tree.id();
+    let abandon_source = new_source_tree.id() == parent_tree.id();
     if abandon_source {
         tx.mut_repo().record_abandoned_commit(source.id().clone());
     } else {
         tx.mut_repo()
             .rewrite_commit(command.settings(), &source)
-            .set_tree(new_source_tree_id)
+            .set_tree(new_source_tree.id().clone())
             .write()?;
     }
     if tx.repo().index().is_ancestor(source.id(), destination.id()) {
@@ -2399,7 +2395,7 @@ from the source will be moved into the destination.
         destination = tx.mut_repo().store().get_commit(&rebased_destination_id)?;
     }
     // Apply the selected changes onto the destination
-    let new_destination_tree_id = merge_trees(&destination.tree(), &parent_tree, &new_parent_tree)?;
+    let new_destination_tree = merge_trees(&destination.tree(), &parent_tree, &new_parent_tree)?;
     let description = combine_messages(
         tx.base_repo(),
         &source,
@@ -2409,7 +2405,7 @@ from the source will be moved into the destination.
     )?;
     tx.mut_repo()
         .rewrite_commit(command.settings(), &destination)
-        .set_tree(new_destination_tree_id)
+        .set_tree(new_destination_tree.id().clone())
         .set_description(description)
         .write()?;
     tx.finish(ui)?;