inc_complete/db/
handle.rs

1use std::collections::BTreeSet;
2
3use crate::{
4    Cell, Computation, Db, Storage,
5    accumulate::{Accumulate, Accumulated},
6    storage::StorageFor,
7};
8
9use super::DbGet;
10
11/// A handle to the database during some operation.
12///
13/// This wraps calls to the Db so that any `get` calls
14/// will be automatically registered as dependencies of
15/// the current operation.
16pub struct DbHandle<'db, S> {
17    db: &'db Db<S>,
18    current_operation: Cell,
19}
20
21impl<'db, S> DbHandle<'db, S> {
22    pub(crate) fn new(db: &'db Db<S>, current_operation: Cell) -> Self {
23        // We're re-running a cell so remove any past dependencies
24        let mut cell = db.cells.get_mut(&current_operation).unwrap();
25
26        cell.dependencies.clear();
27        cell.dependency_set.clear();
28
29        Self {
30            db,
31            current_operation,
32        }
33    }
34
35    /// Retrieve an immutable reference to this `Db`'s storage
36    ///
37    /// Note that any mutations made to the storage using this are _not_ tracked by the database!
38    /// Using this incorrectly may break correctness!
39    pub fn storage(&self) -> &S {
40        self.db.storage()
41    }
42}
43
44impl<S: Storage> DbHandle<'_, S> {
45    /// Locking behavior: This function locks the cell corresponding to the given computation. This
46    /// can cause a deadlock if the computation recursively depends on itself.
47    pub fn get<C: Computation>(&self, compute: C) -> C::Output
48    where
49        S: StorageFor<C>,
50    {
51        // Register the dependency
52        let dependency = self.db.get_or_insert_cell(compute);
53        self.update_and_register_dependency::<C>(dependency);
54
55        // Fetch the current value of the dependency, running it if out of date
56        self.db.get_with_cell(dependency)
57    }
58
59    /// Registers the given cell as a dependency, running it and updating any required metadata
60    fn update_and_register_dependency<C: Computation>(&self, dependency: Cell) {
61        self.update_and_register_dependency_inner(dependency);
62    }
63
64    fn update_and_register_dependency_inner(&self, dependency: Cell) {
65        let mut cell = self.db.cells.get_mut(&self.current_operation).unwrap();
66
67        // Storing dependency_set separately takes a hit to memory usage but is worth
68        // it for extra runtime performance on this check
69        if cell.dependency_set.insert(dependency) {
70            cell.dependencies.push(dependency);
71        }
72        drop(cell);
73
74        // Run the computation to update its dependencies before we query them afterward
75        self.db.update_cell(dependency);
76    }
77
78    /// Accumulate an item in the current computation. This item can be retrieved along
79    /// with all other accumulated items in this computation and its dependencies via
80    /// a call to `get_accumulated`.
81    ///
82    /// This is most often used for operations like pushing diagnostics or logs.
83    pub fn accumulate<Item>(&self, item: Item)
84    where
85        S: Accumulate<Item>,
86    {
87        self.storage().accumulate(self.current_operation, item);
88    }
89
90    /// Retrieve an accumulated value in a container of the user's choice.
91    /// This will return all the accumulated items after the given computation.
92    ///
93    /// This is most often used for operations like retrieving diagnostics or logs.
94    pub fn get_accumulated<Item, C>(&self, compute: C) -> BTreeSet<Item>
95    where
96        C: Computation,
97        Item: 'static + Ord,
98        S: StorageFor<Accumulated<Item>> + StorageFor<C> + Accumulate<Item>,
99    {
100        let dependency = self.db.get_or_insert_cell(compute);
101        self.get_accumulated_with_cell::<Item>(dependency)
102    }
103
104    /// Retrieve an accumulated value in a container of the user's choice.
105    /// This will return all the accumulated items after the given computation.
106    ///
107    /// This is the implementation of the publically accessible `db.get(Accumulated::<Item>(MyComputation))`.
108    ///
109    /// This is most often used for operations like retrieving diagnostics or logs.
110    pub(crate) fn get_accumulated_with_cell<Item>(&self, cell_id: Cell) -> BTreeSet<Item>
111    where
112        Item: 'static + Ord,
113        S: StorageFor<Accumulated<Item>> + Accumulate<Item>,
114    {
115        self.update_and_register_dependency_inner(cell_id);
116        let dependencies = self.db.with_cell(cell_id, |cell| cell.dependencies.clone());
117
118        // Collect `Accumulator` results from each dependency. This should also ensure we
119        // rerun this if any dependency changes, even if `cell_id` is updated such that it
120        // uses different dependencies but its output remains the same.
121        let computation_id = Accumulated::<Item>::computation_id();
122        let mut result: BTreeSet<Item> = dependencies
123            .into_iter()
124            // Filter out `Accumulated<Item>` cells from the dep list. They exist for staleness
125            // tracking only and must not be traversed for value collection, or we'd get duplicates.
126            .filter(|&dep| self.db.with_cell(dep, |cell| cell.computation_id) != computation_id)
127            .flat_map(|dependency| self.get(Accumulated::<Item>::new(dependency)))
128            .collect();
129
130        result.extend(self.storage().get_accumulated::<Vec<Item>>(cell_id));
131        result
132    }
133}
134
135impl<'db, S, C> DbGet<C> for DbHandle<'db, S>
136where
137    C: Computation,
138    S: Storage + StorageFor<C>,
139{
140    fn get(&self, key: C) -> C::Output {
141        self.get(key)
142    }
143}
inc_complete/db/handle.rs

inc_complete/db/
handle.rs
