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
554701
#[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
2112654
    pub fn open_tree<R: Root>(
26
2112654
        &mut self,
27
2112654
        name: &str,
28
2112654
        #[cfg(any(feature = "encryption", feature = "compression"))] vault: Option<TreeVault>,
29
2112654
    ) {
30
2112654
        if !self.trees_index_by_name.contains_key(name) {
31
1727794
            self.trees_index_by_name
32
1727794
                .insert(name.to_string(), self.trees.len());
33
1727794
            let mut tree = R::tree(name.to_string());
34

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

            
40
1727764
            self.trees.push(Box::new(tree));
41
384860
        }
42
2112624
    }
43

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

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