jj_lib/merged_tree_builder.rs
1// Copyright 2020 The Jujutsu Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Helps build a new `MergedTree` from a base tree and overrides.
16
17use std::collections::BTreeMap;
18use std::iter::zip;
19
20use futures::future::try_join_all;
21
22use crate::backend::BackendResult;
23use crate::backend::TreeId;
24use crate::conflict_labels::ConflictLabels;
25use crate::merge::Merge;
26use crate::merge::MergedTreeValue;
27use crate::merged_tree::MergedTree;
28use crate::repo_path::RepoPathBuf;
29use crate::tree_builder::TreeBuilder;
30
31/// Helper for writing trees with conflicts.
32///
33/// You start by creating an instance of this type with one or more
34/// base trees. You then add overrides on top. The overrides may be
35/// conflicts. Then you can write the result as a merge of trees.
36#[derive(Debug)]
37pub struct MergedTreeBuilder {
38 base_tree: MergedTree,
39 overrides: BTreeMap<RepoPathBuf, MergedTreeValue>,
40}
41
42impl MergedTreeBuilder {
43 /// Create a new builder with the given trees as base.
44 pub fn new(base_tree: MergedTree) -> Self {
45 Self {
46 base_tree,
47 overrides: BTreeMap::new(),
48 }
49 }
50
51 /// Set an override compared to the base tree. The `values` merge must
52 /// either be resolved (i.e. have 1 side) or have the same number of
53 /// sides as the `base_tree_ids` used to construct this builder. Use
54 /// `Merge::absent()` to remove a value from the tree.
55 pub fn set_or_remove(&mut self, path: RepoPathBuf, values: MergedTreeValue) {
56 self.overrides.insert(path, values);
57 }
58
59 /// Create new tree(s) from the base tree(s) and overrides.
60 pub async fn write_tree(self) -> BackendResult<MergedTree> {
61 let store = self.base_tree.store().clone();
62 let labels = self.base_tree.labels().clone();
63 let new_tree_ids = self.write_merged_trees().await?;
64 let labels = if labels.num_sides() == Some(new_tree_ids.num_sides()) {
65 labels
66 } else {
67 // If the number of sides changed, we need to discard the conflict labels,
68 // otherwise `MergedTree::new` would panic.
69 // TODO: we should preserve conflict labels when setting conflicted tree values
70 // originating from a different tree than the base tree.
71 ConflictLabels::unlabeled()
72 };
73 let (labels, new_tree_ids) = labels.simplify_with(&new_tree_ids);
74 match new_tree_ids.into_resolved() {
75 Ok(single_tree_id) => Ok(MergedTree::resolved(store, single_tree_id)),
76 Err(tree_ids) => {
77 let tree = MergedTree::new(store, tree_ids, labels);
78 tree.resolve().await
79 }
80 }
81 }
82
83 async fn write_merged_trees(self) -> BackendResult<Merge<TreeId>> {
84 let store = self.base_tree.store().clone();
85 let mut base_tree_ids = self.base_tree.into_tree_ids();
86 let num_sides = self
87 .overrides
88 .values()
89 .map(|value| value.num_sides())
90 .max()
91 .unwrap_or(0);
92 base_tree_ids.pad_to(num_sides, store.empty_tree_id());
93 // Create a single-tree builder for each base tree
94 let mut tree_builders =
95 base_tree_ids.into_map(|base_tree_id| TreeBuilder::new(store.clone(), base_tree_id));
96 for (path, values) in self.overrides {
97 match values.into_resolved() {
98 Ok(value) => {
99 // This path was overridden with a resolved value. Apply that to all
100 // builders.
101 for builder in &mut tree_builders {
102 builder.set_or_remove(path.clone(), value.clone());
103 }
104 }
105 Err(mut values) => {
106 values.pad_to(num_sides, &None);
107 // This path was overridden with a conflicted value. Apply each term to
108 // its corresponding builder.
109 for (builder, value) in zip(&mut tree_builders, values) {
110 builder.set_or_remove(path.clone(), value);
111 }
112 }
113 }
114 }
115 // TODO: This can be made more efficient. If there's a single resolved conflict
116 // in `dir/file`, we shouldn't have to write the `dir/` and root trees more than
117 // once.
118 let tree_ids = try_join_all(
119 tree_builders
120 .into_iter()
121 .map(|builder| builder.write_tree()),
122 )
123 .await?;
124 let merge = Merge::from_vec(tree_ids);
125 Ok(merge)
126 }
127}