jujube_lib/
tree_builder.rs

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

jujube_lib/
tree_builder.rs
