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 itertools::Itertools as _;
21use pollster::FutureExt as _;
22
23use crate::backend::BackendResult;
24use crate::backend::TreeId;
25use crate::conflict_labels::ConflictLabels;
26use crate::merge::Merge;
27use crate::merge::MergeBuilder;
28use crate::merge::MergedTreeValue;
29use crate::merged_tree::MergedTree;
30use crate::repo_path::RepoPathBuf;
31use crate::tree_builder::TreeBuilder;
32
33/// Helper for writing trees with conflicts.
34///
35/// You start by creating an instance of this type with one or more
36/// base trees. You then add overrides on top. The overrides may be
37/// conflicts. Then you can write the result as a merge of trees.
38#[derive(Debug)]
39pub struct MergedTreeBuilder {
40    base_tree: MergedTree,
41    overrides: BTreeMap<RepoPathBuf, MergedTreeValue>,
42}
43
44impl MergedTreeBuilder {
45    /// Create a new builder with the given trees as base.
46    pub fn new(base_tree: MergedTree) -> Self {
47        Self {
48            base_tree,
49            overrides: BTreeMap::new(),
50        }
51    }
52
53    /// Set an override compared to  the base tree. The `values` merge must
54    /// either be resolved (i.e. have 1 side) or have the same number of
55    /// sides as the `base_tree_ids` used to construct this builder. Use
56    /// `Merge::absent()` to remove a value from the tree.
57    pub fn set_or_remove(&mut self, path: RepoPathBuf, values: MergedTreeValue) {
58        self.overrides.insert(path, values);
59    }
60
61    /// Create new tree(s) from the base tree(s) and overrides.
62    pub fn write_tree(self) -> BackendResult<MergedTree> {
63        let store = self.base_tree.store().clone();
64        let labels = self.base_tree.labels().clone();
65        let new_tree_ids = self.write_merged_trees()?;
66        let labels = if labels.num_sides() == Some(new_tree_ids.num_sides()) {
67            labels
68        } else {
69            // If the number of sides changed, we need to discard the conflict labels,
70            // otherwise `MergedTree::new` would panic.
71            // TODO: we should preserve conflict labels when setting conflicted tree values
72            // originating from a different tree than the base tree.
73            ConflictLabels::unlabeled()
74        };
75        let (labels, new_tree_ids) = labels.simplify_with(&new_tree_ids);
76        match new_tree_ids.into_resolved() {
77            Ok(single_tree_id) => Ok(MergedTree::resolved(store, single_tree_id)),
78            Err(tree_ids) => {
79                let tree = MergedTree::new(store, tree_ids, labels);
80                tree.resolve().block_on()
81            }
82        }
83    }
84
85    fn write_merged_trees(self) -> BackendResult<Merge<TreeId>> {
86        let store = self.base_tree.store().clone();
87        let mut base_tree_ids = self.base_tree.into_tree_ids();
88        let num_sides = self
89            .overrides
90            .values()
91            .map(|value| value.num_sides())
92            .max()
93            .unwrap_or(0);
94        base_tree_ids.pad_to(num_sides, store.empty_tree_id());
95        // Create a single-tree builder for each base tree
96        let mut tree_builders =
97            base_tree_ids.map(|base_tree_id| TreeBuilder::new(store.clone(), base_tree_id.clone()));
98        for (path, values) in self.overrides {
99            match values.into_resolved() {
100                Ok(value) => {
101                    // This path was overridden with a resolved value. Apply that to all
102                    // builders.
103                    for builder in &mut tree_builders {
104                        builder.set_or_remove(path.clone(), value.clone());
105                    }
106                }
107                Err(mut values) => {
108                    values.pad_to(num_sides, &None);
109                    // This path was overridden with a conflicted value. Apply each term to
110                    // its corresponding builder.
111                    for (builder, value) in zip(&mut tree_builders, values) {
112                        builder.set_or_remove(path.clone(), value);
113                    }
114                }
115            }
116        }
117        // TODO: This can be made more efficient. If there's a single resolved conflict
118        // in `dir/file`, we shouldn't have to write the `dir/` and root trees more than
119        // once.
120        let merge_builder: MergeBuilder<TreeId> = tree_builders
121            .into_iter()
122            .map(|builder| builder.write_tree())
123            .try_collect()?;
124        Ok(merge_builder.build())
125    }
126}
jj_lib/merged_tree_builder.rs

jj_lib/
merged_tree_builder.rs
