Introduce a builder for merged trees and use it in jj chmod

cli/src/commands/mod.rs

@@ -37,7 +37,7 @@ use jj_lib::dag_walk::topo_order_reverse;
 use jj_lib::git_backend::GitBackend;
 use jj_lib::matchers::EverythingMatcher;
 use jj_lib::merge::Merge;
-use jj_lib::merged_tree::MergedTree;
+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;
@@ -2657,55 +2657,50 @@ fn cmd_chmod(ui: &mut Ui, command: &CommandHelper, args: &ChmodArgs) -> Result<(
         },
         commit.id().hex(),
     ));
-    let tree = commit.tree();
+    let tree = commit.merged_tree()?;
     let store = tree.store();
-    let mut tree_builder = store.tree_builder(tree.id().clone());
+    let mut tree_builder = MergedTreeBuilder::new(store.clone(), commit.merged_tree_id().clone());
     for repo_path in repo_paths {
         let user_error_with_path = |msg: &str| {
             user_error(format!(
                 "{msg} at '{}'.",
                 tx.base_workspace_helper().format_file_path(&repo_path)
             ))
         };
-        let new_tree_value = match tree.path_value(&repo_path) {
-            None => return Err(user_error_with_path("No such path")),
-            Some(TreeValue::File { id, executable: _ }) => TreeValue::File {
-                id,
+        let tree_value = tree.path_value(&repo_path);
+        if tree_value.is_absent() {
+            return Err(user_error_with_path("No such path"));
+        }
+        let all_files = tree_value
+            .adds()
+            .iter()
+            .flatten()
+            .all(|tree_value| matches!(tree_value, TreeValue::File { .. }));
+        if !all_files {
+            let message = if tree_value.is_resolved() {
+                "Found neither a file nor a conflict"
+            } else {
+                "Some of the sides of the conflict are not files"
+            };
+            return Err(user_error_with_path(message));
+        }
+        let new_tree_value = tree_value.map(|value| match value {
+            Some(TreeValue::File { id, executable: _ }) => Some(TreeValue::File {
+                id: id.clone(),
                 executable: executable_bit,
-            },
-            Some(TreeValue::Conflict(id)) => {
-                let conflict = tree.store().read_conflict(&repo_path, &id)?;
-                let all_files = conflict
-                    .adds()
-                    .iter()
-                    .flatten()
-                    .all(|tree_value| matches!(tree_value, TreeValue::File { .. }));
-                if !all_files {
-                    return Err(user_error_with_path(
-                        "Some of the sides of the conflict are not files",
-                    ));
-                }
-                let new_conflict = conflict.map(|value| match value {
-                    Some(TreeValue::File { id, executable: _ }) => Some(TreeValue::File {
-                        id: id.clone(),
-                        executable: executable_bit,
-                    }),
-                    Some(TreeValue::Conflict(_)) => {
-                        panic!("Conflict sides must not themselves be conflicts")
-                    }
-                    value => value.clone(),
-                });
-                let new_conflict_id = store.write_conflict(&repo_path, &new_conflict)?;
-                TreeValue::Conflict(new_conflict_id)
+            }),
+            Some(TreeValue::Conflict(_)) => {
+                panic!("Conflict sides must not themselves be conflicts")
             }
-            Some(_) => return Err(user_error_with_path("Found neither a file nor a conflict")),
-        };
-        tree_builder.set(repo_path, new_tree_value);
+            value => value.clone(),
+        });
+        tree_builder.set_or_remove(repo_path, new_tree_value);
     }
 
+    let new_tree_id = tree_builder.write_tree()?;
     tx.mut_repo()
         .rewrite_commit(command.settings(), &commit)
-        .set_tree(tree_builder.write_tree())
+        .set_tree_id(new_tree_id)
         .write()?;
     tx.finish(ui)
 }

lib/src/commit.rs

@@ -128,6 +128,10 @@ impl Commit {
         self.data.root_tree.as_legacy_tree_id()
     }
 
+    pub fn merged_tree_id(&self) -> &MergedTreeId {
+        &self.data.root_tree
+    }
+
     pub fn change_id(&self) -> &ChangeId {
         &self.data.change_id
     }

lib/src/commit_builder.rs

@@ -109,11 +109,18 @@ impl CommitBuilder<'_> {
         self.commit.root_tree.as_legacy_tree_id()
     }
 
+    // TODO(#1624): delete when all callers have been updated to support tree-level
+    // conflicts
     pub fn set_tree(mut self, tree_id: TreeId) -> Self {
         self.commit.root_tree = MergedTreeId::Legacy(tree_id);
         self
     }
 
+    pub fn set_tree_id(mut self, tree_id: MergedTreeId) -> Self {
+        self.commit.root_tree = tree_id;
+        self
+    }
+
     pub fn change_id(&self) -> &ChangeId {
         &self.commit.change_id
     }

lib/src/merged_tree.rs

@@ -15,13 +15,14 @@
 //! A lazily merged view of a set of trees.
 
 use std::cmp::max;
+use std::collections::BTreeMap;
 use std::sync::Arc;
 use std::{iter, vec};
 
 use itertools::Itertools;
 
 use crate::backend;
-use crate::backend::{ConflictId, TreeId, TreeValue};
+use crate::backend::{BackendResult, ConflictId, MergedTreeId, TreeId, TreeValue};
 use crate::matchers::Matcher;
 use crate::merge::Merge;
 use crate::repo_path::{RepoPath, RepoPathComponent, RepoPathJoin};
@@ -819,3 +820,72 @@ impl Iterator for TreeDiffIterator<'_> {
         None
     }
 }
+
+/// Helps with writing trees with conflicts. You start by creating an instance
+/// of this type with one or more base trees. You then add overrides on top. The
+/// overrides may be conflicts. Then you can write the result as a legacy tree
+/// (allowing path-level conflicts) or as multiple conflict-free trees.
+pub struct MergedTreeBuilder {
+    store: Arc<Store>,
+    base_tree_id: MergedTreeId,
+    overrides: BTreeMap<RepoPath, Merge<Option<TreeValue>>>,
+}
+
+impl MergedTreeBuilder {
+    /// Create a new builder with the given trees as base.
+    pub fn new(store: Arc<Store>, base_tree_id: MergedTreeId) -> Self {
+        MergedTreeBuilder {
+            store,
+            base_tree_id,
+            overrides: BTreeMap::new(),
+        }
+    }
+
+    /// Set an override compared to  the base tree. The `values` merge must
+    /// either be resolved (i.e. have 1 side) or have the same number of
+    /// sides as the `base_tree_ids` used to construct this builder. Use
+    /// `Merge::absent()` to remove a value from the tree. When the base tree is
+    /// a legacy tree, conflicts can be written either as a multi-way `Merge`
+    /// value or as a resolved `Merge` value using `TreeValue::Conflict`.
+    pub fn set_or_remove(&mut self, path: RepoPath, values: Merge<Option<TreeValue>>) {
+        if let MergedTreeId::Merge(base_tree_ids) = &self.base_tree_id {
+            if !values.is_resolved() {
+                assert_eq!(values.num_sides(), base_tree_ids.num_sides());
+            }
+            assert!(!values
+                .iter()
+                .flatten()
+                .any(|value| matches!(value, TreeValue::Conflict(_))));
+        }
+        self.overrides.insert(path, values);
+    }
+
+    /// Create new tree(s) from the base tree(s) and overrides.
+    ///
+    /// When the base tree was a legacy tree, then the result will be another
+    /// legacy tree. Overrides with conflicts will result in conflict objects
+    /// being written to the store.
+    pub fn write_tree(self) -> BackendResult<MergedTreeId> {
+        match self.base_tree_id {
+            MergedTreeId::Legacy(base_tree_id) => {
+                let mut tree_builder = TreeBuilder::new(self.store.clone(), base_tree_id);
+                for (path, values) in self.overrides {
+                    let values = values.simplify();
+                    match values.into_resolved() {
+                        Ok(value) => {
+                            tree_builder.set_or_remove(path, value);
+                        }
+                        Err(values) => {
+                            let conflict_id = self.store.write_conflict(&path, &values)?;
+                            tree_builder.set(path, TreeValue::Conflict(conflict_id));
+                        }
+                    }
+                }
+                Ok(MergedTreeId::Legacy(tree_builder.write_tree()))
+            }
+            MergedTreeId::Merge(_) => {
+                todo!()
+            }
+        }
+    }
+}

lib/src/working_copy.rs

@@ -40,7 +40,7 @@ use thiserror::Error;
 use tracing::{instrument, trace_span};
 
 use crate::backend::{
-    BackendError, FileId, MillisSinceEpoch, ObjectId, SymlinkId, TreeId, TreeValue,
+    BackendError, FileId, MergedTreeId, MillisSinceEpoch, ObjectId, SymlinkId, TreeId, TreeValue,
 };
 use crate::conflicts;
 #[cfg(feature = "watchman")]
@@ -52,7 +52,7 @@ use crate::matchers::{
     DifferenceMatcher, EverythingMatcher, IntersectionMatcher, Matcher, PrefixMatcher,
 };
 use crate::merge::Merge;
-use crate::merged_tree::MergedTree;
+use crate::merged_tree::{MergedTree, MergedTreeBuilder};
 use crate::op_store::{OperationId, WorkspaceId};
 use crate::repo_path::{FsPathParseError, RepoPath, RepoPathComponent, RepoPathJoin};
 use crate::settings::HumanByteSize;
@@ -673,7 +673,10 @@ impl TreeState {
             )
         })?;
 
-        let mut tree_builder = self.store.tree_builder(self.tree_id.clone());
+        let mut tree_builder = MergedTreeBuilder::new(
+            self.store.clone(),
+            MergedTreeId::Legacy(self.tree_id.clone()),
+        );
         let mut deleted_files: HashSet<_> =
             trace_span!("collecting existing files").in_scope(|| {
                 self.file_states
@@ -686,15 +689,7 @@ impl TreeState {
             });
         trace_span!("process tree entries").in_scope(|| -> Result<(), SnapshotError> {
             while let Ok((path, tree_values)) = tree_entries_rx.recv() {
-                match tree_values.into_resolved() {
-                    Ok(tree_value) => {
-                        tree_builder.set(path, tree_value.unwrap());
-                    }
-                    Err(conflict) => {
-                        let conflict_id = self.store.write_conflict(&path, &conflict)?;
-                        tree_builder.set(path, TreeValue::Conflict(conflict_id));
-                    }
-                }
+                tree_builder.set_or_remove(path, tree_values);
             }
             Ok(())
         })?;
@@ -713,11 +708,15 @@ impl TreeState {
             for file in &deleted_files {
                 is_dirty = true;
                 self.file_states.remove(file);
-                tree_builder.remove(file.clone());
+                tree_builder.set_or_remove(file.clone(), Merge::absent());
             }
         });
         trace_span!("write tree").in_scope(|| {
-            let new_tree_id = tree_builder.write_tree();
+            let new_tree_id = tree_builder
+                .write_tree()
+                .unwrap()
+                .as_legacy_tree_id()
+                .clone();
             is_dirty |= new_tree_id != self.tree_id;
             self.tree_id = new_tree_id;
         });