[refactor] Simplify fetch_hot (#792)

* [refactor] Simplify `fetch_hot`

* Review feedback

* Reduce diff

Author: MichaReiser

src/function/fetch.rs

@@ -4,8 +4,8 @@ use crate::function::{Configuration, IngredientImpl, VerifyResult};
 use crate::runtime::StampedValue;
 use crate::table::sync::ClaimResult;
 use crate::zalsa::{MemoIngredientIndex, Zalsa, ZalsaDatabase};
-use crate::zalsa_local::QueryRevisions;
-use crate::{AsDynDatabase as _, Id};
+use crate::zalsa_local::{ActiveQueryGuard, QueryRevisions, ZalsaLocal};
+use crate::{AsDynDatabase as _, DatabaseKeyIndex, Id};
 
 impl<C> IngredientImpl<C>
 where
@@ -87,17 +87,15 @@ where
 
         let shallow_update = self.shallow_verify_memo(zalsa, database_key_index, memo)?;
 
-        if self.validate_may_be_provisional(db, zalsa, database_key_index, memo)
-            || self.validate_same_iteration(db, database_key_index, memo)
-        {
+        if self.validate_may_be_provisional(db, zalsa, database_key_index, memo) {
             self.update_shallow(db, zalsa, database_key_index, memo, shallow_update);
 
             // SAFETY: memo is present in memo_map and we have verified that it is
             // still valid for the current revision.
-            return unsafe { Some(self.extend_memo_lifetime(memo)) };
+            unsafe { Some(self.extend_memo_lifetime(memo)) }
+        } else {
+            None
         }
-
-        None
     }
 
     fn fetch_cold<'db>(
@@ -173,15 +171,14 @@ where
             ClaimResult::Claimed(guard) => guard,
         };
 
-        // Push the query on the stack.
-        let active_query = db.zalsa_local().push_query(database_key_index, 0);
+        let mut active_query = LazyActiveQueryGuard::new(database_key_index);
 
         // Now that we've claimed the item, check again to see if there's a "hot" value.
         let opt_old_memo = self.get_memo_from_table_for(zalsa, id, memo_ingredient_index);
         if let Some(old_memo) = opt_old_memo {
             if old_memo.value.is_some() {
                 if let VerifyResult::Unchanged(_, cycle_heads) =
-                    self.deep_verify_memo(db, zalsa, old_memo, &active_query)
+                    self.deep_verify_memo(db, zalsa, old_memo, &mut active_query)
                 {
                     if cycle_heads.is_empty() {
                         // SAFETY: memo is present in memo_map and we have verified that it is
@@ -192,8 +189,36 @@ where
             }
         }
 
-        let memo = self.execute(db, active_query, opt_old_memo);
+        let memo = self.execute(db, active_query.into_inner(db.zalsa_local()), opt_old_memo);
 
         Some(memo)
     }
 }
+
+pub(super) struct LazyActiveQueryGuard<'me> {
+    guard: Option<ActiveQueryGuard<'me>>,
+    database_key_index: DatabaseKeyIndex,
+}
+
+impl<'me> LazyActiveQueryGuard<'me> {
+    pub(super) fn new(database_key_index: DatabaseKeyIndex) -> Self {
+        Self {
+            guard: None,
+            database_key_index,
+        }
+    }
+
+    pub(super) const fn database_key_index(&self) -> DatabaseKeyIndex {
+        self.database_key_index
+    }
+
+    pub(super) fn guard(&mut self, zalsa_local: &'me ZalsaLocal) -> &ActiveQueryGuard<'me> {
+        self.guard
+            .get_or_insert_with(|| zalsa_local.push_query(self.database_key_index, 0))
+    }
+
+    pub(super) fn into_inner(self, zalsa_local: &'me ZalsaLocal) -> ActiveQueryGuard<'me> {
+        self.guard
+            .unwrap_or_else(|| zalsa_local.push_query(self.database_key_index, 0))
+    }
+}

src/function/maybe_changed_after.rs

@@ -2,12 +2,13 @@ use std::sync::atomic::Ordering;
 
 use crate::accumulator::accumulated_map::InputAccumulatedValues;
 use crate::cycle::{CycleHeads, CycleRecoveryStrategy};
+use crate::function::fetch::LazyActiveQueryGuard;
 use crate::function::memo::Memo;
 use crate::function::{Configuration, IngredientImpl};
 use crate::key::DatabaseKeyIndex;
 use crate::table::sync::ClaimResult;
 use crate::zalsa::{MemoIngredientIndex, Zalsa, ZalsaDatabase};
-use crate::zalsa_local::{ActiveQueryGuard, QueryEdge, QueryOrigin};
+use crate::zalsa_local::{QueryEdge, QueryOrigin};
 use crate::{AsDynDatabase as _, Id, Revision};
 
 /// Result of memo validation.
@@ -135,9 +136,9 @@ where
         );
 
         // Check if the inputs are still valid. We can just compare `changed_at`.
-        let active_query = db.zalsa_local().push_query(database_key_index, 0);
+        let mut active_query = LazyActiveQueryGuard::new(database_key_index);
         if let VerifyResult::Unchanged(_, cycle_heads) =
-            self.deep_verify_memo(db, zalsa, old_memo, &active_query)
+            self.deep_verify_memo(db, zalsa, old_memo, &mut active_query)
         {
             return Some(if old_memo.revisions.changed_at > revision {
                 VerifyResult::Changed
@@ -151,7 +152,11 @@ where
         // backdated. In that case, although we will have computed a new memo,
         // the value has not logically changed.
         if old_memo.value.is_some() {
-            let memo = self.execute(db, active_query, Some(old_memo));
+            let memo = self.execute(
+                db,
+                active_query.into_inner(db.zalsa_local()),
+                Some(old_memo),
+            );
             let changed_at = memo.revisions.changed_at;
 
             return Some(if changed_at > revision {
@@ -317,14 +322,14 @@ where
     /// Takes an [`ActiveQueryGuard`] argument because this function recursively
     /// walks dependencies of `old_memo` and may even execute them to see if their
     /// outputs have changed.
-    pub(super) fn deep_verify_memo(
+    pub(super) fn deep_verify_memo<'db>(
         &self,
-        db: &C::DbView,
+        db: &'db C::DbView,
         zalsa: &Zalsa,
         old_memo: &Memo<C::Output<'_>>,
-        active_query: &ActiveQueryGuard<'_>,
+        active_query: &mut LazyActiveQueryGuard<'db>,
     ) -> VerifyResult {
-        let database_key_index = active_query.database_key_index;
+        let database_key_index = active_query.database_key_index();
 
         tracing::debug!(
             "{database_key_index:?}: deep_verify_memo(old_memo = {old_memo:#?})",
@@ -333,9 +338,10 @@ where
 
         let shallow_update = self.shallow_verify_memo(zalsa, database_key_index, old_memo);
         let shallow_update_possible = shallow_update.is_some();
-
         if let Some(shallow_update) = shallow_update {
-            if self.validate_may_be_provisional(db, zalsa, database_key_index, old_memo) {
+            if self.validate_may_be_provisional(db, zalsa, database_key_index, old_memo)
+                || self.validate_same_iteration(db, database_key_index, old_memo)
+            {
                 self.update_shallow(db, zalsa, database_key_index, old_memo, shallow_update);
 
                 return VerifyResult::unchanged();
@@ -365,11 +371,14 @@ where
             }
             QueryOrigin::Derived(edges) => {
                 let is_provisional = old_memo.may_be_provisional();
+
                 // If the value is from the same revision but is still provisional, consider it changed
                 if shallow_update_possible && is_provisional {
                     return VerifyResult::Changed;
                 }
 
+                let _guard = active_query.guard(db.zalsa_local());
+
                 let mut cycle_heads = CycleHeads::default();
                 'cycle: loop {
                     // Fully tracked inputs? Iterate over the inputs and check them, one by one.