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