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        db.cells
22            .get_mut(&current_operation)
23            .unwrap()
24            .dependencies
25            .clear();
26
27        Self {
28            db,
29            current_operation,
30        }
31    }
32
33    /// Retrieve an immutable reference to this `Db`'s storage
34    ///
35    /// Note that any mutations made to the storage using this are _not_ tracked by the database!
36    /// Using this incorrectly may break correctness!
37    pub fn storage(&self) -> &S {
38        self.db.storage()
39    }
40}
41
42impl<S: Storage> DbHandle<'_, S> {
43    /// Locking behavior: This function locks the cell corresponding to the given computation. This
44    /// can cause a deadlock if the computation recursively depends on itself.
45    #[cfg(not(feature = "async"))]
46    pub fn get<C: OutputType + ComputationId>(&self, compute: C) -> C::Output
47    where
48        S: StorageFor<C>,
49    {
50        // Register the dependency
51        let dependency = self.db.get_or_insert_cell(compute);
52        let mut cell = self.db.cells.get_mut(&self.current_operation).unwrap();
53        cell.dependencies.push(dependency);
54        drop(cell);
55
56        // Fetch the current value of the dependency
57        self.db.get_with_cell(dependency)
58    }
59
60    #[cfg(feature = "async")]
61    pub fn get<C: OutputType + ComputationId>(
62        &self,
63        compute: C,
64    ) -> impl Future<Output = C::Output> + Send
65    where
66        S: StorageFor<C> + Sync,
67    {
68        // Register the dependency
69        let dependency = self.db.get_or_insert_cell(compute);
70        let mut cell = self.db.cells.get(&self.current_operation).unwrap();
71        cell.dependencies.push(dependency);
72        drop(cell);
73
74        // Fetch the current value of the dependency
75        self.db.get_with_cell(dependency)
76    }
77}
78
79#[cfg(not(feature = "async"))]
80impl<'db, S, C> DbGet<C> for DbHandle<'db, S>
81where
82    C: OutputType + ComputationId,
83    S: Storage + StorageFor<C>,
84{
85    fn get(&self, key: C) -> C::Output {
86        self.get(key)
87    }
88}
89
90#[cfg(feature = "async")]
91impl<'db, S, C> DbGet<C> for DbHandle<'db, S>
92where
93    C: OutputType + ComputationId,
94    S: Storage + StorageFor<C> + Sync,
95{
96    fn get(&self, key: C) -> impl Future<Output = C::Output> + Send {
97        self.get(key)
98    }
99}