merged_tree: introduce a type for a set of trees to merge on the fly

In order to store conflicts in the commit, as conflicts between a set
of trees, we want to be able merge those trees on the fly. This
introduces a type for that. It has a `Merge(Conflict(Tree))` variant,
where the individual trees cannot have path-level conflicts. It also
has a `Legacy(Tree)` variant, which does allow path-level conflicts. I
think that should help us with the migration.

lib/src/lib.rs

@@ -43,6 +43,7 @@ pub mod local_backend;
 pub mod lock;
 pub mod matchers;
 pub mod merge;
+pub mod merged_tree;
 pub mod op_heads_store;
 pub mod op_store;
 pub mod operation;

lib/src/merged_tree.rs

@@ -0,0 +1,157 @@
+// Copyright 2023 The Jujutsu Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! A lazily merged view of a set of trees.
+
+use std::cmp::max;
+
+use itertools::Itertools;
+
+use crate::backend::TreeValue;
+use crate::conflicts::Conflict;
+use crate::repo_path::{RepoPath, RepoPathComponent};
+use crate::store::Store;
+use crate::tree::Tree;
+use crate::tree_builder::TreeBuilder;
+
+/// Presents a view of a merged set of trees.
+pub enum MergedTree {
+    /// A single tree, possibly with path-level conflicts.
+    Legacy(Tree),
+    /// A merge of multiple trees, or just a single tree. The individual trees
+    /// have no path-level conflicts.
+    Merge(Conflict<Tree>),
+}
+
+/// The value at a given path in a `MergedTree`.
+#[derive(PartialEq, Eq, Hash, Clone, Debug)]
+pub enum MergedTreeValue<'a> {
+    /// A single non-conflicted value.
+    Resolved(Option<&'a TreeValue>),
+    /// TODO: Make this a `Conflict<Option<&'a TreeValue>>` (reference to the
+    /// value) once we have removed the `MergedTree::Legacy` variant.
+    Conflict(Conflict<Option<TreeValue>>),
+}
+
+impl MergedTree {
+    /// Creates a new `MergedTree` representing a single tree without conflicts.
+    pub fn resolved(tree: Tree) -> Self {
+        MergedTree::new(Conflict::resolved(tree))
+    }
+
+    /// Creates a new `MergedTree` representing a merge of a set of trees. The
+    /// individual trees must not have any conflicts.
+    pub fn new(conflict: Conflict<Tree>) -> Self {
+        debug_assert!(!conflict.removes().iter().any(|t| t.has_conflict()));
+        debug_assert!(!conflict.adds().iter().any(|t| t.has_conflict()));
+        debug_assert!(itertools::chain(conflict.removes(), conflict.adds())
+            .map(|tree| tree.dir())
+            .all_equal());
+        debug_assert!(itertools::chain(conflict.removes(), conflict.adds())
+            .map(|tree| Arc::as_ptr(tree.store()))
+            .all_equal());
+        MergedTree::Merge(conflict)
+    }
+
+    /// Creates a new `MergedTree` backed by a tree with path-level conflicts.
+    pub fn legacy(tree: Tree) -> Self {
+        MergedTree::Legacy(tree)
+    }
+
+    /// Takes a tree in the legacy format (with path-level conflicts in the
+    /// tree) and returns a `MergedTree` with any conflicts converted to
+    /// tree-level conflicts.
+    pub fn from_legacy_tree(tree: Tree) -> Self {
+        let conflict_ids = tree.conflicts();
+        if conflict_ids.is_empty() {
+            return MergedTree::resolved(tree);
+        }
+        // Find the number of removes in the most complex conflict. We will then
+        // build `2*num_removes + 1` trees
+        let mut max_num_removes = 0;
+        let store = tree.store();
+        let mut conflicts: Vec<(&RepoPath, Conflict<Option<TreeValue>>)> = vec![];
+        for (path, conflict_id) in &conflict_ids {
+            let conflict = store.read_conflict(path, conflict_id).unwrap();
+            max_num_removes = max(max_num_removes, conflict.removes().len());
+            conflicts.push((path, conflict));
+        }
+        let mut removes = vec![];
+        let mut adds = vec![store.tree_builder(tree.id().clone())];
+        for _ in 0..max_num_removes {
+            removes.push(store.tree_builder(tree.id().clone()));
+            adds.push(store.tree_builder(tree.id().clone()));
+        }
+        for (path, conflict) in conflicts {
+            let num_removes = conflict.removes().len();
+            // If there are fewer terms in this conflict than in some other conflict, we can
+            // add canceling removes and adds of any value. The simplest value is an absent
+            // value, so we use that.
+            for i in num_removes..max_num_removes {
+                removes[i].remove(path.clone());
+                adds[i + 1].remove(path.clone());
+            }
+            // Now add the terms that were present in the conflict to the appropriate trees.
+            for (i, term) in conflict.removes().iter().enumerate() {
+                match term {
+                    Some(value) => removes[i].set(path.clone(), value.clone()),
+                    None => removes[i].remove(path.clone()),
+                }
+            }
+            for (i, term) in conflict.adds().iter().enumerate() {
+                match term {
+                    Some(value) => adds[i].set(path.clone(), value.clone()),
+                    None => adds[i].remove(path.clone()),
+                }
+            }
+        }
+
+        let write_tree = |builder: TreeBuilder| {
+            let tree_id = builder.write_tree();
+            store.get_tree(&RepoPath::root(), &tree_id).unwrap()
+        };
+
+        MergedTree::Merge(Conflict::new(
+            removes.into_iter().map(write_tree).collect(),
+            adds.into_iter().map(write_tree).collect(),
+        ))
+    }
+
+    /// The value at the given basename. The value can be `Resolved` even if
+    /// `self` is a `Conflict`, which happens if the value at the path can be
+    /// trivially merged. Does not recurse, so if `basename` refers to a Tree,
+    /// then a `TreeValue::Tree` will be returned.
+    pub fn value(&self, basename: &RepoPathComponent) -> MergedTreeValue {
+        match self {
+            MergedTree::Legacy(tree) => match tree.value(basename) {
+                Some(TreeValue::Conflict(conflict_id)) => {
+                    let conflict = tree.store().read_conflict(tree.dir(), conflict_id).unwrap();
+                    MergedTreeValue::Conflict(conflict)
+                }
+                other => MergedTreeValue::Resolved(other),
+            },
+            MergedTree::Merge(conflict) => {
+                if let Some(tree) = conflict.as_resolved() {
+                    return MergedTreeValue::Resolved(tree.value(basename));
+                }
+                let value = conflict.map(|tree| tree.value(basename));
+                if let Some(resolved) = value.resolve_trivial() {
+                    return MergedTreeValue::Resolved(*resolved);
+                }
+
+                MergedTreeValue::Conflict(value.map(|x| x.cloned()))
+            }
+        }
+    }
+}

lib/tests/test_merged_tree.rs

@@ -0,0 +1,180 @@
+// Copyright 2023 The Jujutsu Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use jj_lib::backend::{FileId, TreeValue};
+use jj_lib::conflicts::Conflict;
+use jj_lib::merged_tree::{MergedTree, MergedTreeValue};
+use jj_lib::repo::Repo;
+use jj_lib::repo_path::{RepoPath, RepoPathComponent, RepoPathJoin};
+use testutils::{write_file, write_normal_file, TestRepo};
+
+fn file_value(file_id: &FileId) -> TreeValue {
+    TreeValue::File {
+        id: file_id.clone(),
+        executable: false,
+    }
+}
+
+#[test]
+fn test_from_legacy_tree() {
+    let test_repo = TestRepo::init(true);
+    let repo = &test_repo.repo;
+    let store = repo.store();
+
+    let mut tree_builder = store.tree_builder(repo.store().empty_tree_id().clone());
+
+    // file1: regular file without conflicts
+    let file1_path = RepoPath::from_internal_string("no_conflict");
+    let file1_id = write_normal_file(&mut tree_builder, &file1_path, "foo");
+
+    // file2: 3-way conflict
+    let file2_path = RepoPath::from_internal_string("3way");
+    let file2_v1_id = write_file(store.as_ref(), &file2_path, "file2_v1");
+    let file2_v2_id = write_file(store.as_ref(), &file2_path, "file2_v2");
+    let file2_v3_id = write_file(store.as_ref(), &file2_path, "file2_v3");
+    let file2_conflict = Conflict::new(
+        vec![Some(file_value(&file2_v1_id))],
+        vec![
+            Some(file_value(&file2_v2_id)),
+            Some(file_value(&file2_v3_id)),
+        ],
+    );
+    let file2_conflict_id = store.write_conflict(&file2_path, &file2_conflict).unwrap();
+    tree_builder.set(file2_path.clone(), TreeValue::Conflict(file2_conflict_id));
+
+    // file3: modify/delete conflict
+    let file3_path = RepoPath::from_internal_string("modify_delete");
+    let file3_v1_id = write_file(store.as_ref(), &file3_path, "file3_v1");
+    let file3_v2_id = write_file(store.as_ref(), &file3_path, "file3_v2");
+    let file3_conflict = Conflict::new(
+        vec![Some(file_value(&file3_v1_id))],
+        vec![Some(file_value(&file3_v2_id)), None],
+    );
+    let file3_conflict_id = store.write_conflict(&file3_path, &file3_conflict).unwrap();
+    tree_builder.set(file3_path.clone(), TreeValue::Conflict(file3_conflict_id));
+
+    // file4: add/add conflict
+    let file4_path = RepoPath::from_internal_string("add_add");
+    let file4_v1_id = write_file(store.as_ref(), &file4_path, "file4_v1");
+    let file4_v2_id = write_file(store.as_ref(), &file4_path, "file4_v2");
+    let file4_conflict = Conflict::new(
+        vec![None],
+        vec![
+            Some(file_value(&file4_v1_id)),
+            Some(file_value(&file4_v2_id)),
+        ],
+    );
+    let file4_conflict_id = store.write_conflict(&file4_path, &file4_conflict).unwrap();
+    tree_builder.set(file4_path.clone(), TreeValue::Conflict(file4_conflict_id));
+
+    // file5: 5-way conflict
+    let file5_path = RepoPath::from_internal_string("5way");
+    let file5_v1_id = write_file(store.as_ref(), &file5_path, "file5_v1");
+    let file5_v2_id = write_file(store.as_ref(), &file5_path, "file5_v2");
+    let file5_v3_id = write_file(store.as_ref(), &file5_path, "file5_v3");
+    let file5_v4_id = write_file(store.as_ref(), &file5_path, "file5_v4");
+    let file5_v5_id = write_file(store.as_ref(), &file5_path, "file5_v5");
+    let file5_conflict = Conflict::new(
+        vec![
+            Some(file_value(&file5_v1_id)),
+            Some(file_value(&file5_v2_id)),
+        ],
+        vec![
+            Some(file_value(&file5_v3_id)),
+            Some(file_value(&file5_v4_id)),
+            Some(file_value(&file5_v5_id)),
+        ],
+    );
+    let file5_conflict_id = store.write_conflict(&file5_path, &file5_conflict).unwrap();
+    tree_builder.set(file5_path.clone(), TreeValue::Conflict(file5_conflict_id));
+
+    // dir1: directory without conflicts
+    let dir1_basename = RepoPathComponent::from("dir1");
+    write_normal_file(
+        &mut tree_builder,
+        &RepoPath::root()
+            .join(&dir1_basename)
+            .join(&RepoPathComponent::from("file")),
+        "foo",
+    );
+
+    let tree_id = tree_builder.write_tree();
+    let tree = store.get_tree(&RepoPath::root(), &tree_id).unwrap();
+
+    let merged_tree = MergedTree::from_legacy_tree(tree.clone());
+    assert_eq!(
+        merged_tree.value(&RepoPathComponent::from("missing")),
+        MergedTreeValue::Resolved(None)
+    );
+    // file1: regular file without conflicts
+    assert_eq!(
+        merged_tree.value(&file1_path.components()[0]),
+        MergedTreeValue::Resolved(Some(&TreeValue::File {
+            id: file1_id,
+            executable: false,
+        }))
+    );
+    // file2: 3-way conflict
+    assert_eq!(
+        merged_tree.value(&file2_path.components()[0]),
+        MergedTreeValue::Conflict(Conflict::new(
+            vec![Some(file_value(&file2_v1_id)), None],
+            vec![
+                Some(file_value(&file2_v2_id)),
+                Some(file_value(&file2_v3_id)),
+                None,
+            ],
+        ))
+    );
+    // file3: modify/delete conflict
+    assert_eq!(
+        merged_tree.value(&file3_path.components()[0]),
+        MergedTreeValue::Conflict(Conflict::new(
+            vec![Some(file_value(&file3_v1_id)), None],
+            vec![Some(file_value(&file3_v2_id)), None, None],
+        ))
+    );
+    // file4: add/add conflict
+    assert_eq!(
+        merged_tree.value(&file4_path.components()[0]),
+        MergedTreeValue::Conflict(Conflict::new(
+            vec![None, None],
+            vec![
+                Some(file_value(&file4_v1_id)),
+                Some(file_value(&file4_v2_id)),
+                None
+            ],
+        ))
+    );
+    // file5: 5-way conflict
+    assert_eq!(
+        merged_tree.value(&file5_path.components()[0]),
+        MergedTreeValue::Conflict(Conflict::new(
+            vec![
+                Some(file_value(&file5_v1_id)),
+                Some(file_value(&file5_v2_id)),
+            ],
+            vec![
+                Some(file_value(&file5_v3_id)),
+                Some(file_value(&file5_v4_id)),
+                Some(file_value(&file5_v5_id)),
+            ],
+        ))
+    );
+    // file6: directory without conflicts
+    assert_eq!(
+        merged_tree.value(&dir1_basename),
+        MergedTreeValue::Resolved(tree.value(&dir1_basename))
+    );
+}

lib/testutils/src/lib.rs

@@ -176,15 +176,20 @@ pub fn write_file(store: &Store, path: &RepoPath, contents: &str) -> FileId {
     store.write_file(path, &mut contents.as_bytes()).unwrap()
 }
 
-pub fn write_normal_file(tree_builder: &mut TreeBuilder, path: &RepoPath, contents: &str) {
+pub fn write_normal_file(
+    tree_builder: &mut TreeBuilder,
+    path: &RepoPath,
+    contents: &str,
+) -> FileId {
     let id = write_file(tree_builder.store(), path, contents);
     tree_builder.set(
         path.clone(),
         TreeValue::File {
-            id,
+            id: id.clone(),
             executable: false,
         },
     );
+    id
 }
 
 pub fn write_executable_file(tree_builder: &mut TreeBuilder, path: &RepoPath, contents: &str) {