1
use std::collections::HashMap;
2

            
3
use bonsaidb_core::schema::{CollectionName, Schematic};
4
use nebari::{
5
    io::any::AnyFile,
6
    tree::{AnyTreeRoot, Root, Unversioned, Versioned},
7
};
8

            
9
#[cfg(any(feature = "encryption", feature = "compression"))]
10
use crate::storage::TreeVault;
11
use crate::{
12
    database::document_tree_name,
13
    views::{view_document_map_tree_name, view_entries_tree_name, view_invalidated_docs_tree_name},
14
};
15

            
16
546460
#[derive(Default)]
17
pub(crate) struct OpenTrees {
18
    pub trees: Vec<Box<dyn AnyTreeRoot<AnyFile>>>,
19
    pub trees_index_by_name: HashMap<String, usize>,
20
}
21

            
22
impl OpenTrees {
23
    #[cfg_attr(not(feature = "encryption"), allow(unused_mut))]
24
    #[cfg_attr(feature = "encryption", allow(clippy::unnecessary_wraps))]
25
2253425
    pub fn open_tree<R: Root>(
26
2253425
        &mut self,
27
2253425
        name: &str,
28
2253425
        #[cfg(any(feature = "encryption", feature = "compression"))] vault: Option<TreeVault>,
29
2253425
    ) {
30
2253425
        if !self.trees_index_by_name.contains_key(name) {
31
1719795
            self.trees_index_by_name
32
1719795
                .insert(name.to_string(), self.trees.len());
33
1719795
            let mut tree = R::tree(name.to_string());
34

            
35
            #[cfg(any(feature = "encryption", feature = "compression"))]
36
1719795
            if let Some(vault) = vault {
37
744069
                tree = tree.with_vault(vault);
38
997389
            }
39

            
40
1719793
            self.trees.push(Box::new(tree));
41
533630
        }
42
2253423
    }
43

            
44
    #[allow(clippy::needless_pass_by_value)]
45
885578
    pub fn open_trees_for_document_change(
46
885578
        &mut self,
47
885578
        collection: &CollectionName,
48
885578
        schema: &Schematic,
49
885578
        #[cfg(any(feature = "encryption", feature = "compression"))] vault: Option<TreeVault>,
50
885578
    ) {
51
885578
        self.open_tree::<Versioned>(
52
885578
            &document_tree_name(collection),
53
885578
            #[cfg(any(feature = "encryption", feature = "compression"))]
54
885578
            vault.clone(),
55
885578
        );
56

            
57
885578
        if let Some(views) = schema.views_in_collection(collection) {
58
1682597
            for view in views {
59
1266700
                let view_name = view.view_name();
60
1266700
                if view.unique() {
61
101085
                    self.open_tree::<Unversioned>(
62
101085
                        &view_document_map_tree_name(&view_name),
63
101085
                        #[cfg(any(feature = "encryption", feature = "compression"))]
64
101085
                        vault.clone(),
65
101085
                    );
66
101085
                    self.open_tree::<Unversioned>(
67
101085
                        &view_entries_tree_name(&view_name),
68
101085
                        #[cfg(any(feature = "encryption", feature = "compression"))]
69
101085
                        vault.clone(),
70
101085
                    );
71
1165615
                } else {
72
1165615
                    self.open_tree::<Unversioned>(
73
1165615
                        &view_invalidated_docs_tree_name(&view_name),
74
1165615
                        #[cfg(any(feature = "encryption", feature = "compression"))]
75
1165615
                        vault.clone(),
76
1165615
                    );
77
1165615
                }
78
            }
79
469681
        }
80
885578
    }
81
}