jujutsu_lib/
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
15use std::collections::{BTreeMap, HashSet};
16use std::sync::Arc;
17
18use crate::backend;
19use crate::backend::{TreeId, TreeValue};
20use crate::repo_path::{RepoPath, RepoPathJoin};
21use crate::store::Store;
22use crate::tree::Tree;
23
24#[derive(Debug)]
25enum Override {
26    Tombstone,
27    Replace(TreeValue),
28}
29
30#[derive(Debug)]
31pub struct TreeBuilder {
32    store: Arc<Store>,
33    base_tree_id: TreeId,
34    overrides: BTreeMap<RepoPath, Override>,
35}
36
37impl TreeBuilder {
38    pub fn new(store: Arc<Store>, base_tree_id: TreeId) -> TreeBuilder {
39        let overrides = BTreeMap::new();
40        TreeBuilder {
41            store,
42            base_tree_id,
43            overrides,
44        }
45    }
46
47    pub fn store(&self) -> &Store {
48        self.store.as_ref()
49    }
50
51    pub fn has_overrides(&self) -> bool {
52        !self.overrides.is_empty()
53    }
54
55    pub fn set(&mut self, path: RepoPath, value: TreeValue) {
56        self.overrides.insert(path, Override::Replace(value));
57    }
58
59    pub fn remove(&mut self, path: RepoPath) {
60        self.overrides.insert(path, Override::Tombstone);
61    }
62
63    pub fn write_tree(mut self) -> TreeId {
64        let mut trees_to_write = self.get_base_trees();
65        if trees_to_write.is_empty() {
66            return self.base_tree_id;
67        }
68
69        // Update entries in parent trees for file overrides
70        for (path, file_override) in self.overrides {
71            if let Some((dir, basename)) = path.split() {
72                let tree = trees_to_write.get_mut(&dir).unwrap();
73                match file_override {
74                    Override::Replace(value) => {
75                        tree.set(basename.clone(), value);
76                    }
77                    Override::Tombstone => {
78                        tree.remove(basename);
79                    }
80                }
81            }
82        }
83
84        // Write trees level by level, starting with trees without children.
85        let store = self.store.as_ref();
86        loop {
87            let mut dirs_to_write: HashSet<RepoPath> = trees_to_write.keys().cloned().collect();
88
89            for dir in trees_to_write.keys() {
90                if let Some(parent) = dir.parent() {
91                    dirs_to_write.remove(&parent);
92                }
93            }
94
95            for dir in dirs_to_write {
96                let tree = trees_to_write.remove(&dir).unwrap();
97
98                if let Some((parent, basename)) = dir.split() {
99                    let parent_tree = trees_to_write.get_mut(&parent).unwrap();
100                    if tree.is_empty() {
101                        parent_tree.remove(basename);
102                    } else {
103                        let tree_id = store.write_tree(&dir, &tree).unwrap();
104                        parent_tree.set(basename.clone(), TreeValue::Tree(tree_id));
105                    }
106                } else {
107                    // We're writing the root tree. Write it even if empty. Return its id.
108                    return store.write_tree(&dir, &tree).unwrap();
109                }
110            }
111        }
112    }
113
114    fn get_base_trees(&mut self) -> BTreeMap<RepoPath, backend::Tree> {
115        let mut tree_cache = BTreeMap::new();
116        let mut base_trees = BTreeMap::new();
117        let store = self.store.clone();
118
119        let mut populate_trees = |dir: &RepoPath| {
120            let mut current_dir = RepoPath::root();
121
122            if !tree_cache.contains_key(&current_dir) {
123                let tree = store.get_tree(&current_dir, &self.base_tree_id).unwrap();
124                let store_tree = tree.data().clone();
125                tree_cache.insert(current_dir.clone(), tree);
126                base_trees.insert(current_dir.clone(), store_tree);
127            }
128
129            for component in dir.components() {
130                let next_dir = current_dir.join(component);
131                let current_tree = tree_cache.get(&current_dir).unwrap();
132                if !tree_cache.contains_key(&next_dir) {
133                    let tree = current_tree
134                        .sub_tree(component)
135                        .unwrap_or_else(|| Tree::null(self.store.clone(), next_dir.clone()));
136                    let store_tree = tree.data().clone();
137                    tree_cache.insert(next_dir.clone(), tree);
138                    base_trees.insert(next_dir.clone(), store_tree);
139                }
140                current_dir = next_dir;
141            }
142        };
143        for path in self.overrides.keys() {
144            if let Some(parent) = path.parent() {
145                populate_trees(&parent);
146            }
147        }
148
149        base_trees
150    }
151}
jujutsu_lib/tree_builder.rs

jujutsu_lib/
tree_builder.rs
