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(&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    #[cfg(not(feature = "async"))]
44    pub fn get<C: OutputType + ComputationId>(&self, compute: C) -> C::Output
45    where
46        S: StorageFor<C>,
47    {
48        // Register the dependency
49        let dependency = self.db.get_or_insert_cell(compute);
50        let mut cell = self.db.cells.get(&self.current_operation).unwrap();
51        cell.dependencies.push(dependency);
52        drop(cell);
53
54        // Fetch the current value of the dependency
55        self.db.get_with_cell(dependency)
56    }
57
58    #[cfg(feature = "async")]
59    pub fn get<C: OutputType + ComputationId>(
60        &self,
61        compute: C,
62    ) -> impl Future<Output = C::Output> + Send
63    where
64        S: StorageFor<C> + Sync,
65    {
66        // Register the dependency
67        let dependency = self.db.get_or_insert_cell(compute);
68        let mut cell = self.db.cells.get(&self.current_operation).unwrap();
69        cell.dependencies.push(dependency);
70        drop(cell);
71
72        // Fetch the current value of the dependency
73        self.db.get_with_cell(dependency)
74    }
75}
76
77#[cfg(not(feature = "async"))]
78impl<'db, S, C> DbGet<C> for DbHandle<'db, S> where
79    C: OutputType + ComputationId,
80    S: Storage + StorageFor<C>
81{
82    fn get(&self, key: C) -> C::Output {
83        self.get(key)
84    }
85}
86
87#[cfg(feature = "async")]
88impl<'db, S, C> DbGet<C> for DbHandle<'db, S> where
89    C: OutputType + ComputationId,
90    S: Storage + StorageFor<C> + Sync
91{
92    fn get(&self, key: C) -> impl Future<Output = C::Output> + Send {
93        self.get(key)
94    }
95}