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