inc_complete/db/
handle.rs

1use petgraph::visit::EdgeRef;
2
3use crate::{
4    Cell, Db, OutputType, Storage,
5    storage::{ComputationId, StorageFor},
6};
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 mut Db<S>,
15    current_operation: Cell,
16}
17
18impl<'db, S> DbHandle<'db, S> {
19    pub(crate) fn new(db: &'db mut Db<S>, current_operation: Cell) -> Self {
20        // We're re-running a cell so remove any past dependencies
21        let edges = db
22            .cells
23            .edges(current_operation.index())
24            .map(|edge| edge.id())
25            .collect::<Vec<_>>();
26
27        for edge in edges {
28            db.cells.remove_edge(edge);
29        }
30        Self {
31            db,
32            current_operation,
33        }
34    }
35
36    /// Retrieve an immutable reference to this `Db`'s storage
37    pub fn storage(&self) -> &S {
38        self.db.storage()
39    }
40
41    /// Retrieve a mutable reference to the internal database's storage.
42    ///
43    /// Note that any mutations made to the storage using this are _not_ tracked by the database!
44    /// Using this incorrectly may break correctness!
45    pub fn storage_mut(&mut self) -> &mut S {
46        self.db.storage_mut()
47    }
48}
49
50impl<'db, S: Storage> DbHandle<'db, S> {
51    pub fn get<C: OutputType + ComputationId>(&mut self, compute: C) -> &C::Output
52    where
53        S: StorageFor<C>,
54    {
55        // Register the dependency
56        let dependency = self.db.get_or_insert_cell(compute);
57        self.db
58            .cells
59            .update_edge(self.current_operation.index(), dependency.index(), ());
60
61        // Fetch the current value of the dependency
62        self.db.get_with_cell(dependency)
63    }
64}