perf: keep the fetch path generic

Author: MichaReiser

src/function/execute.rs

@@ -80,6 +80,28 @@ impl<'db, C: Configuration> QueryStateImpl<'db, C> {
     }
 }
 
+impl<C: Configuration> IngredientImpl<C> {
+    /// Executes this query through the shared query lifecycle and restores its typed memo.
+    pub(super) fn execute<'db>(
+        &'db self,
+        db: &'db C::DbView,
+        claim_guard: ClaimGuard<'db>,
+        opt_old_memo: Option<&'db Memo<'db, C>>,
+        memo_ingredient_index: MemoIngredientIndex,
+    ) -> Option<&'db Memo<'db, C>> {
+        let mut state = QueryStateImpl::new(self, db);
+        let memo = execute_erased(
+            &mut state,
+            claim_guard,
+            opt_old_memo.map(ErasedMemo::new),
+            memo_ingredient_index,
+            C::CYCLE_STRATEGY,
+        )?;
+        // SAFETY: `state` and all memos it returns belong to `IngredientImpl<C>`.
+        Some(unsafe { memo.downcast::<C>() })
+    }
+}
+
 impl<'db, C: Configuration> QueryState<'db> for QueryStateImpl<'db, C> {
     fn execute_query(&mut self, zalsa: &'db Zalsa, id: Id) {
         self.value = Some(C::execute(self.db, C::id_to_input(zalsa, id)));

src/function/fetch.rs

@@ -1,9 +1,8 @@
 use crate::cycle::{CycleRecoveryStrategy, IterationStamp};
-use crate::database::RawDatabase;
 use crate::function::eviction::EvictionPolicy;
-use crate::function::execute::{QueryState, QueryStateImpl, execute_erased};
+use crate::function::execute::{QueryState, QueryStateImpl};
 use crate::function::memo::{ErasedMemo, Memo};
-use crate::function::sync::{ClaimResult, SyncTable};
+use crate::function::sync::ClaimResult;
 use crate::function::{Configuration, IngredientImpl, Reentrancy};
 use crate::zalsa::{MemoIngredientIndex, Zalsa};
 use crate::zalsa_local::{QueryRevisions, ZalsaLocal};
@@ -112,134 +111,132 @@ where
         memo_ingredient_index: MemoIngredientIndex,
     ) -> Option<&'db Memo<'db, C>> {
         let database_key_index = self.database_key_index(id);
-        let mut state = QueryStateImpl::new(self, db);
-        let memo = fetch_cold_erased(
-            &mut state,
-            &self.sync_table,
-            zalsa,
-            zalsa_local,
-            db.into(),
-            database_key_index,
-            memo_ingredient_index,
-            C::CYCLE_STRATEGY,
-        )?;
-        // SAFETY: `state` and all memos it returns belong to `IngredientImpl<C>`.
-        Some(unsafe { memo.downcast::<C>() })
-    }
-}
+        // Try to claim this query: if someone else has claimed it already, go back and start again.
+        let claim_guard = match self
+            .sync_table
+            .try_claim(zalsa, zalsa_local, id, Reentrancy::Allow)
+        {
+            ClaimResult::Claimed(guard) => guard,
+            ClaimResult::Running(blocked_on) => {
+                let _ = blocked_on.block_on(zalsa);
+                return None;
+            }
+            ClaimResult::Cycle { .. } => {
+                return Some(self.fetch_cold_cycle(
+                    db,
+                    zalsa,
+                    zalsa_local,
+                    database_key_index,
+                    memo_ingredient_index,
+                ));
+            }
+        };
 
-#[allow(clippy::too_many_arguments)]
-fn fetch_cold_erased<'db>(
-    state: &mut dyn QueryState<'db>,
-    sync_table: &'db SyncTable,
-    zalsa: &'db Zalsa,
-    zalsa_local: &'db ZalsaLocal,
-    db: RawDatabase<'db>,
-    database_key_index: DatabaseKeyIndex,
-    memo_ingredient_index: MemoIngredientIndex,
-    strategy: CycleRecoveryStrategy,
-) -> Option<ErasedMemo<'db>> {
-    let id = database_key_index.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);
 
-    // Try to claim this query: if someone else has claimed it already, go back and start again.
-    let claim_guard = match sync_table.try_claim(zalsa, zalsa_local, id, Reentrancy::Allow) {
-        ClaimResult::Claimed(guard) => guard,
-        ClaimResult::Running(blocked_on) => {
-            let _ = blocked_on.block_on(zalsa);
-            return None;
-        }
-        ClaimResult::Cycle { .. } => {
-            return Some(fetch_cold_cycle_erased(
-                state,
-                zalsa,
-                zalsa_local,
-                database_key_index,
-                memo_ingredient_index,
-                strategy,
-            ));
+        if let Some(old_memo) = opt_old_memo {
+            if old_memo.value.is_some()
+                && old_memo
+                    .header
+                    .verify_memo(db.into(), &claim_guard, C::CYCLE_STRATEGY, true)
+            {
+                // SAFETY: The memo is present in the memo table, and we verified that it is valid
+                // for the current revision.
+                return unsafe { Some(self.extend_memo_lifetime(old_memo)) };
+            }
         }
-    };
 
-    // Now that we've claimed the item, check again to see if there's a hot value.
-    let opt_old_memo = state.get_memo(zalsa, id, memo_ingredient_index);
+        self.execute(db, claim_guard, opt_old_memo, memo_ingredient_index)
+    }
 
-    if let Some(old_memo) = opt_old_memo {
-        if old_memo.value_is_some && old_memo.verify_memo(db, &claim_guard, strategy, true) {
-            return Some(old_memo);
+    #[cold]
+    #[inline(never)]
+    fn fetch_cold_cycle<'db>(
+        &'db self,
+        db: &'db C::DbView,
+        zalsa: &'db Zalsa,
+        zalsa_local: &'db ZalsaLocal,
+        database_key_index: DatabaseKeyIndex,
+        memo_ingredient_index: MemoIngredientIndex,
+    ) -> &'db Memo<'db, C> {
+        match C::CYCLE_STRATEGY {
+            CycleRecoveryStrategy::Panic => fetch_cold_cycle_panic(zalsa_local, database_key_index),
+            CycleRecoveryStrategy::FallbackImmediate | CycleRecoveryStrategy::Fixpoint => {
+                let mut state = QueryStateImpl::new(self, db);
+                let memo = fetch_cold_cycle_recoverable_erased(
+                    &mut state,
+                    zalsa,
+                    database_key_index,
+                    memo_ingredient_index,
+                );
+                // SAFETY: `state` and all memos it returns belong to `IngredientImpl<C>`.
+                unsafe { memo.downcast::<C>() }
+            }
         }
     }
+}
 
-    execute_erased(
-        state,
-        claim_guard,
-        opt_old_memo,
-        memo_ingredient_index,
-        strategy,
-    )
+#[cold]
+fn fetch_cold_cycle_panic(zalsa_local: &ZalsaLocal, database_key_index: DatabaseKeyIndex) -> ! {
+    // SAFETY: We do not access the query stack reentrantly.
+    unsafe {
+        zalsa_local.with_query_stack_unchecked(|stack| {
+            panic!(
+                "dependency graph cycle when querying {database_key_index:#?}, \
+                set cycle_fn/cycle_initial to fixpoint iterate.\n\
+                Query stack:\n{stack:#?}",
+            );
+        })
+    }
 }
 
 #[cold]
-fn fetch_cold_cycle_erased<'db>(
+fn fetch_cold_cycle_recoverable_erased<'db>(
     state: &mut dyn QueryState<'db>,
     zalsa: &'db Zalsa,
-    zalsa_local: &'db ZalsaLocal,
     database_key_index: DatabaseKeyIndex,
     memo_ingredient_index: MemoIngredientIndex,
-    strategy: CycleRecoveryStrategy,
 ) -> ErasedMemo<'db> {
     let id = database_key_index.key_index();
 
-    match strategy {
-        // SAFETY: We do not access the query stack reentrantly.
-        CycleRecoveryStrategy::Panic => unsafe {
-            zalsa_local.with_query_stack_unchecked(|stack| {
-                panic!(
-                    "dependency graph cycle when querying {database_key_index:#?}, \
-                    set cycle_fn/cycle_initial to fixpoint iterate.\n\
-                    Query stack:\n{stack:#?}",
-                );
-            })
-        },
-        CycleRecoveryStrategy::Fixpoint | CycleRecoveryStrategy::FallbackImmediate => {
-            let cancellation_count = zalsa.runtime().cancellation_count();
-            // Don't validate provisional memos here: an existing value should be reused.
-            let current_memo = state
-                .get_memo(zalsa, id, memo_ingredient_index)
-                .filter(|memo| {
-                    memo.verified_at.load() == zalsa.current_revision()
-                        && memo.value_is_some
-                        && memo.revisions.iteration().cancellation_count() == cancellation_count
-                });
-
-            // Ideally, any current provisional value could be reused. Reusing a value that was
-            // not a cycle head in the last iteration would require inserting itself as a head,
-            // which in turn requires cloning the value or making the cycle-head list concurrent.
-            if let Some(memo) = current_memo
-                .filter(|memo| memo.revisions.cycle_heads().contains(&database_key_index))
-            {
-                memo.revisions
-                    .cycle_heads()
-                    .remove_all_except(database_key_index);
-
-                crate::tracing::debug!(
-                    "hit cycle at {database_key_index:#?}, \
-                            returning last provisional value: {:#?}",
-                    memo.revisions
-                );
-                return memo;
-            }
-
-            crate::tracing::debug!(
-                "hit cycle at {database_key_index:#?}, \
-                inserting and returning fixpoint initial value"
-            );
-
-            let iteration = current_memo
-                .map(|memo| memo.revisions.iteration())
-                .unwrap_or_else(|| IterationStamp::initial(cancellation_count));
-            let revisions = QueryRevisions::fixpoint_initial(database_key_index, iteration);
-            state.use_fallback(zalsa, id);
-            state.insert_memo(zalsa, id, revisions, memo_ingredient_index)
-        }
+    let cancellation_count = zalsa.runtime().cancellation_count();
+    // Don't validate provisional memos here: an existing value should be reused.
+    let current_memo = state
+        .get_memo(zalsa, id, memo_ingredient_index)
+        .filter(|memo| {
+            memo.verified_at.load() == zalsa.current_revision()
+                && memo.value_is_some
+                && memo.revisions.iteration().cancellation_count() == cancellation_count
+        });
+
+    // Ideally, any current provisional value could be reused. Reusing a value that was not a
+    // cycle head in the last iteration would require inserting itself as a head, which in turn
+    // requires cloning the value or making the cycle-head list concurrent.
+    if let Some(memo) =
+        current_memo.filter(|memo| memo.revisions.cycle_heads().contains(&database_key_index))
+    {
+        memo.revisions
+            .cycle_heads()
+            .remove_all_except(database_key_index);
+
+        crate::tracing::debug!(
+            "hit cycle at {database_key_index:#?}, \
+                    returning last provisional value: {:#?}",
+            memo.revisions
+        );
+        return memo;
     }
+
+    crate::tracing::debug!(
+        "hit cycle at {database_key_index:#?}, \
+        inserting and returning fixpoint initial value"
+    );
+
+    let iteration = current_memo
+        .map(|memo| memo.revisions.iteration())
+        .unwrap_or_else(|| IterationStamp::initial(cancellation_count));
+    let revisions = QueryRevisions::fixpoint_initial(database_key_index, iteration);
+    state.use_fallback(zalsa, id);
+    state.insert_memo(zalsa, id, revisions, memo_ingredient_index)
 }