jj_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
15#![allow(missing_docs)]
16
17use std::collections::BTreeMap;
18use std::sync::Arc;
19
20use pollster::FutureExt as _;
21
22use crate::backend;
23use crate::backend::BackendResult;
24use crate::backend::TreeId;
25use crate::backend::TreeValue;
26use crate::repo_path::RepoPath;
27use crate::repo_path::RepoPathBuf;
28use crate::store::Store;
29use crate::tree::Tree;
30
31#[derive(Debug)]
32enum Override {
33    Tombstone,
34    Replace(TreeValue),
35}
36
37#[derive(Debug)]
38pub struct TreeBuilder {
39    store: Arc<Store>,
40    base_tree_id: TreeId,
41    overrides: BTreeMap<RepoPathBuf, Override>,
42}
43
44impl TreeBuilder {
45    pub fn new(store: Arc<Store>, base_tree_id: TreeId) -> TreeBuilder {
46        let overrides = BTreeMap::new();
47        TreeBuilder {
48            store,
49            base_tree_id,
50            overrides,
51        }
52    }
53
54    pub fn store(&self) -> &Store {
55        self.store.as_ref()
56    }
57
58    pub fn set(&mut self, path: RepoPathBuf, value: TreeValue) {
59        assert!(!path.is_root());
60        self.overrides.insert(path, Override::Replace(value));
61    }
62
63    pub fn remove(&mut self, path: RepoPathBuf) {
64        assert!(!path.is_root());
65        self.overrides.insert(path, Override::Tombstone);
66    }
67
68    pub fn set_or_remove(&mut self, path: RepoPathBuf, value: Option<TreeValue>) {
69        assert!(!path.is_root());
70        if let Some(value) = value {
71            self.overrides.insert(path, Override::Replace(value));
72        } else {
73            self.overrides.insert(path, Override::Tombstone);
74        }
75    }
76
77    pub fn write_tree(self) -> BackendResult<TreeId> {
78        if self.overrides.is_empty() {
79            return Ok(self.base_tree_id);
80        }
81
82        let mut trees_to_write = self.get_base_trees()?;
83
84        // Update entries in parent trees for file overrides
85        for (path, file_override) in self.overrides {
86            let (dir, basename) = path.split().unwrap();
87            let tree = trees_to_write.get_mut(dir).unwrap();
88            match file_override {
89                Override::Replace(value) => {
90                    tree.set(basename.to_owned(), value);
91                }
92                Override::Tombstone => {
93                    tree.remove(basename);
94                }
95            }
96        }
97
98        // Write trees in reverse lexicographical order, starting with trees without
99        // children.
100        // TODO: Writing trees concurrently should help on high-latency backends
101        let store = &self.store;
102        while let Some((dir, tree)) = trees_to_write.pop_last() {
103            if let Some((parent, basename)) = dir.split() {
104                let parent_tree = trees_to_write.get_mut(parent).unwrap();
105                if tree.is_empty() {
106                    if let Some(TreeValue::Tree(_)) = parent_tree.value(basename) {
107                        parent_tree.remove(basename);
108                    } else {
109                        // Entry would have been replaced with file (see above)
110                    }
111                } else {
112                    let tree = store.write_tree(&dir, tree).block_on()?;
113                    parent_tree.set(basename.to_owned(), TreeValue::Tree(tree.id().clone()));
114                }
115            } else {
116                // We're writing the root tree. Write it even if empty. Return its id.
117                assert!(trees_to_write.is_empty());
118                let written_tree = store.write_tree(&dir, tree).block_on()?;
119                return Ok(written_tree.id().clone());
120            }
121        }
122
123        unreachable!("trees_to_write must contain the root tree");
124    }
125
126    fn get_base_trees(&self) -> BackendResult<BTreeMap<RepoPathBuf, backend::Tree>> {
127        let store = &self.store;
128        let mut tree_cache = {
129            let dir = RepoPathBuf::root();
130            let tree = store.get_tree(dir.clone(), &self.base_tree_id)?;
131            BTreeMap::from([(dir, tree)])
132        };
133
134        fn populate_trees<'a>(
135            tree_cache: &'a mut BTreeMap<RepoPathBuf, Tree>,
136            store: &Arc<Store>,
137            dir: &RepoPath,
138        ) -> BackendResult<&'a Tree> {
139            // `if let Some(tree) = ...` doesn't pass lifetime check as of Rust 1.84.0
140            if tree_cache.contains_key(dir) {
141                return Ok(tree_cache.get(dir).unwrap());
142            }
143            let (parent, basename) = dir.split().expect("root must be populated");
144            let tree = populate_trees(tree_cache, store, parent)?
145                .sub_tree(basename)?
146                .unwrap_or_else(|| Tree::empty(store.clone(), dir.to_owned()));
147            Ok(tree_cache.entry(dir.to_owned()).or_insert(tree))
148        }
149
150        for path in self.overrides.keys() {
151            let parent = path.parent().unwrap();
152            populate_trees(&mut tree_cache, store, parent)?;
153        }
154
155        Ok(tree_cache
156            .into_iter()
157            .map(|(dir, tree)| (dir, tree.data().clone()))
158            .collect())
159    }
160}
jj_lib/tree_builder.rs

jj_lib/
tree_builder.rs
