inc_complete/db/
handle.rs

1use crate::{
2    Cell, Db, OutputType, Storage,
3    storage::{ComputationId, StorageFor},
4};
5
6use super::DbGet;
7
8/// A handle to the database during some operation.
9///
10/// This wraps calls to the Db so that any `get` calls
11/// will be automatically registered as dependencies of
12/// the current operation.
13pub struct DbHandle<'db, S> {
14    db: &'db Db<S>,
15    current_operation: Cell,
16}
17
18impl<'db, S> DbHandle<'db, S> {
19    pub(crate) fn new(db: &'db Db<S>, current_operation: Cell) -> Self {
20        // We're re-running a cell so remove any past dependencies
21        let mut cell = db.cells
22            .get_mut(&current_operation)
23            .unwrap();
24        
25        cell.dependencies.clear();
26        cell.input_dependencies.clear();
27
28        Self {
29            db,
30            current_operation,
31        }
32    }
33
34    /// Retrieve an immutable reference to this `Db`'s storage
35    ///
36    /// Note that any mutations made to the storage using this are _not_ tracked by the database!
37    /// Using this incorrectly may break correctness!
38    pub fn storage(&self) -> &S {
39        self.db.storage()
40    }
41}
42
43impl<S: Storage> DbHandle<'_, S> {
44    /// Locking behavior: This function locks the cell corresponding to the given computation. This
45    /// can cause a deadlock if the computation recursively depends on itself.
46    #[cfg(not(feature = "async"))]
47    pub fn get<C: OutputType + ComputationId>(&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        let mut cell = self.db.cells.get_mut(&self.current_operation).unwrap();
54
55        // If `dependency` is an input it must be remembered both as a dependency
56        // and as an input dependency. Otherwise we cannot differentiate between
57        // computations which directly depend on inputs and those that only indirectly
58        // depend on them.
59        cell.dependencies.push(dependency);
60        if C::IS_INPUT {
61            cell.input_dependencies.insert(dependency);
62        }
63
64        drop(cell);
65
66        // Fetch the current value of the dependency
67        let output = self.db.get_with_cell(dependency);
68
69        let dependency = self.db.cells.get(&dependency).unwrap();
70        let dependency_inputs = dependency.input_dependencies.clone();
71        drop(dependency);
72
73        let mut cell = self.db.cells.get_mut(&self.current_operation).unwrap();
74        for input in dependency_inputs {
75            cell.input_dependencies.insert(input);
76        }
77
78        output
79    }
80
81    #[cfg(feature = "async")]
82    pub fn get<C: OutputType + ComputationId>(
83        &self,
84        compute: C,
85    ) -> impl Future<Output = C::Output> + Send
86    where
87        S: StorageFor<C> + Sync,
88    {
89        // Register the dependency
90        let dependency = self.db.get_or_insert_cell(compute);
91        let mut cell = self.db.cells.get(&self.current_operation).unwrap();
92        cell.dependencies.push(dependency);
93        drop(cell);
94
95        // Fetch the current value of the dependency
96        self.db.get_with_cell(dependency)
97    }
98}
99
100#[cfg(not(feature = "async"))]
101impl<'db, S, C> DbGet<C> for DbHandle<'db, S>
102where
103    C: OutputType + ComputationId,
104    S: Storage + StorageFor<C>,
105{
106    fn get(&self, key: C) -> C::Output {
107        self.get(key)
108    }
109}
110
111#[cfg(feature = "async")]
112impl<'db, S, C> DbGet<C> for DbHandle<'db, S>
113where
114    C: OutputType + ComputationId,
115    S: Storage + StorageFor<C> + Sync,
116{
117    fn get(&self, key: C) -> impl Future<Output = C::Output> + Send {
118        self.get(key)
119    }
120}