allow reuse of cached provisional memos within the same cycle iteration (#786)

* test for caching provisional values

* add iteration-count to cycle heads

* CycleHeads insert/extend checks iteration count match

* update iteration count in cycle heads

* all tests passing

* remove debug prints

* just walk active query stack once

* switch to tracking active cycle iterations on ZalsaLocal

* Revert "switch to tracking active cycle iterations on ZalsaLocal"

This reverts commit 4ea3d85.

* Revert "just walk active query stack once"

This reverts commit 2d79486.

* make ActiveQuery::iteration_count private with accessor

* iterate active query stack in reverse

* use tracing::trace! in hot path

* try a cold annotation on validate_same_iteration

* Revert "try a cold annotation on validate_same_iteration"

This reverts commit 49ceb84.

Author: carljm

src/active_query.rs

@@ -60,6 +60,9 @@ pub(crate) struct ActiveQuery {
 
     /// Provisional cycle results that this query depends on.
     cycle_heads: CycleHeads,
+
+    /// If this query is a cycle head, iteration count of that cycle.
+    iteration_count: u32,
 }
 
 impl ActiveQuery {
@@ -126,10 +129,14 @@ impl ActiveQuery {
             changed_at: self.changed_at,
         }
     }
+
+    pub(super) fn iteration_count(&self) -> u32 {
+        self.iteration_count
+    }
 }
 
 impl ActiveQuery {
-    fn new(database_key_index: DatabaseKeyIndex) -> Self {
+    fn new(database_key_index: DatabaseKeyIndex, iteration_count: u32) -> Self {
         ActiveQuery {
             database_key_index,
             durability: Durability::MAX,
@@ -141,6 +148,7 @@ impl ActiveQuery {
             accumulated: Default::default(),
             accumulated_inputs: Default::default(),
             cycle_heads: Default::default(),
+            iteration_count,
         }
     }
 
@@ -156,6 +164,7 @@ impl ActiveQuery {
             ref mut accumulated,
             accumulated_inputs,
             ref mut cycle_heads,
+            iteration_count: _,
         } = self;
 
         let edges = QueryEdges::new(input_outputs.drain(..));
@@ -196,15 +205,17 @@ impl ActiveQuery {
             accumulated,
             accumulated_inputs: _,
             cycle_heads,
+            iteration_count,
         } = self;
         input_outputs.clear();
         disambiguator_map.clear();
         tracked_struct_ids.clear();
         accumulated.clear();
         *cycle_heads = Default::default();
+        *iteration_count = 0;
     }
 
-    fn reset_for(&mut self, new_database_key_index: DatabaseKeyIndex) {
+    fn reset_for(&mut self, new_database_key_index: DatabaseKeyIndex, new_iteration_count: u32) {
         let Self {
             database_key_index,
             durability,
@@ -216,12 +227,14 @@ impl ActiveQuery {
             accumulated,
             accumulated_inputs,
             cycle_heads,
+            iteration_count,
         } = self;
         *database_key_index = new_database_key_index;
         *durability = Durability::MAX;
         *changed_at = Revision::start();
         *untracked_read = false;
         *accumulated_inputs = Default::default();
+        *iteration_count = new_iteration_count;
         debug_assert!(
             input_outputs.is_empty(),
             "`ActiveQuery::clear` or `ActiveQuery::into_revisions` should've been called"
@@ -266,11 +279,16 @@ impl ops::DerefMut for QueryStack {
 }
 
 impl QueryStack {
-    pub(crate) fn push_new_query(&mut self, database_key_index: DatabaseKeyIndex) {
+    pub(crate) fn push_new_query(
+        &mut self,
+        database_key_index: DatabaseKeyIndex,
+        iteration_count: u32,
+    ) {
         if self.len < self.stack.len() {
-            self.stack[self.len].reset_for(database_key_index);
+            self.stack[self.len].reset_for(database_key_index, iteration_count);
         } else {
-            self.stack.push(ActiveQuery::new(database_key_index));
+            self.stack
+                .push(ActiveQuery::new(database_key_index, iteration_count));
         }
         self.len += 1;
     }

src/cycle.rs

@@ -86,12 +86,19 @@ pub enum CycleRecoveryStrategy {
 /// A "cycle head" is the query at which we encounter a cycle; that is, if A -> B -> C -> A, then A
 /// would be the cycle head. It returns an "initial value" when the cycle is encountered (if
 /// fixpoint iteration is enabled for that query), and then is responsible for re-iterating the
-/// cycle until it converges. Any provisional value generated by any query in the cycle will track
-/// the cycle head(s) (can be plural in case of nested cycles) representing the cycles it is part
-/// of. This struct tracks these cycle heads.
+/// cycle until it converges.
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
+pub struct CycleHead {
+    pub database_key_index: DatabaseKeyIndex,
+    pub iteration_count: u32,
+}
+
+/// Any provisional value generated by any query in a cycle will track the cycle head(s) (can be
+/// plural in case of nested cycles) representing the cycles it is part of, and the current
+/// iteration count for each cycle head. This struct tracks these cycle heads.
 #[derive(Clone, Debug, Default)]
 #[allow(clippy::box_collection)]
-pub struct CycleHeads(Option<Box<Vec<DatabaseKeyIndex>>>);
+pub struct CycleHeads(Option<Box<Vec<CycleHead>>>);
 
 impl CycleHeads {
     pub(crate) fn is_empty(&self) -> bool {
@@ -100,15 +107,25 @@ impl CycleHeads {
         self.0.is_none()
     }
 
+    pub(crate) fn initial(database_key_index: DatabaseKeyIndex) -> Self {
+        Self(Some(Box::new(vec![CycleHead {
+            database_key_index,
+            iteration_count: 0,
+        }])))
+    }
+
     pub(crate) fn contains(&self, value: &DatabaseKeyIndex) -> bool {
-        self.0.as_ref().is_some_and(|heads| heads.contains(value))
+        self.into_iter()
+            .any(|head| head.database_key_index == *value)
     }
 
     pub(crate) fn remove(&mut self, value: &DatabaseKeyIndex) -> bool {
         let Some(cycle_heads) = &mut self.0 else {
             return false;
         };
-        let found = cycle_heads.iter().position(|&head| head == *value);
+        let found = cycle_heads
+            .iter()
+            .position(|&head| head.database_key_index == *value);
         let Some(found) = found else { return false };
         cycle_heads.swap_remove(found);
         if cycle_heads.is_empty() {
@@ -117,43 +134,63 @@ impl CycleHeads {
         true
     }
 
+    pub(crate) fn update_iteration_count(
+        &mut self,
+        cycle_head_index: DatabaseKeyIndex,
+        new_iteration_count: u32,
+    ) {
+        if let Some(cycle_head) = self.0.as_mut().and_then(|cycle_heads| {
+            cycle_heads
+                .iter_mut()
+                .find(|cycle_head| cycle_head.database_key_index == cycle_head_index)
+        }) {
+            cycle_head.iteration_count = new_iteration_count;
+        }
+    }
+
     #[inline]
-    pub(crate) fn insert_into(self, cycle_heads: &mut Vec<DatabaseKeyIndex>) {
+    pub(crate) fn insert_into(self, cycle_heads: &mut Vec<CycleHead>) {
         if let Some(heads) = self.0 {
-            for head in *heads {
-                if !cycle_heads.contains(&head) {
-                    cycle_heads.push(head);
-                }
-            }
+            insert_into_impl(&heads, cycle_heads);
         }
     }
 
     pub(crate) fn extend(&mut self, other: &Self) {
         if let Some(other) = &other.0 {
             let heads = &mut **self.0.get_or_insert_with(|| Box::new(Vec::new()));
-            heads.reserve(other.len());
-            other.iter().for_each(|&head| {
-                if !heads.contains(&head) {
-                    heads.push(head);
-                }
-            });
+            insert_into_impl(other, heads);
+        }
+    }
+}
+
+#[inline]
+fn insert_into_impl(insert_from: &Vec<CycleHead>, insert_into: &mut Vec<CycleHead>) {
+    insert_into.reserve(insert_from.len());
+    for head in insert_from {
+        if let Some(existing) = insert_into
+            .iter()
+            .find(|candidate| candidate.database_key_index == head.database_key_index)
+        {
+            assert!(existing.iteration_count == head.iteration_count);
+        } else {
+            insert_into.push(*head);
         }
     }
 }
 
 impl IntoIterator for CycleHeads {
-    type Item = DatabaseKeyIndex;
+    type Item = CycleHead;
     type IntoIter = <Vec<Self::Item> as IntoIterator>::IntoIter;
 
     fn into_iter(self) -> Self::IntoIter {
         self.0.map(|heads| *heads).unwrap_or_default().into_iter()
     }
 }
 
-pub struct CycleHeadsIter<'a>(std::slice::Iter<'a, DatabaseKeyIndex>);
+pub struct CycleHeadsIter<'a>(std::slice::Iter<'a, CycleHead>);
 
 impl Iterator for CycleHeadsIter<'_> {
-    type Item = DatabaseKeyIndex;
+    type Item = CycleHead;
 
     fn next(&mut self) -> Option<Self::Item> {
         self.0.next().copied()
@@ -167,7 +204,7 @@ impl Iterator for CycleHeadsIter<'_> {
 impl std::iter::FusedIterator for CycleHeadsIter<'_> {}
 
 impl<'a> std::iter::IntoIterator for &'a CycleHeads {
-    type Item = DatabaseKeyIndex;
+    type Item = CycleHead;
     type IntoIter = CycleHeadsIter<'a>;
 
     fn into_iter(self) -> Self::IntoIter {
@@ -180,14 +217,14 @@ impl<'a> std::iter::IntoIterator for &'a CycleHeads {
     }
 }
 
-impl From<DatabaseKeyIndex> for CycleHeads {
-    fn from(value: DatabaseKeyIndex) -> Self {
+impl From<CycleHead> for CycleHeads {
+    fn from(value: CycleHead) -> Self {
         Self(Some(Box::new(vec![value])))
     }
 }
 
-impl From<Vec<DatabaseKeyIndex>> for CycleHeads {
-    fn from(value: Vec<DatabaseKeyIndex>) -> Self {
+impl From<Vec<CycleHead>> for CycleHeads {
+    fn from(value: Vec<CycleHead>) -> Self {
         Self(if value.is_empty() {
             None
         } else {

src/function.rs

@@ -243,7 +243,11 @@ where
     fn is_provisional_cycle_head<'db>(&'db self, db: &'db dyn Database, input: Id) -> bool {
         let zalsa = db.zalsa();
         self.get_memo_from_table_for(zalsa, input, self.memo_ingredient_index(zalsa, input))
-            .is_some_and(|memo| memo.cycle_heads().contains(&self.database_key_index(input)))
+            .is_some_and(|memo| {
+                memo.cycle_heads()
+                    .into_iter()
+                    .any(|head| head.database_key_index == self.database_key_index(input))
+            })
     }
 
     /// Attempts to claim `key_index`, returning `false` if a cycle occurs.

src/function/execute.rs

@@ -123,14 +123,19 @@ where
                     if iteration_count > MAX_ITERATIONS {
                         panic!("{database_key_index:?}: execute: too many cycle iterations");
                     }
+                    revisions
+                        .cycle_heads
+                        .update_iteration_count(database_key_index, iteration_count);
                     opt_last_provisional = Some(self.insert_memo(
                         zalsa,
                         id,
                         Memo::new(Some(new_value), revision_now, revisions),
                         memo_ingredient_index,
                     ));
 
-                    active_query = db.zalsa_local().push_query(database_key_index);
+                    active_query = db
+                        .zalsa_local()
+                        .push_query(database_key_index, iteration_count);
 
                     continue;
                 }

src/function/fetch.rs

@@ -83,7 +83,8 @@ where
         if let Some(memo) = memo_guard {
             let database_key_index = self.database_key_index(id);
             if memo.value.is_some()
-                && self.validate_may_be_provisional(db, zalsa, database_key_index, memo)
+                && (self.validate_may_be_provisional(db, zalsa, database_key_index, memo)
+                    || self.validate_same_iteration(db, database_key_index, memo))
                 && self.shallow_verify_memo(db, zalsa, database_key_index, memo)
             {
                 // SAFETY: memo is present in memo_map and we have verified that it is
@@ -158,7 +159,7 @@ where
         };
 
         // Push the query on the stack.
-        let active_query = db.zalsa_local().push_query(database_key_index);
+        let active_query = db.zalsa_local().push_query(database_key_index, 0);
 
         // 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);

src/function/maybe_changed_after.rs

@@ -112,7 +112,7 @@ where
                 CycleRecoveryStrategy::Fixpoint => {
                     return Some(VerifyResult::Unchanged(
                         InputAccumulatedValues::Empty,
-                        CycleHeads::from(database_key_index),
+                        CycleHeads::initial(database_key_index),
                     ));
                 }
             },
@@ -131,7 +131,7 @@ 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);
+        let active_query = db.zalsa_local().push_query(database_key_index, 0);
         if let VerifyResult::Unchanged(_, cycle_heads) =
             self.deep_verify_memo(db, zalsa, old_memo, &active_query)
         {
@@ -243,14 +243,17 @@ where
         database_key_index: DatabaseKeyIndex,
         memo: &Memo<C::Output<'_>>,
     ) -> bool {
-        tracing::debug!(
+        tracing::trace!(
             "{database_key_index:?}: validate_provisional(memo = {memo:#?})",
             memo = memo.tracing_debug()
         );
         if (&memo.revisions.cycle_heads).into_iter().any(|cycle_head| {
             zalsa
-                .lookup_ingredient(cycle_head.ingredient_index())
-                .is_provisional_cycle_head(db.as_dyn_database(), cycle_head.key_index())
+                .lookup_ingredient(cycle_head.database_key_index.ingredient_index())
+                .is_provisional_cycle_head(
+                    db.as_dyn_database(),
+                    cycle_head.database_key_index.key_index(),
+                )
         }) {
             return false;
         }
@@ -260,6 +263,33 @@ where
         true
     }
 
+    /// If this is a provisional memo, validate that it was cached in the same iteration of the
+    /// same cycle(s) that we are still executing. If so, it is valid for reuse. This avoids
+    /// runaway re-execution of the same queries within a fixpoint iteration.
+    pub(super) fn validate_same_iteration(
+        &self,
+        db: &C::DbView,
+        database_key_index: DatabaseKeyIndex,
+        memo: &Memo<C::Output<'_>>,
+    ) -> bool {
+        tracing::trace!(
+            "{database_key_index:?}: validate_same_iteration(memo = {memo:#?})",
+            memo = memo.tracing_debug()
+        );
+        for cycle_head in &memo.revisions.cycle_heads {
+            if !db.zalsa_local().with_query_stack(|stack| {
+                stack.iter().rev().any(|entry| {
+                    entry.database_key_index == cycle_head.database_key_index
+                        && entry.iteration_count() == cycle_head.iteration_count
+                })
+            }) {
+                return false;
+            }
+        }
+
+        true
+    }
+
     /// VerifyResult::Unchanged if the memo's value and `changed_at` time is up-to-date in the
     /// current revision. When this returns Unchanged with no cycle heads, it also updates the
     /// memo's `verified_at` field if needed to make future calls cheaper.
@@ -390,7 +420,7 @@ where
 
                     let in_heads = cycle_heads
                         .iter()
-                        .position(|&head| head == database_key_index)
+                        .position(|&head| head.database_key_index == database_key_index)
                         .inspect(|&head| _ = cycle_heads.swap_remove(head))
                         .is_some();