fix: Access to tracked-struct that was freed during fixpoint (#817)

* Add test for untracked read on tracked struct created in previous cycle

* Initial fix

* Restrict seeding to memos from the same revision

* Reduce changes

* seed_outputs

* Cleanup test

* Add assertion

* Try

* Try merging outputs after query executed

* Assert logs from first execution

* Enable trace level logging

* Use `FxIndexSet` in `diff_outputs`

* Log more events

* Cleanup

* Append outputs only once

Author: MichaReiser

src/function/diff_outputs.rs

@@ -1,6 +1,6 @@
 use crate::function::memo::Memo;
 use crate::function::{Configuration, IngredientImpl};
-use crate::hash::FxHashSet;
+use crate::hash::FxIndexSet;
 use crate::zalsa::Zalsa;
 use crate::zalsa_local::QueryRevisions;
 use crate::{AsDynDatabase as _, Database, DatabaseKeyIndex, Event, EventKind};
@@ -27,22 +27,28 @@ where
         provisional: bool,
     ) {
         // Iterate over the outputs of the `old_memo` and put them into a hashset
-        let mut old_outputs: FxHashSet<_> = old_memo.revisions.origin.outputs().collect();
+        let mut old_outputs: FxIndexSet<_> = old_memo.revisions.origin.outputs().collect();
+
+        if old_outputs.is_empty() {
+            return;
+        }
 
         // Iterate over the outputs of the current query
         // and remove elements from `old_outputs` when we find them
         for new_output in revisions.origin.outputs() {
-            old_outputs.remove(&new_output);
+            old_outputs.swap_remove(&new_output);
         }
 
-        if !old_outputs.is_empty() {
-            // Remove the outputs that are no longer present in the current revision
-            // to prevent that the next revision is seeded with a id mapping that no longer exists.
-            revisions.tracked_struct_ids.retain(|&k, &mut value| {
-                !old_outputs.contains(&DatabaseKeyIndex::new(k.ingredient_index(), value))
-            });
+        if old_outputs.is_empty() {
+            return;
         }
 
+        // Remove the outputs that are no longer present in the current revision
+        // to prevent that the next revision is seeded with an id mapping that no longer exists.
+        revisions.tracked_struct_ids.retain(|&k, &mut value| {
+            !old_outputs.contains(&DatabaseKeyIndex::new(k.ingredient_index(), value))
+        });
+
         for old_output in old_outputs {
             Self::report_stale_output(zalsa, db, key, old_output, provisional);
         }

src/function/execute.rs

@@ -59,6 +59,18 @@ where
             // Query was not previously executed, or value is potentially
             // stale, or value is absent. Let's execute!
             let mut new_value = C::execute(db, C::id_to_input(db, id));
+
+            if let Some(old_memo) = opt_old_memo {
+                // Copy over all outputs from a previous iteration.
+                // This is necessary to ensure that tracked struct created during the previous iteration
+                // (and are owned by the query) are alive even if the query in this iteration no longer creates them.
+                // The query not re-creating the tracked struct doesn't guarantee that there
+                // aren't any other queries depending on it.
+                if old_memo.may_be_provisional() && old_memo.verified_at.load() == revision_now {
+                    active_query.append_outputs(old_memo.revisions.origin.outputs());
+                }
+            }
+
             let mut revisions = active_query.pop();
 
             // Did the new result we got depend on our own provisional value, in a cycle?

src/tracked_struct.rs

@@ -406,8 +406,10 @@ where
         let current_revision = zalsa.current_revision();
         match zalsa_local.tracked_struct_id(&identity) {
             Some(id) => {
+                let index = self.database_key_index(id);
+                tracing::trace!("Reuse tracked struct {id:?}", id = index);
                 // The struct already exists in the intern map.
-                zalsa_local.add_output(self.database_key_index(id));
+                zalsa_local.add_output(index);
                 self.update(zalsa, current_revision, id, &current_deps, fields);
                 FromId::from_id(id)
             }
@@ -416,6 +418,7 @@ where
                 // This is a new tracked struct, so create an entry in the struct map.
                 let id = self.allocate(zalsa, zalsa_local, current_revision, &current_deps, fields);
                 let key = self.database_key_index(id);
+                tracing::trace!("Allocated new tracked struct {id:?}", id = key);
                 zalsa_local.add_output(key);
                 zalsa_local.store_tracked_struct_id(identity, id);
                 FromId::from_id(id)

src/zalsa_local.rs

@@ -490,6 +490,22 @@ impl ActiveQueryGuard<'_> {
         })
     }
 
+    /// Append the given `outputs` to the query's output list.
+    pub(crate) fn append_outputs<I>(&self, outputs: I)
+    where
+        I: IntoIterator<Item = DatabaseKeyIndex> + UnwindSafe,
+    {
+        self.local_state.with_query_stack(|stack| {
+            #[cfg(debug_assertions)]
+            assert_eq!(stack.len(), self.push_len);
+            let frame = stack.last_mut().unwrap();
+
+            for output in outputs {
+                frame.add_output(output);
+            }
+        })
+    }
+
     /// Invoked when the query has successfully completed execution.
     fn complete(self) -> QueryRevisions {
         let query = self.local_state.with_query_stack(|stack| {

tests/common/mod.rs

@@ -41,7 +41,7 @@ pub trait LogDatabase: HasLogger + Database {
     /// it is meant to be run from outside any tracked functions.
     fn assert_logs_len(&self, expected: usize) {
         let logs = std::mem::take(&mut *self.logger().logs.lock().unwrap());
-        assert_eq!(logs.len(), expected);
+        assert_eq!(logs.len(), expected, "Actual logs: {logs:#?}");
     }
 }
 

tests/cycle_output.rs

@@ -102,7 +102,10 @@ impl salsa::Database for Database {
             | salsa::EventKind::DidValidateMemoizedValue { .. } => {
                 self.push_log(format!("salsa_event({:?})", event.kind));
             }
-            _ => {}
+            salsa::EventKind::WillCheckCancellation => {}
+            _ => {
+                self.push_log(format!("salsa_event({:?})", event.kind));
+            }
         }
     }
 }
@@ -127,7 +130,7 @@ fn revalidate_no_changes() {
     assert_eq!(query_c(&db, c_input), 10);
     assert_eq!(query_b(&db, ab_input), 3);
 
-    db.assert_logs_len(11);
+    db.assert_logs_len(15);
 
     // trigger a new revision, but one that doesn't touch the query_a/query_b cycle
     c_input.set_value(&mut db).to(20);
@@ -136,9 +139,12 @@ fn revalidate_no_changes() {
 
     db.assert_logs(expect![[r#"
         [
-            "salsa_event(DidValidateMemoizedValue { database_key: read_value(Id(401)) })",
+            "salsa_event(DidSetCancellationFlag)",
+            "salsa_event(DidValidateMemoizedValue { database_key: read_value(Id(403)) })",
+            "salsa_event(DidReinternValue { key: Configuration(Id(800)), revision: R2 })",
             "salsa_event(DidValidateMemoizedValue { database_key: query_d(Id(800)) })",
             "salsa_event(DidValidateMemoizedValue { database_key: query_b(Id(0)) })",
+            "salsa_event(DidReinternValue { key: Configuration(Id(800)), revision: R2 })",
             "salsa_event(DidValidateMemoizedValue { database_key: query_a(Id(0)) })",
             "salsa_event(DidValidateMemoizedValue { database_key: query_b(Id(0)) })",
         ]"#]]);
@@ -154,7 +160,7 @@ fn revalidate_with_change_after_output_read() {
 
     assert_eq!(query_b(&db, ab_input), 3);
 
-    db.assert_logs_len(10);
+    db.assert_logs_len(14);
 
     // trigger a new revision that changes the output of query_d
     d_input.set_value(&mut db).to(20);
@@ -163,15 +169,29 @@ fn revalidate_with_change_after_output_read() {
 
     db.assert_logs(expect![[r#"
         [
-            "salsa_event(DidValidateMemoizedValue { database_key: read_value(Id(401)) })",
+            "salsa_event(DidSetCancellationFlag)",
+            "salsa_event(DidValidateMemoizedValue { database_key: read_value(Id(403)) })",
+            "salsa_event(DidReinternValue { key: Configuration(Id(800)), revision: R2 })",
             "salsa_event(WillExecute { database_key: query_d(Id(800)) })",
             "salsa_event(WillExecute { database_key: query_a(Id(0)) })",
-            "salsa_event(WillExecute { database_key: read_value(Id(400)) })",
+            "salsa_event(DidValidateMemoizedValue { database_key: read_value(Id(400)) })",
+            "salsa_event(WillDiscardStaleOutput { execute_key: query_a(Id(0)), output_key: Output(Id(403)) })",
+            "salsa_event(DidDiscard { key: Output(Id(403)) })",
+            "salsa_event(DidDiscard { key: read_value(Id(403)) })",
+            "salsa_event(WillDiscardStaleOutput { execute_key: query_a(Id(0)), output_key: Output(Id(402)) })",
+            "salsa_event(DidDiscard { key: Output(Id(402)) })",
+            "salsa_event(DidDiscard { key: read_value(Id(402)) })",
+            "salsa_event(WillDiscardStaleOutput { execute_key: query_a(Id(0)), output_key: Output(Id(401)) })",
+            "salsa_event(DidDiscard { key: Output(Id(401)) })",
+            "salsa_event(DidDiscard { key: read_value(Id(401)) })",
             "salsa_event(WillExecute { database_key: query_b(Id(0)) })",
+            "salsa_event(WillIterateCycle { database_key: query_b(Id(0)), iteration_count: 1, fell_back: false })",
             "salsa_event(WillExecute { database_key: query_a(Id(0)) })",
-            "salsa_event(WillExecute { database_key: read_value(Id(401)) })",
+            "salsa_event(WillExecute { database_key: read_value(Id(403)) })",
+            "salsa_event(WillIterateCycle { database_key: query_b(Id(0)), iteration_count: 2, fell_back: false })",
             "salsa_event(WillExecute { database_key: query_a(Id(0)) })",
-            "salsa_event(WillExecute { database_key: read_value(Id(400)) })",
+            "salsa_event(WillExecute { database_key: read_value(Id(402)) })",
+            "salsa_event(WillIterateCycle { database_key: query_b(Id(0)), iteration_count: 3, fell_back: false })",
             "salsa_event(WillExecute { database_key: query_a(Id(0)) })",
             "salsa_event(WillExecute { database_key: read_value(Id(401)) })",
         ]"#]]);

tests/cycle_tracked_own_input.rs

@@ -0,0 +1,130 @@
+//! Test for cycle handling where a tracked struct created in the first revision
+//! is stored in the final value of the cycle but isn't recreated in the second
+//! iteration of the creating query.
+//!
+//! It's important that the creating query in the last iteration keeps *owning* the
+//! tracked struct from the previous iteration, otherwise Salsa will discard it
+//! and dereferencing the value panics.
+mod common;
+
+use crate::common::{EventLoggerDatabase, LogDatabase};
+use expect_test::expect;
+use salsa::{CycleRecoveryAction, Database, Setter};
+
+#[salsa::input(debug)]
+struct ClassNode {
+    name: String,
+    type_params: Option<TypeParamNode>,
+}
+
+#[salsa::input(debug)]
+struct TypeParamNode {
+    name: String,
+    constraint: Option<ClassNode>,
+}
+
+#[salsa::interned(debug)]
+struct Class<'db> {
+    name: String,
+    type_params: Option<TypeParam<'db>>,
+}
+
+#[salsa::tracked(debug)]
+struct TypeParam<'db> {
+    name: String,
+    constraint: Option<Type<'db>>,
+}
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, salsa::Update)]
+enum Type<'db> {
+    Class(Class<'db>),
+    Unknown,
+}
+
+impl Type<'_> {
+    fn class(&self) -> Option<Class> {
+        match self {
+            Type::Class(class) => Some(*class),
+            Type::Unknown => None,
+        }
+    }
+}
+
+#[salsa::tracked(cycle_fn=infer_class_recover, cycle_initial=infer_class_initial)]
+fn infer_class<'db>(db: &'db dyn salsa::Database, node: ClassNode) -> Type<'db> {
+    Type::Class(Class::new(
+        db,
+        node.name(db),
+        node.type_params(db).map(|tp| infer_type_param(db, tp)),
+    ))
+}
+
+#[salsa::tracked]
+fn infer_type_param<'db>(db: &'db dyn salsa::Database, node: TypeParamNode) -> TypeParam<'db> {
+    if let Some(constraint) = node.constraint(db) {
+        // Reuse the type param from the class if any.
+        // The example is a bit silly, because it's a reduction of what we have in Astral's type checker
+        // but including all the details doesn't make sense. What's important for the test is
+        // that this query doesn't re-create the `TypeParam` tracked struct in the second iteration
+        // and instead returns the one from the first iteration which
+        // then is returned in the overall result (Class).
+        match infer_class(db, constraint) {
+            Type::Class(class) => class
+                .type_params(db)
+                .unwrap_or_else(|| TypeParam::new(db, node.name(db), Some(Type::Unknown))),
+            Type::Unknown => TypeParam::new(db, node.name(db), Some(Type::Unknown)),
+        }
+    } else {
+        TypeParam::new(db, node.name(db), None)
+    }
+}
+
+fn infer_class_initial(_db: &dyn Database, _node: ClassNode) -> Type {
+    Type::Unknown
+}
+
+fn infer_class_recover<'db>(
+    _db: &'db dyn Database,
+    _type: &Type<'db>,
+    _count: u32,
+    _inputs: ClassNode,
+) -> CycleRecoveryAction<Type<'db>> {
+    CycleRecoveryAction::Iterate
+}
+
+#[test]
+fn main() {
+    let mut db = EventLoggerDatabase::default();
+
+    // Class with a type parameter that's constrained to itself.
+    // class Test[T: Test]: ...
+    let class_node = ClassNode::new(&db, "Test".to_string(), None);
+    let type_param_node = TypeParamNode::new(&db, "T".to_string(), Some(class_node));
+    class_node
+        .set_type_params(&mut db)
+        .to(Some(type_param_node));
+
+    let ty = infer_class(&db, class_node);
+
+    db.assert_logs(expect![[r#"
+        [
+            "DidSetCancellationFlag",
+            "WillCheckCancellation",
+            "WillExecute { database_key: infer_class(Id(0)) }",
+            "WillCheckCancellation",
+            "WillExecute { database_key: infer_type_param(Id(400)) }",
+            "WillCheckCancellation",
+            "DidInternValue { key: Class(Id(c00)), revision: R2 }",
+            "WillIterateCycle { database_key: infer_class(Id(0)), iteration_count: 1, fell_back: false }",
+            "WillCheckCancellation",
+            "WillExecute { database_key: infer_type_param(Id(400)) }",
+            "WillCheckCancellation",
+        ]"#]]);
+
+    let class = ty.class().unwrap();
+    let type_param = class.type_params(&db).unwrap();
+
+    // Now read the name from the type param struct that was created in the first iteration of
+    // `infer_type_param`. This should not panic!
+    assert_eq!(type_param.name(&db), "T");
+}