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

jj_lib/
tree_builder.rs
