Track cycle function dependenciees as part of the cyclic query

Author: MichaReiser

src/active_query.rs

@@ -91,6 +91,10 @@ impl ActiveQuery {
             .mark_all_active(active_tracked_ids.iter().copied());
     }
 
+    pub(super) fn take_cycle_heads(&mut self) -> CycleHeads {
+        std::mem::take(&mut self.cycle_heads)
+    }
+
     pub(super) fn add_read(
         &mut self,
         input: DatabaseKeyIndex,

src/cycle.rs

@@ -490,4 +490,8 @@ impl<'db> ProvisionalStatus<'db> {
             _ => empty_cycle_heads(),
         }
     }
+
+    pub(crate) const fn is_provisional(&self) -> bool {
+        matches!(self, ProvisionalStatus::Provisional { .. }
+    }
 }

src/function/execute.rs

@@ -56,20 +56,25 @@ where
         });
 
         let (new_value, mut completed_query) = match C::CYCLE_STRATEGY {
-            CycleRecoveryStrategy::Panic => Self::execute_query(
-                db,
-                zalsa,
-                zalsa_local.push_query(database_key_index, IterationCount::initial()),
-                opt_old_memo,
-            ),
+            CycleRecoveryStrategy::Panic => {
+                let (new_value, active_query) = Self::execute_query(
+                    db,
+                    zalsa,
+                    zalsa_local.push_query(database_key_index, IterationCount::initial()),
+                    opt_old_memo,
+                );
+                (new_value, active_query.pop())
+            }
             CycleRecoveryStrategy::FallbackImmediate => {
-                let (mut new_value, mut completed_query) = Self::execute_query(
+                let (mut new_value, active_query) = Self::execute_query(
                     db,
                     zalsa,
                     zalsa_local.push_query(database_key_index, IterationCount::initial()),
                     opt_old_memo,
                 );
 
+                let mut completed_query = active_query.pop();
+
                 if let Some(cycle_heads) = completed_query.revisions.cycle_heads_mut() {
                     // Did the new result we got depend on our own provisional value, in a cycle?
                     if cycle_heads.contains(&database_key_index) {
@@ -196,9 +201,10 @@ where
 
         let _poison_guard =
             PoisonProvisionalIfPanicking::new(self, zalsa, id, memo_ingredient_index);
-        let mut active_query = zalsa_local.push_query(database_key_index, iteration_count);
 
         let (new_value, completed_query) = loop {
+            let active_query = zalsa_local.push_query(database_key_index, iteration_count);
+
             // Tracked struct ids that existed in the previous revision
             // but weren't recreated in the last iteration. It's important that we seed the next
             // query with these ids because the query might re-create them as part of the next iteration.
@@ -207,21 +213,22 @@ where
             // if they aren't recreated when reaching the final iteration.
             active_query.seed_tracked_struct_ids(&last_stale_tracked_ids);
 
-            let (mut new_value, mut completed_query) = Self::execute_query(
+            let (mut new_value, mut active_query) = Self::execute_query(
                 db,
                 zalsa,
                 active_query,
                 last_provisional_memo.or(opt_old_memo),
             );
 
-            // If there are no cycle heads, break out of the loop (`cycle_heads_mut` returns `None` if the cycle head list is empty)
-            let Some(cycle_heads) = completed_query.revisions.cycle_heads_mut() else {
+            // Take the cycle heads to not-fight-rust's-borrow-checker.
+            let mut cycle_heads = active_query.take_cycle_heads();
+
+            // If there are no cycle heads, break out of the loop.
+            if cycle_heads.is_empty() {
                 claim_guard.set_release_mode(ReleaseMode::SelfOnly);
-                break (new_value, completed_query);
-            };
+                break (new_value, active_query.pop());
+            }
 
-            // Take the cycle heads to not-fight-rust's-borrow-checker.
-            let mut cycle_heads = std::mem::take(cycle_heads);
             let mut missing_heads: SmallVec<[(DatabaseKeyIndex, IterationCount); 1]> =
                 SmallVec::new_const();
             let mut max_iteration_count = iteration_count;
@@ -252,6 +259,10 @@ where
                     .provisional_status(zalsa, head.database_key_index.key_index())
                     .expect("cycle head memo must have been created during the execution");
 
+                // A cycle head isn't allowed to be final because that would mean that this query
+                // dependent on the initial value (or last provisional)
+                assert!(provisional_status.is_provisional());
+
                 for nested_head in provisional_status.cycle_heads() {
                     let nested_as_tuple = (
                         nested_head.database_key_index,
@@ -288,6 +299,8 @@ where
                     claim_guard.set_release_mode(ReleaseMode::SelfOnly);
                 }
 
+                let mut completed_query = active_query.pop();
+                *completed_query.revisions.verified_final.get_mut() = false;
                 completed_query.revisions.set_cycle_heads(cycle_heads);
                 break (new_value, completed_query);
             }
@@ -359,8 +372,17 @@ where
                         this_converged = C::values_equal(&new_value, last_provisional_value);
                     }
                 }
+
+                let new_cycle_heads = active_query.take_cycle_heads();
+                for head in new_cycle_heads {
+                    if !cycle_heads.contains(&head.database_key_index) {
+                        panic!("Cycle recovery function for {database_key_index:?} introduced a cycle, depending on {:?}. This is not allowed.", head.database_key_index);
+                    }
+                }
             }
 
+            let mut completed_query = active_query.pop();
+
             if let Some(outer_cycle) = outer_cycle {
                 tracing::info!(
                     "Detected nested cycle {database_key_index:?}, iterate it as part of the outer cycle {outer_cycle:?}"
@@ -371,6 +393,7 @@ where
                 completed_query
                     .revisions
                     .set_cycle_converged(this_converged);
+                *completed_query.revisions.verified_final.get_mut() = false;
 
                 // Transfer ownership of this query to the outer cycle, so that it can claim it
                 // and other threads don't compete for the same lock.
@@ -409,9 +432,9 @@ where
                 }
 
                 *completed_query.revisions.verified_final.get_mut() = true;
-
                 break (new_value, completed_query);
             }
+            *completed_query.revisions.verified_final.get_mut() = false;
 
             // The fixpoint iteration hasn't converged. Iterate again...
             iteration_count = iteration_count.increment().unwrap_or_else(|| {
@@ -465,7 +488,6 @@ where
             last_provisional_memo = Some(new_memo);
 
             last_stale_tracked_ids = completed_query.stale_tracked_structs;
-            active_query = zalsa_local.push_query(database_key_index, iteration_count);
 
             continue;
         };
@@ -484,7 +506,7 @@ where
         zalsa: &'db Zalsa,
         active_query: ActiveQueryGuard<'db>,
         opt_old_memo: Option<&Memo<'db, C>>,
-    ) -> (C::Output<'db>, CompletedQuery) {
+    ) -> (C::Output<'db>, ActiveQueryGuard<'db>) {
         if let Some(old_memo) = opt_old_memo {
             // If we already executed this query once, then use the tracked-struct ids from the
             // previous execution as the starting point for the new one.
@@ -509,7 +531,7 @@ where
             C::id_to_input(zalsa, active_query.database_key_index.key_index()),
         );
 
-        (new_value, active_query.pop())
+        (new_value, active_query)
     }
 }
 

src/zalsa_local.rs

@@ -1202,6 +1202,18 @@ impl ActiveQueryGuard<'_> {
         }
     }
 
+    pub(crate) fn take_cycle_heads(&mut self) -> CycleHeads {
+        // SAFETY: We do not access the query stack reentrantly.
+        unsafe {
+            self.local_state.with_query_stack_unchecked_mut(|stack| {
+                #[cfg(debug_assertions)]
+                assert_eq!(stack.len(), self.push_len);
+                let frame = stack.last_mut().unwrap();
+                frame.take_cycle_heads()
+            })
+        }
+    }
+
     /// Invoked when the query has successfully completed execution.
     fn complete(self) -> CompletedQuery {
         // SAFETY: We do not access the query stack reentrantly.