Skip to main content

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}