Expose the query ID and the last provisional value to the cycle recovery function (#1012)

* Expose the query ID and the last provisional value to the cycle recovery function

* Mark cycle as converged if fallback value is the same as the last provisional

* Make `cycle_fn` optional

Author: MichaReiser

benches/dataflow.rs

@@ -76,6 +76,8 @@ fn def_cycle_initial(_db: &dyn Db, _def: Definition) -> Type {
 
 fn def_cycle_recover(
     _db: &dyn Db,
+    _id: salsa::Id,
+    _last_provisional_value: &Type,
     value: &Type,
     count: u32,
     _def: Definition,
@@ -89,6 +91,8 @@ fn use_cycle_initial(_db: &dyn Db, _use: Use) -> Type {
 
 fn use_cycle_recover(
     _db: &dyn Db,
+    _id: salsa::Id,
+    _last_provisional_value: &Type,
     value: &Type,
     count: u32,
     _use: Use,

book/src/cycles.md

@@ -21,6 +21,8 @@ fn initial(_db: &dyn KnobsDatabase) -> u32 {
 }
 ```
 
+The `cycle_fn` is optional. The default implementation always returns `Iterate`.
+
 If `query` becomes the head of a cycle (that is, `query` is executing and on the active query stack, it calls `query2`, `query2` calls `query3`, and `query3` calls `query` again -- there could be any number of queries involved in the cycle), the `initial_fn` will be called to generate an "initial" value for `query` in the fixed-point computation. (The initial value should usually be the "bottom" value in the partial order.) All queries in the cycle will compute a provisional result based on this initial value for the cycle head. That is, `query3` will compute a provisional result using the initial value for `query`, `query2` will compute a provisional result using this provisional value for `query3`. When `cycle2` returns its provisional result back to `cycle`, `cycle` will observe that it has received a provisional result from its own cycle, and will call the `cycle_fn` (with the current value and the number of iterations that have occurred so far). The `cycle_fn` can return `salsa::CycleRecoveryAction::Iterate` to indicate that the cycle should iterate again, or `salsa::CycleRecoveryAction::Fallback(value)` to indicate that fixpoint iteration should resume starting with the given value (which should be a value that will converge quickly).
 
 The cycle will iterate until it converges: that is, until two successive iterations produce the same result.

components/salsa-macro-rules/src/setup_tracked_fn.rs

@@ -308,11 +308,13 @@ macro_rules! setup_tracked_fn {
 
                 fn recover_from_cycle<$db_lt>(
                     db: &$db_lt dyn $Db,
+                    id: salsa::Id,
+                    last_provisional_value: &Self::Output<$db_lt>,
                     value: &Self::Output<$db_lt>,
-                    count: u32,
+                    iteration_count: u32,
                     ($($input_id),*): ($($interned_input_ty),*)
                 ) -> $zalsa::CycleRecoveryAction<Self::Output<$db_lt>> {
-                    $($cycle_recovery_fn)*(db, value, count, $($input_id),*)
+                    $($cycle_recovery_fn)*(db, id, last_provisional_value, value, iteration_count, $($input_id),*)
                 }
 
                 fn id_to_input<$db_lt>(zalsa: &$db_lt $zalsa::Zalsa, key: salsa::Id) -> Self::Input<$db_lt> {

components/salsa-macro-rules/src/unexpected_cycle_recovery.rs

@@ -3,10 +3,10 @@
 // a macro because it can take a variadic number of arguments.
 #[macro_export]
 macro_rules! unexpected_cycle_recovery {
-    ($db:ident, $value:ident, $count:ident, $($other_inputs:ident),*) => {{
-        std::mem::drop($db);
+    ($db:ident, $id:ident, $last_provisional_value:ident, $new_value:ident, $count:ident, $($other_inputs:ident),*) => {{
+        let (_db, _id, _last_provisional_value, _new_value, _count) = ($db, $id, $last_provisional_value, $new_value, $count);
         std::mem::drop(($($other_inputs,)*));
-        panic!("cannot recover from cycle")
+        salsa::CycleRecoveryAction::Iterate
     }};
 }
 

components/salsa-macros/src/tracked_fn.rs

@@ -286,9 +286,10 @@ impl Macro {
                 self.args.cycle_fn.as_ref().unwrap(),
                 "must provide `cycle_initial` along with `cycle_fn`",
             )),
-            (None, Some(_), None) => Err(syn::Error::new_spanned(
-                self.args.cycle_initial.as_ref().unwrap(),
-                "must provide `cycle_fn` along with `cycle_initial`",
+            (None, Some(cycle_initial), None) => Ok((
+                quote!((salsa::plumbing::unexpected_cycle_recovery!)),
+                quote!((#cycle_initial)),
+                quote!(Fixpoint),
             )),
             (None, None, Some(cycle_result)) => Ok((
                 quote!((salsa::plumbing::unexpected_cycle_recovery!)),

src/cycle.rs

@@ -70,7 +70,11 @@ pub enum CycleRecoveryAction<T> {
     /// Iterate the cycle again to look for a fixpoint.
     Iterate,
 
-    /// Cut off iteration and use the given result value for this query.
+    /// Use the given value as the result for the current iteration instead
+    /// of the value computed by the query function.
+    ///
+    /// Returning `Fallback` doesn't stop the fixpoint iteration. It only
+    /// allows the iterate function to return a different value.
     Fallback(T),
 }
 

src/function.rs

@@ -94,20 +94,32 @@ pub trait Configuration: Any {
     /// value from the latest iteration of this cycle. `count` is the number of cycle iterations
     /// completed so far.
     ///
-    /// # Iteration count semantics
+    /// # Id
     ///
-    /// The `count` parameter isn't guaranteed to start from zero or to be contiguous:
+    /// The id can be used to uniquely identify the query instance. This can be helpful
+    /// if the cycle function has to re-identify a value it returned previously.
     ///
-    /// * **Initial value**: `count` may be non-zero on the first call for a given query if that
+    /// # Values
+    ///
+    /// The `last_provisional_value` is the value from the previous iteration of this cycle
+    /// and `value` is the new value that was computed in the current iteration.
+    ///
+    /// # Iteration count
+    ///
+    /// The `iteration` parameter isn't guaranteed to start from zero or to be contiguous:
+    ///
+    /// * **Initial value**: `iteration` may be non-zero on the first call for a given query if that
     ///   query becomes the outermost cycle head after a nested cycle complete a few iterations. In this case,
-    ///   `count` continues from the nested cycle's iteration count rather than resetting to zero.
+    ///   `iteration` continues from the nested cycle's iteration count rather than resetting to zero.
     /// * **Non-contiguous values**: This function isn't called if this cycle is part of an outer cycle
     ///   and the value for this query remains unchanged for one iteration. But the outer cycle might
     ///   keep iterating because other heads keep changing.
     fn recover_from_cycle<'db>(
         db: &'db Self::DbView,
-        value: &Self::Output<'db>,
-        count: u32,
+        id: Id,
+        last_provisional_value: &Self::Output<'db>,
+        new_value: &Self::Output<'db>,
+        iteration: u32,
         input: Self::Input<'db>,
     ) -> CycleRecoveryAction<Self::Output<'db>>;
 

src/function/execute.rs

@@ -320,7 +320,7 @@ where
                 I am a cycle head, comparing last provisional value with new value"
             );
 
-            let this_converged = C::values_equal(&new_value, last_provisional_value);
+            let mut this_converged = C::values_equal(&new_value, last_provisional_value);
 
             // If this is the outermost cycle, use the maximum iteration count of all cycles.
             // This is important for when later iterations introduce new cycle heads (that then
@@ -341,6 +341,8 @@ where
                 // cycle-recovery function what to do:
                 match C::recover_from_cycle(
                     db,
+                    id,
+                    last_provisional_value,
                     &new_value,
                     iteration_count.as_u32(),
                     C::id_to_input(zalsa, id),
@@ -351,6 +353,8 @@ where
                             "{database_key_index:?}: execute: user cycle_fn says to fall back"
                         );
                         new_value = fallback_value;
+
+                        this_converged = C::values_equal(&new_value, last_provisional_value);
                     }
                 }
             }

src/function/memo.rs

@@ -557,6 +557,8 @@ mod _memory_usage {
 
         fn recover_from_cycle<'db>(
             _: &'db Self::DbView,
+            _: Id,
+            _: &Self::Output<'db>,
             _: &Self::Output<'db>,
             _: u32,
             _: Self::Input<'db>,

tests/backtrace.rs

@@ -42,7 +42,7 @@ fn query_f(db: &dyn Database, thing: Thing) -> String {
     query_cycle(db, thing)
 }
 
-#[salsa::tracked(cycle_fn=cycle_fn, cycle_initial=cycle_initial)]
+#[salsa::tracked(cycle_initial=cycle_initial)]
 fn query_cycle(db: &dyn Database, thing: Thing) -> String {
     let backtrace = query_cycle(db, thing);
     if backtrace.is_empty() {
@@ -56,15 +56,6 @@ fn cycle_initial(_db: &dyn salsa::Database, _thing: Thing) -> String {
     String::new()
 }
 
-fn cycle_fn(
-    _db: &dyn salsa::Database,
-    _value: &str,
-    _count: u32,
-    _thing: Thing,
-) -> salsa::CycleRecoveryAction<String> {
-    salsa::CycleRecoveryAction::Iterate
-}
-
 #[test]
 fn backtrace_works() {
     let db = DatabaseImpl::default();