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

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

            
40
1719375
            self.trees.push(Box::new(tree));
41
367910
        }
42
2087285
    }
43

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

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