even lazier

crates/ty_python_semantic/src/types/constraints.rs

@@ -1004,9 +1004,8 @@ impl<'db> Node<'db> {
             Node::AlwaysTrue => {}
             Node::AlwaysFalse => {}
             Node::Interior(interior) => {
-                let map = interior.sequent_map(db);
-                let mut path = PathAssignments::default();
-                self.for_each_path_inner(db, &mut f, map, &mut path);
+                let mut path = interior.path_assignments(db);
+                self.for_each_path_inner(db, &mut f, &mut path);
             }
         }
     }
@@ -1015,7 +1014,6 @@ impl<'db> Node<'db> {
         self,
         db: &'db dyn Db,
         f: &mut dyn FnMut(&PathAssignments<'db>),
-        map: &SequentMap<'db>,
         path: &mut PathAssignments<'db>,
     ) {
         match self {
@@ -1024,11 +1022,11 @@ impl<'db> Node<'db> {
             Node::Interior(interior) => {
                 let constraint = interior.constraint(db);
                 let source_order = interior.source_order(db);
-                path.walk_edge(db, map, constraint.when_true(), source_order, |path, _| {
-                    interior.if_true(db).for_each_path_inner(db, f, map, path);
+                path.walk_edge(db, constraint.when_true(), source_order, |path, _| {
+                    interior.if_true(db).for_each_path_inner(db, f, path);
                 });
-                path.walk_edge(db, map, constraint.when_false(), source_order, |path, _| {
-                    interior.if_false(db).for_each_path_inner(db, f, map, path);
+                path.walk_edge(db, constraint.when_false(), source_order, |path, _| {
+                    interior.if_false(db).for_each_path_inner(db, f, path);
                 });
             }
         }
@@ -1040,19 +1038,13 @@ impl<'db> Node<'db> {
             Node::AlwaysTrue => true,
             Node::AlwaysFalse => false,
             Node::Interior(interior) => {
-                let map = interior.sequent_map(db);
-                let mut path = PathAssignments::default();
-                self.is_always_satisfied_inner(db, map, &mut path)
+                let mut path = interior.path_assignments(db);
+                self.is_always_satisfied_inner(db, &mut path)
             }
         }
     }
 
-    fn is_always_satisfied_inner(
-        self,
-        db: &'db dyn Db,
-        map: &SequentMap<'db>,
-        path: &mut PathAssignments<'db>,
-    ) -> bool {
+    fn is_always_satisfied_inner(self, db: &'db dyn Db, path: &mut PathAssignments<'db>) -> bool {
         match self {
             Node::AlwaysTrue => true,
             Node::AlwaysFalse => false,
@@ -1063,10 +1055,8 @@ impl<'db> Node<'db> {
                 let constraint = interior.constraint(db);
                 let source_order = interior.source_order(db);
                 let true_always_satisfied = path
-                    .walk_edge(db, map, constraint.when_true(), source_order, |path, _| {
-                        interior
-                            .if_true(db)
-                            .is_always_satisfied_inner(db, map, path)
+                    .walk_edge(db, constraint.when_true(), source_order, |path, _| {
+                        interior.if_true(db).is_always_satisfied_inner(db, path)
                     })
                     .unwrap_or(true);
                 if !true_always_satisfied {
@@ -1074,10 +1064,8 @@ impl<'db> Node<'db> {
                 }
 
                 // Ditto for the if_false branch
-                path.walk_edge(db, map, constraint.when_false(), source_order, |path, _| {
-                    interior
-                        .if_false(db)
-                        .is_always_satisfied_inner(db, map, path)
+                path.walk_edge(db, constraint.when_false(), source_order, |path, _| {
+                    interior.if_false(db).is_always_satisfied_inner(db, path)
                 })
                 .unwrap_or(true)
             }
@@ -1090,19 +1078,13 @@ impl<'db> Node<'db> {
             Node::AlwaysTrue => false,
             Node::AlwaysFalse => true,
             Node::Interior(interior) => {
-                let map = interior.sequent_map(db);
-                let mut path = PathAssignments::default();
-                self.is_never_satisfied_inner(db, map, &mut path)
+                let mut path = interior.path_assignments(db);
+                self.is_never_satisfied_inner(db, &mut path)
             }
         }
     }
 
-    fn is_never_satisfied_inner(
-        self,
-        db: &'db dyn Db,
-        map: &SequentMap<'db>,
-        path: &mut PathAssignments<'db>,
-    ) -> bool {
+    fn is_never_satisfied_inner(self, db: &'db dyn Db, path: &mut PathAssignments<'db>) -> bool {
         match self {
             Node::AlwaysTrue => false,
             Node::AlwaysFalse => true,
@@ -1113,8 +1095,8 @@ impl<'db> Node<'db> {
                 let constraint = interior.constraint(db);
                 let source_order = interior.source_order(db);
                 let true_never_satisfied = path
-                    .walk_edge(db, map, constraint.when_true(), source_order, |path, _| {
-                        interior.if_true(db).is_never_satisfied_inner(db, map, path)
+                    .walk_edge(db, constraint.when_true(), source_order, |path, _| {
+                        interior.if_true(db).is_never_satisfied_inner(db, path)
                     })
                     .unwrap_or(true);
                 if !true_never_satisfied {
@@ -1122,10 +1104,8 @@ impl<'db> Node<'db> {
                 }
 
                 // Ditto for the if_false branch
-                path.walk_edge(db, map, constraint.when_false(), source_order, |path, _| {
-                    interior
-                        .if_false(db)
-                        .is_never_satisfied_inner(db, map, path)
+                path.walk_edge(db, constraint.when_false(), source_order, |path, _| {
+                    interior.if_false(db).is_never_satisfied_inner(db, path)
                 })
                 .unwrap_or(true)
             }
@@ -1411,13 +1391,12 @@ impl<'db> Node<'db> {
         self,
         db: &'db dyn Db,
         should_remove: &mut dyn FnMut(ConstrainedTypeVar<'db>) -> bool,
-        map: &SequentMap<'db>,
         path: &mut PathAssignments<'db>,
     ) -> Self {
         match self {
             Node::AlwaysTrue => Node::AlwaysTrue,
             Node::AlwaysFalse => Node::AlwaysFalse,
-            Node::Interior(interior) => interior.abstract_one_inner(db, should_remove, map, path),
+            Node::Interior(interior) => interior.abstract_one_inner(db, should_remove, path),
         }
     }
 
@@ -2007,8 +1986,7 @@ impl<'db> InteriorNode<'db> {
 
     #[salsa::tracked(heap_size=ruff_memory_usage::heap_size)]
     fn exists_one(self, db: &'db dyn Db, bound_typevar: BoundTypeVarIdentity<'db>) -> Node<'db> {
-        let map = self.sequent_map(db);
-        let mut path = PathAssignments::default();
+        let mut path = self.path_assignments(db);
         let mentions_typevar = |ty: Type<'db>| match ty {
             Type::TypeVar(haystack) => haystack.identity(db) == bound_typevar,
             _ => false,
@@ -2034,15 +2012,13 @@ impl<'db> InteriorNode<'db> {
                 }
                 false
             },
-            map,
             &mut path,
         )
     }
 
     #[salsa::tracked(heap_size=ruff_memory_usage::heap_size)]
     fn retain_one(self, db: &'db dyn Db, bound_typevar: BoundTypeVarIdentity<'db>) -> Node<'db> {
-        let map = self.sequent_map(db);
-        let mut path = PathAssignments::default();
+        let mut path = self.path_assignments(db);
         self.abstract_one_inner(
             db,
             // Remove any node that constrains some other typevar than `bound_typevar`, and any
@@ -2058,7 +2034,6 @@ impl<'db> InteriorNode<'db> {
                 }
                 false
             },
-            map,
             &mut path,
         )
     }
@@ -2067,7 +2042,6 @@ impl<'db> InteriorNode<'db> {
         self,
         db: &'db dyn Db,
         should_remove: &mut dyn FnMut(ConstrainedTypeVar<'db>) -> bool,
-        map: &SequentMap<'db>,
         path: &mut PathAssignments<'db>,
     ) -> Node<'db> {
         let self_constraint = self.constraint(db);
@@ -2089,13 +2063,10 @@ impl<'db> InteriorNode<'db> {
             let if_true = path
                 .walk_edge(
                     db,
-                    map,
                     self_constraint.when_true(),
                     self_source_order,
                     |path, new_range| {
-                        let branch =
-                            self.if_true(db)
-                                .abstract_one_inner(db, should_remove, map, path);
+                        let branch = self.if_true(db).abstract_one_inner(db, should_remove, path);
                         path.assignments[new_range]
                             .iter()
                             .filter(|(assignment, _)| {
@@ -2115,13 +2086,12 @@ impl<'db> InteriorNode<'db> {
             let if_false = path
                 .walk_edge(
                     db,
-                    map,
                     self_constraint.when_false(),
                     self_source_order,
                     |path, new_range| {
-                        let branch =
-                            self.if_false(db)
-                                .abstract_one_inner(db, should_remove, map, path);
+                        let branch = self
+                            .if_false(db)
+                            .abstract_one_inner(db, should_remove, path);
                         path.assignments[new_range]
                             .iter()
                             .filter(|(assignment, _)| {
@@ -2144,24 +2114,19 @@ impl<'db> InteriorNode<'db> {
             let if_true = path
                 .walk_edge(
                     db,
-                    map,
                     self_constraint.when_true(),
                     self_source_order,
-                    |path, _| {
-                        self.if_true(db)
-                            .abstract_one_inner(db, should_remove, map, path)
-                    },
+                    |path, _| self.if_true(db).abstract_one_inner(db, should_remove, path),
                 )
                 .unwrap_or(Node::AlwaysFalse);
             let if_false = path
                 .walk_edge(
                     db,
-                    map,
                     self_constraint.when_false(),
                     self_source_order,
                     |path, _| {
                         self.if_false(db)
-                            .abstract_one_inner(db, should_remove, map, path)
+                            .abstract_one_inner(db, should_remove, path)
                     },
                 )
                 .unwrap_or(Node::AlwaysFalse);
@@ -2233,11 +2198,14 @@ impl<'db> InteriorNode<'db> {
         });
         constraints.sort_unstable_by_key(|(_, source_order)| *source_order);
 
-        let mut map = SequentMap::default();
-        for (constraint, _) in constraints {
-            map.add(db, constraint);
+        SequentMap::new(constraints.into_iter().map(|(constraint, _)| constraint))
+    }
+
+    fn path_assignments(self, db: &'db dyn Db) -> PathAssignments<'db> {
+        PathAssignments {
+            map: self.sequent_map(db).clone(),
+            assignments: FxOrderMap::default(),
         }
-        map
     }
 
     /// Returns a simplified version of a BDD.
@@ -2783,7 +2751,7 @@ impl<'db> ConstraintAssignment<'db> {
 ///
 /// - `C → D`: This indicates that `C` on its own is enough to imply `D`. Any path that assumes `C`
 ///   holds but `D` does _not_ is impossible and can be pruned.
-#[derive(Debug, Default, Eq, PartialEq, get_size2::GetSize, salsa::Update)]
+#[derive(Clone, Debug, Default, Eq, PartialEq, get_size2::GetSize, salsa::Update)]
 struct SequentMap<'db> {
     /// Sequents of the form `¬C₁ → false`
     single_tautologies: FxHashSet<ConstrainedTypeVar<'db>>,
@@ -2796,54 +2764,67 @@ struct SequentMap<'db> {
     >,
     /// Sequents of the form `C → D`
     single_implications: FxHashMap<ConstrainedTypeVar<'db>, FxOrderSet<ConstrainedTypeVar<'db>>>,
-    /// Constraints that we have already processed
-    processed: FxHashSet<ConstrainedTypeVar<'db>>,
-    /// Constraints that enqueued to be processed
-    enqueued: Vec<ConstrainedTypeVar<'db>>,
+    /// Constraints that we have discovered, mapped to whether we have processed them yet. (This
+    /// ensures a stable order for all of the derived constraints that we create, while still
+    /// letting us create them lazily.)
+    discovered: FxOrderMap<ConstrainedTypeVar<'db>, bool>,
 }
 
 impl<'db> SequentMap<'db> {
-    fn add(&mut self, db: &'db dyn Db, constraint: ConstrainedTypeVar<'db>) {
-        self.enqueue_constraint(constraint);
-
-        while let Some(constraint) = self.enqueued.pop() {
-            // If we've already processed this constraint, we can skip it.
-            if !self.processed.insert(constraint) {
-                continue;
-            }
-
-            // First see if we can create any sequents from the constraint on its own.
-            tracing::trace!(
-                target: "ty_python_semantic::types::constraints::SequentMap",
-                constraint = %constraint.display(db),
-                "add sequents for constraint",
-            );
-            self.add_sequents_for_single(db, constraint);
-
-            // Then check this constraint against all of the other ones we've seen so far, seeing
-            // if they're related to each other.
-            let processed = std::mem::take(&mut self.processed);
-            for other in &processed {
-                if constraint != *other {
-                    tracing::trace!(
-                        target: "ty_python_semantic::types::constraints::SequentMap",
-                        left = %constraint.display(db),
-                        right = %other.display(db),
-                        "add sequents for constraint pair",
-                    );
-                    self.add_sequents_for_pair(db, constraint, *other);
-                }
-            }
-            self.processed = processed;
+    fn new(constraints: impl IntoIterator<Item = ConstrainedTypeVar<'db>>) -> Self {
+        let discovered = constraints
+            .into_iter()
+            .map(|constraint| (constraint, false))
+            .collect();
+        Self {
+            single_tautologies: FxHashSet::default(),
+            pair_impossibilities: FxHashSet::default(),
+            pair_implications: FxHashMap::default(),
+            single_implications: FxHashMap::default(),
+            discovered,
         }
     }
 
-    fn enqueue_constraint(&mut self, constraint: ConstrainedTypeVar<'db>) {
+    fn add(&mut self, db: &'db dyn Db, constraint: ConstrainedTypeVar<'db>) {
         // If we've already processed this constraint, we can skip it.
-        if self.processed.contains(&constraint) {
+        let (new_index, existing) = self.discovered.insert_full(constraint, true);
+        let already_processed =
+            existing.expect("should not process constraint before discovering it");
+        if already_processed {
             return;
         }
-        self.enqueued.push(constraint);
+
+        // First see if we can create any sequents from the constraint on its own.
+        tracing::trace!(
+            target: "ty_python_semantic::types::constraints::SequentMap",
+            constraint = %constraint.display(db),
+            "add sequents for constraint",
+        );
+        self.add_sequents_for_single(db, constraint);
+
+        // Then check this constraint against all of the other ones we've seen so far, seeing
+        // if they're related to each other.
+        for other_index in 0..self.discovered.len() {
+            if new_index != other_index {
+                let other_constraint = self.discovered.keys()[other_index];
+                let (left, right) = if new_index < other_index {
+                    (constraint, other_constraint)
+                } else {
+                    (other_constraint, constraint)
+                };
+                tracing::trace!(
+                    target: "ty_python_semantic::types::constraints::SequentMap",
+                    left = %left.display(db),
+                    right = %right.display(db),
+                    "add sequents for constraint pair",
+                );
+                self.add_sequents_for_pair(db, left, right);
+            }
+        }
+    }
+
+    fn discover_constraint(&mut self, constraint: ConstrainedTypeVar<'db>) {
+        self.discovered.insert(constraint, false);
     }
 
     fn pair_key(
@@ -2893,6 +2874,7 @@ impl<'db> SequentMap<'db> {
         ante2: ConstrainedTypeVar<'db>,
         post: ConstrainedTypeVar<'db>,
     ) {
+        self.discover_constraint(post);
         // If either antecedent implies the consequent on its own, this new sequent is redundant.
         if ante1.implies(db, post) || ante2.implies(db, post) {
             return;
@@ -2925,6 +2907,7 @@ impl<'db> SequentMap<'db> {
         if ante == post {
             return;
         }
+        self.discover_constraint(post);
         if self
             .single_implications
             .entry(ante)
@@ -2987,7 +2970,6 @@ impl<'db> SequentMap<'db> {
         };
 
         self.add_single_implication(db, constraint, post_constraint);
-        self.enqueue_constraint(post_constraint);
     }
 
     fn add_sequents_for_pair(
@@ -3127,7 +3109,6 @@ impl<'db> SequentMap<'db> {
         let post_constraint =
             ConstrainedTypeVar::new(db, constrained_typevar, new_lower, new_upper);
         self.add_pair_implication(db, left_constraint, right_constraint, post_constraint);
-        self.enqueue_constraint(post_constraint);
     }
 
     fn add_mutual_sequents_for_same_typevars(
@@ -3177,7 +3158,6 @@ impl<'db> SequentMap<'db> {
                     _ => return,
                 };
                 self.add_pair_implication(db, left_constraint, right_constraint, post_constraint);
-                self.enqueue_constraint(post_constraint);
             };
 
         try_one_direction(left_constraint, right_constraint);
@@ -3231,7 +3211,6 @@ impl<'db> SequentMap<'db> {
                 );
                 self.add_single_implication(db, intersection_constraint, left_constraint);
                 self.add_single_implication(db, intersection_constraint, right_constraint);
-                self.enqueue_constraint(intersection_constraint);
             }
 
             // The sequent map only needs to include constraints that might appear in a BDD. If the
@@ -3318,8 +3297,9 @@ impl<'db> SequentMap<'db> {
 
 /// The collection of constraints that we know to be true or false at a certain point when
 /// traversing a BDD.
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub(crate) struct PathAssignments<'db> {
+    map: SequentMap<'db>,
     assignments: FxOrderMap<ConstraintAssignment<'db>, usize>,
 }
 
@@ -3349,7 +3329,6 @@ impl<'db> PathAssignments<'db> {
     fn walk_edge<R>(
         &mut self,
         db: &'db dyn Db,
-        map: &SequentMap<'db>,
         assignment: ConstraintAssignment<'db>,
         source_order: usize,
         f: impl FnOnce(&mut Self, Range<usize>) -> R,
@@ -3369,7 +3348,7 @@ impl<'db> PathAssignments<'db> {
             edge = %assignment.display(db),
             "walk edge",
         );
-        let found_conflict = self.add_assignment(db, map, assignment, source_order);
+        let found_conflict = self.add_assignment(db, assignment, source_order);
         let result = if found_conflict.is_err() {
             // If that results in the path now being impossible due to a contradiction, return
             // without invoking the callback.
@@ -3419,7 +3398,6 @@ impl<'db> PathAssignments<'db> {
     fn add_assignment(
         &mut self,
         db: &'db dyn Db,
-        map: &SequentMap<'db>,
         assignment: ConstraintAssignment<'db>,
         source_order: usize,
     ) -> Result<(), PathAssignmentConflict> {
@@ -3453,7 +3431,9 @@ impl<'db> PathAssignments<'db> {
         // don't anticipate the sequent maps to be very large. We might consider avoiding the
         // brute-force search.
 
-        for ante in &map.single_tautologies {
+        self.map.add(db, assignment.constraint());
+
+        for ante in &self.map.single_tautologies {
             if self.assignment_holds(ante.when_false()) {
                 // The sequent map says (ante1) is always true, and the current path asserts that
                 // it's false.
@@ -3470,7 +3450,7 @@ impl<'db> PathAssignments<'db> {
             }
         }
 
-        for (ante1, ante2) in &map.pair_impossibilities {
+        for (ante1, ante2) in &self.map.pair_impossibilities {
             if self.assignment_holds(ante1.when_true()) && self.assignment_holds(ante2.when_true())
             {
                 // The sequent map says (ante1 ∧ ante2) is an impossible combination, and the
@@ -3489,24 +3469,29 @@ impl<'db> PathAssignments<'db> {
             }
         }
 
-        for ((ante1, ante2), posts) in &map.pair_implications {
+        let mut new_constraints = Vec::new();
+        for ((ante1, ante2), posts) in &self.map.pair_implications {
             for post in posts {
                 if self.assignment_holds(ante1.when_true())
                     && self.assignment_holds(ante2.when_true())
                 {
-                    self.add_assignment(db, map, post.when_true(), source_order)?;
+                    new_constraints.push(*post);
                 }
             }
         }
 
-        for (ante, posts) in &map.single_implications {
+        for (ante, posts) in &self.map.single_implications {
             for post in posts {
                 if self.assignment_holds(ante.when_true()) {
-                    self.add_assignment(db, map, post.when_true(), source_order)?;
+                    new_constraints.push(*post);
                 }
             }
         }
 
+        for new_constraint in new_constraints {
+            self.add_assignment(db, new_constraint.when_true(), source_order)?;
+        }
+
         Ok(())
     }
 }