use std::collections::{BTreeMap, HashSet};
use std::sync::Arc;
use crate::backend;
use crate::backend::{TreeId, TreeValue};
use crate::repo_path::{RepoPath, RepoPathJoin};
use crate::store::Store;
use crate::tree::Tree;
#[derive(Debug)]
enum Override {
Tombstone,
Replace(TreeValue),
}
#[derive(Debug)]
pub struct TreeBuilder {
store: Arc<Store>,
base_tree_id: TreeId,
overrides: BTreeMap<RepoPath, Override>,
}
impl TreeBuilder {
pub fn new(store: Arc<Store>, base_tree_id: TreeId) -> TreeBuilder {
let overrides = BTreeMap::new();
TreeBuilder {
store,
base_tree_id,
overrides,
}
}
pub fn store(&self) -> &Store {
self.store.as_ref()
}
pub fn has_overrides(&self) -> bool {
!self.overrides.is_empty()
}
pub fn set(&mut self, path: RepoPath, value: TreeValue) {
self.overrides.insert(path, Override::Replace(value));
}
pub fn remove(&mut self, path: RepoPath) {
self.overrides.insert(path, Override::Tombstone);
}
pub fn write_tree(mut self) -> TreeId {
let mut trees_to_write = self.get_base_trees();
if trees_to_write.is_empty() {
return self.base_tree_id;
}
for (path, file_override) in self.overrides {
if let Some((dir, basename)) = path.split() {
let tree = trees_to_write.get_mut(&dir).unwrap();
match file_override {
Override::Replace(value) => {
tree.set(basename.clone(), value);
}
Override::Tombstone => {
tree.remove(basename);
}
}
}
}
let store = self.store.as_ref();
loop {
let mut dirs_to_write: HashSet<RepoPath> =
trees_to_write.keys().cloned().into_iter().collect();
for dir in trees_to_write.keys() {
if let Some(parent) = dir.parent() {
dirs_to_write.remove(&parent);
}
}
for dir in dirs_to_write {
let tree = trees_to_write.remove(&dir).unwrap();
if let Some((parent, basename)) = dir.split() {
let parent_tree = trees_to_write.get_mut(&parent).unwrap();
if tree.is_empty() {
parent_tree.remove(basename);
} else {
let tree_id = store.write_tree(&dir, &tree).unwrap();
parent_tree.set(basename.clone(), TreeValue::Tree(tree_id));
}
} else {
return store.write_tree(&dir, &tree).unwrap();
}
}
}
}
fn get_base_trees(&mut self) -> BTreeMap<RepoPath, backend::Tree> {
let mut tree_cache = BTreeMap::new();
let mut base_trees = BTreeMap::new();
let store = self.store.clone();
let mut populate_trees = |dir: &RepoPath| {
let mut current_dir = RepoPath::root();
if !tree_cache.contains_key(¤t_dir) {
let tree = store.get_tree(¤t_dir, &self.base_tree_id).unwrap();
let store_tree = tree.data().clone();
tree_cache.insert(current_dir.clone(), tree);
base_trees.insert(current_dir.clone(), store_tree);
}
for component in dir.components() {
let next_dir = current_dir.join(component);
let current_tree = tree_cache.get(¤t_dir).unwrap();
if !tree_cache.contains_key(&next_dir) {
let tree = current_tree
.sub_tree(component)
.unwrap_or_else(|| Tree::null(self.store.clone(), next_dir.clone()));
let store_tree = tree.data().clone();
tree_cache.insert(next_dir.clone(), tree);
base_trees.insert(next_dir.clone(), store_tree);
}
current_dir = next_dir;
}
};
for path in self.overrides.keys() {
if let Some(parent) = path.parent() {
populate_trees(&parent);
}
}
base_trees
}
}