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::{
14
        view_document_map_tree_name, view_entries_tree_name, view_invalidated_docs_tree_name,
15
        view_omitted_docs_tree_name,
16
    },
17
};
18

            
19
398850
#[derive(Default)]
20
pub(crate) struct OpenTrees {
21
    pub trees: Vec<Box<dyn AnyTreeRoot<AnyFile>>>,
22
    pub trees_index_by_name: HashMap<String, usize>,
23
}
24

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

            
38
            #[cfg(any(feature = "encryption", feature = "compression"))]
39
1295991
            if let Some(vault) = vault {
40
477110
                tree = tree.with_vault(vault);
41
829835
            }
42

            
43
1295991
            self.trees.push(Box::new(tree));
44
367710
        }
45
1663701
    }
46

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

            
60
592847
        if let Some(views) = schema.views_in_collection(collection) {
61
1234120
            for view in views {
62
925472
                let view_name = view.view_name();
63
925472
                if view.unique() {
64
72689
                    self.open_tree::<Unversioned>(
65
72689
                        &view_omitted_docs_tree_name(&view_name),
66
72689
                        #[cfg(any(feature = "encryption", feature = "compression"))]
67
72689
                        vault.clone(),
68
72689
                    );
69
72689
                    self.open_tree::<Unversioned>(
70
72689
                        &view_document_map_tree_name(&view_name),
71
72689
                        #[cfg(any(feature = "encryption", feature = "compression"))]
72
72689
                        vault.clone(),
73
72689
                    );
74
72689
                    self.open_tree::<Unversioned>(
75
72689
                        &view_entries_tree_name(&view_name),
76
72689
                        #[cfg(any(feature = "encryption", feature = "compression"))]
77
72689
                        vault.clone(),
78
72689
                    );
79
852785
                } else {
80
852785
                    self.open_tree::<Unversioned>(
81
852785
                        &view_invalidated_docs_tree_name(&view_name),
82
852785
                        #[cfg(any(feature = "encryption", feature = "compression"))]
83
852785
                        vault.clone(),
84
852785
                    );
85
852785
                }
86
            }
87
284201
        }
88
592849
    }
89
}