Port over changes from other branch

Author: MichaReiser

Cargo.toml

@@ -22,7 +22,7 @@ intrusive-collections = "0.9.7"
 parking_lot = "0.12"
 portable-atomic = "1"
 rustc-hash = "2"
-smallvec = "1"
+smallvec = { version = "1", features = ["const_new"] }
 thin-vec = { version = "0.2.14" }
 tracing = { version = "0.1", default-features = false, features = ["std"] }
 

src/cycle.rs

@@ -102,10 +102,31 @@ pub enum CycleRecoveryStrategy {
 pub struct CycleHead {
     pub(crate) database_key_index: DatabaseKeyIndex,
     pub(crate) iteration_count: AtomicIterationCount,
+
+    /// Marks a cycle head as removed within its `CycleHeads` container.
+    ///
+    /// Cycle heads are marked as removed when the memo from the last iteration (a provisional memo)
+    /// is used as the initial value for the next iteration. It's necessary to remove all but its own
+    /// head from the `CycleHeads` container, because the query might now depend on fewer cycles
+    /// (in case of conditional dependencies). However, we can't actually remove the cycle head
+    /// within `fetch_cold_cycle` because we only have a readonly memo. That's what `removed` is used for.
     #[cfg_attr(feature = "persistence", serde(skip))]
     removed: AtomicBool,
 }
 
+impl CycleHead {
+    pub const fn new(
+        database_key_index: DatabaseKeyIndex,
+        iteration_count: IterationCount,
+    ) -> Self {
+        Self {
+            database_key_index,
+            iteration_count: AtomicIterationCount(AtomicU8::new(iteration_count.0)),
+            removed: AtomicBool::new(false),
+        }
+    }
+}
+
 impl Clone for CycleHead {
     fn clone(&self) -> Self {
         Self {
@@ -130,10 +151,19 @@ impl IterationCount {
         self.0 == 0
     }
 
+    /// Iteration count reserved for panicked cycles.
+    ///
+    /// Using a special iteration count ensures that `validate_same_iteration` and `validate_provisional`
+    /// return `false` for queries depending on this panicked cycle, because the iteration count is guaranteed
+    /// to be different (which isn't guaranteed if the panicked memo uses [`Self::initial`]).
     pub(crate) const fn panicked() -> Self {
         Self(u8::MAX)
     }
 
+    pub(crate) const fn is_panicked(self) -> bool {
+        self.0 == u8::MAX
+    }
+
     pub(crate) const fn increment(self) -> Option<Self> {
         let next = Self(self.0 + 1);
         if next.0 <= MAX_ITERATIONS.0 {
@@ -150,7 +180,7 @@ impl IterationCount {
 
 impl std::fmt::Display for IterationCount {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "iteration={}", self.0)
+        self.0.fmt(f)
     }
 }
 
@@ -162,6 +192,10 @@ impl AtomicIterationCount {
         IterationCount(self.0.load(Ordering::Relaxed))
     }
 
+    pub(crate) fn load_mut(&mut self) -> IterationCount {
+        IterationCount(*self.0.get_mut())
+    }
+
     pub(crate) fn store(&self, value: IterationCount) {
         self.0.store(value.0, Ordering::Release);
     }
@@ -214,10 +248,13 @@ impl CycleHeads {
         self.0.is_empty()
     }
 
-    pub(crate) fn initial(database_key_index: DatabaseKeyIndex) -> Self {
+    pub(crate) fn initial(
+        database_key_index: DatabaseKeyIndex,
+        iteration_count: IterationCount,
+    ) -> Self {
         Self(thin_vec![CycleHead {
             database_key_index,
-            iteration_count: IterationCount::initial().into(),
+            iteration_count: iteration_count.into(),
             removed: false.into()
         }])
     }
@@ -228,22 +265,35 @@ impl CycleHeads {
         }
     }
 
+    /// Iterates over all cycle heads that aren't equal to `own`.
+    pub(crate) fn iter_not_eq(&self, own: DatabaseKeyIndex) -> impl Iterator<Item = &CycleHead> {
+        self.iter()
+            .filter(move |head| head.database_key_index != own)
+    }
+
     pub(crate) fn contains(&self, value: &DatabaseKeyIndex) -> bool {
         self.into_iter()
             .any(|head| head.database_key_index == *value && !head.removed.load(Ordering::Relaxed))
     }
 
-    pub(crate) fn clear_except(&self, except: DatabaseKeyIndex) {
+    /// Removes all cycle heads except `except` by marking them as removed.
+    ///
+    /// Note that the heads aren't actually removed. They're only marked as removed and will be
+    /// skipped when iterating. This is because we might not have a mutable reference.
+    pub(crate) fn remove_all_except(&self, except: DatabaseKeyIndex) {
         for head in self.0.iter() {
             if head.database_key_index == except {
                 continue;
             }
 
-            // TODO: verify ordering
             head.removed.store(true, Ordering::Release);
         }
     }
 
+    /// Updates the iteration count for the head `cycle_head_index` to `new_iteration_count`.
+    ///
+    /// Unlike [`update_iteration_count`], this method takes a `&mut self` reference. It should
+    /// be preferred if possible, as it avoids atomic operations.
     pub(crate) fn update_iteration_count_mut(
         &mut self,
         cycle_head_index: DatabaseKeyIndex,
@@ -258,6 +308,9 @@ impl CycleHeads {
         }
     }
 
+    /// Updates the iteration count for the head `cycle_head_index` to `new_iteration_count`.
+    ///
+    /// Unlike [`update_iteration_count_mut`], this method takes a `&self` reference.
     pub(crate) fn update_iteration_count(
         &self,
         cycle_head_index: DatabaseKeyIndex,
@@ -277,15 +330,20 @@ impl CycleHeads {
         self.0.reserve(other.0.len());
 
         for head in other {
-            self.insert(head);
+            debug_assert!(!head.removed.load(Ordering::Relaxed));
+            self.insert(head.database_key_index, head.iteration_count.load());
         }
     }
 
-    pub(crate) fn insert(&mut self, head: &CycleHead) -> bool {
+    pub(crate) fn insert(
+        &mut self,
+        database_key_index: DatabaseKeyIndex,
+        iteration_count: IterationCount,
+    ) -> bool {
         if let Some(existing) = self
             .0
             .iter_mut()
-            .find(|candidate| candidate.database_key_index == head.database_key_index)
+            .find(|candidate| candidate.database_key_index == database_key_index)
         {
             let removed = existing.removed.get_mut();
 
@@ -294,23 +352,19 @@ impl CycleHeads {
 
                 true
             } else {
-                let existing_count = existing.iteration_count.load();
-                let head_count = head.iteration_count.load();
+                let existing_count = existing.iteration_count.load_mut();
 
-                // It's now possible that a query can depend on different iteration counts of the same query
-                // This because some queries (inner) read the provisional value of the last iteration
-                // while outer queries read the value from the last iteration (which is i+1 if the head didn't converge).
                 assert_eq!(
-                    existing_count, head_count,
-                    "Can't merge cycle heads {:?} with different iteration counts ({existing_count:?}, {head_count:?})",
+                    existing_count, iteration_count,
+                    "Can't merge cycle heads {:?} with different iteration counts ({existing_count:?}, {iteration_count:?})",
                     existing.database_key_index
                 );
 
                 false
             }
         } else {
-            debug_assert!(!head.removed.load(Ordering::Relaxed));
-            self.0.push(head.clone());
+            self.0
+                .push(CycleHead::new(database_key_index, iteration_count));
             true
         }
     }

src/function.rs

@@ -94,6 +94,15 @@ pub trait Configuration: Any {
     /// Decide whether to iterate a cycle again or fallback. `value` is the provisional return
     /// value from the latest iteration of this cycle. `count` is the number of cycle iterations
     /// we've already completed.
+    ///
+    /// Note: There is no guarantee that `count` always starts at 0. It's possible that
+    /// the function is called with a non-zero value even if it is the first time around for
+    /// this specific query if the query has become the outermost cycle of a larger cycle.
+    /// In this case, Salsa uses the `count` value of the already iterating cycle as the start.
+    ///
+    /// It's also not guaranteed that `count` values are contiguous. The function might not be called
+    /// if this query converged in this specific iteration OR if the query only participates conditionally
+    /// in the cycle (e.g. every other iteration).
     fn recover_from_cycle<'db>(
         db: &'db Self::DbView,
         value: &Self::Output<'db>,
@@ -372,7 +381,7 @@ where
             .set_iteration_count(Self::database_key_index(self, input), iteration_count);
     }
 
-    fn set_cycle_finalized(&self, zalsa: &Zalsa, input: Id) {
+    fn finalize_cycle_head(&self, zalsa: &Zalsa, input: Id) {
         let Some(memo) =
             self.get_memo_from_table_for(zalsa, input, self.memo_ingredient_index(zalsa, input))
         else {