working_copy: move writing of conflict objects into new tree builder

This introduces a `MergedTreeBuilder` type, which takes a set of base
trees and overrides. The idea is that it will be able to write
multiple trees or a legacy tree. For now, it's only able to write
legacy trees. To show that it works, the working copy's snaphotting
code has been updated to use it.

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;
         });