[ty] Normalize divergent types with oscillating cycle heads

crates/ty_python_semantic/resources/corpus/cyclic_lambdas.py

@@ -0,0 +1,25 @@
+# This test would previous panic with: `infer_definition_types(Id(1406)): execute: too many cycle iterations`.
+
+lambda: name_4
+
+@lambda: name_5
+class name_1: ...
+
+name_2 = [lambda: name_4, name_1]
+
+if name_2:
+    @(*name_2,)
+    class name_3: ...
+    assert unique_name_19
+
+@lambda: name_3
+class name_4[*name_2](0, name_1=name_3): ...
+
+try:
+    [name_5, name_4] = *name_4, = name_4
+except* 0:
+    ...
+else:
+    async def name_4(): ...
+
+for name_3 in name_4: ...

crates/ty_python_semantic/resources/mdtest/cycle.md

@@ -138,7 +138,7 @@ class C:
         # revealed: (positional_only=..., /) -> Unknown | ((positional_only=..., /) -> Unknown) | ((positional_only=..., /) -> Divergent) | ((positional_only=..., /) -> Divergent)
         reveal_type(self.c)
 
-        # revealed: (*, kw_only=...) -> Unknown | ((*, kw_only=...) -> Unknown) | ((*, kw_only=...) -> Divergent) | ((*, kw_only=...) -> Divergent)
+        # revealed: (*, kw_only: Unknown = ...) -> Unknown | ((*, kw_only=...) -> Unknown) | ((*, kw_only=...) -> Divergent)
         reveal_type(self.d)
 ```
 

crates/ty_python_semantic/src/place.rs

@@ -842,7 +842,7 @@ impl<'db> PlaceAndQualifiers<'db> {
             // However, the handling described above may reduce the exactness of reachability analysis,
             // so it may be better to remove it. In that case, this branch is necessary.
             (Place::Undefined, Place::Defined(current)) => Place::Defined(DefinedPlace {
-                ty: current.ty.recursive_type_normalized(db, cycle),
+                ty: current.ty.recursive_type_normalized(db, cycle.head_ids()),
                 definedness: Definedness::PossiblyUndefined,
                 ..current
             }),
@@ -853,7 +853,7 @@ impl<'db> PlaceAndQualifiers<'db> {
                     Place::Undefined
                 } else {
                     Place::Defined(DefinedPlace {
-                        ty: prev.ty.recursive_type_normalized(db, cycle),
+                        ty: prev.ty.recursive_type_normalized(db, cycle.head_ids()),
                         definedness: Definedness::PossiblyUndefined,
                         ..prev
                     })

crates/ty_python_semantic/src/types.rs

@@ -2,6 +2,7 @@ use compact_str::{CompactString, ToCompactString};
 use infer::nearest_enclosing_class;
 use itertools::{Either, Itertools};
 use ruff_diagnostics::{Edit, Fix};
+use rustc_hash::FxHashSet;
 
 use std::borrow::Cow;
 use std::cell::RefCell;
@@ -75,7 +76,7 @@ use crate::types::typed_dict::TypedDictField;
 pub(crate) use crate::types::typed_dict::{TypedDictParams, TypedDictType, walk_typed_dict_type};
 pub use crate::types::variance::TypeVarVariance;
 use crate::types::variance::VarianceInferable;
-use crate::types::visitor::any_over_type;
+use crate::types::visitor::{any_over_type, any_over_type_mut};
 use crate::unpack::EvaluationMode;
 use crate::{Db, FxOrderSet, Program};
 pub use class::KnownClass;
@@ -912,8 +913,8 @@ impl<'db> Type<'db> {
         any_over_type(
             db,
             *self,
-            &|ty| matches!(ty, Type::TypeVar(tv) if tv.typevar(db).is_self(db)),
             false,
+            &|ty| matches!(ty, Type::TypeVar(tv) if tv.typevar(db).is_self(db)),
         )
     }
 
@@ -928,6 +929,8 @@ impl<'db> Type<'db> {
         previous: Self,
         cycle: &salsa::Cycle,
     ) -> Self {
+        let mut cycle_heads: FxHashSet<_> = cycle.head_ids().collect();
+
         // When we encounter a salsa cycle, we want to avoid oscillating between two or more types
         // without converging on a fixed-point result. Most of the time, we union together the
         // types from each cycle iteration to ensure that our result is monotonic, even if we
@@ -956,34 +959,42 @@ impl<'db> Type<'db> {
             }
 
             _ => {
+                let has_divergent_type_in_cycle = |ty| {
+                    any_over_type(db, ty, false, &|ty| {
+                        ty.as_divergent()
+                            .is_some_and(|DivergentType { id }| cycle.head_ids().contains(&id))
+                    })
+                };
+
                 // Also avoid unioning in a previous type which contains a Divergent from the
                 // current cycle, if the most-recent type does not. This cannot cause an
                 // oscillation, since Divergent is only introduced at the start of fixpoint
                 // iteration.
-                let has_divergent_type_in_cycle = |ty| {
-                    any_over_type(
-                        db,
-                        ty,
-                        &|nested_ty| {
-                            matches!(
-                    nested_ty,
-                    Type::Dynamic(DynamicType::Divergent(DivergentType { id }))
-                    if cycle.head_ids().contains(&id))
-                        },
-                        false,
-                    )
-                };
                 if has_divergent_type_in_cycle(previous) && !has_divergent_type_in_cycle(self) {
                     self
                 } else {
+                    // Continue to normalize any divergent types that were found in previous
+                    // cycles, as they may also be present in the most-recent type, despite
+                    // the cycle heads having changed.
+                    //
+                    // Without this, we may encounter oscillation as the cycle heads oscillate,
+                    // despite the same divergent types being inferred in each iteration.
+                    any_over_type_mut(db, self, false, &mut |ty| {
+                        if let Some(DivergentType { id }) = ty.as_divergent() {
+                            cycle_heads.insert(id);
+                        }
+
+                        false
+                    });
+
                     // The current type is unioned to the previous type. Unioning in the reverse order can cause the fixed-point iterations to converge slowly or even fail.
                     // Consider the case where the order of union types is different between the previous and current cycle.
                     // We should use the previous union type as the base and only add new element types in this cycle, if any.
                     UnionType::from_elements_cycle_recovery(db, [previous, self])
                 }
             }
         }
-        .recursive_type_normalized(db, cycle)
+        .recursive_type_normalized(db, cycle_heads.into_iter())
     }
 
     fn is_none(&self, db: &'db dyn Db) -> bool {
@@ -1224,6 +1235,13 @@ impl<'db> Type<'db> {
         )
     }
 
+    pub(crate) const fn as_divergent(self) -> Option<DivergentType> {
+        match self {
+            Type::Dynamic(DynamicType::Divergent(divergent)) => Some(divergent),
+            _ => None,
+        }
+    }
+
     pub(crate) const fn is_type_var(self) -> bool {
         matches!(self, Type::TypeVar(_))
     }
@@ -1236,7 +1254,7 @@ impl<'db> Type<'db> {
     }
 
     pub(crate) fn has_typevar(self, db: &'db dyn Db) -> bool {
-        any_over_type(db, self, &|ty| matches!(ty, Type::TypeVar(_)), false)
+        any_over_type(db, self, false, &|ty| matches!(ty, Type::TypeVar(_)))
     }
 
     pub(crate) const fn as_special_form(self) -> Option<SpecialFormType> {
@@ -1738,8 +1756,12 @@ impl<'db> Type<'db> {
     /// If this continues, the query will not converge, so this method is called in the cycle recovery function.
     /// Then `tuple[tuple[Divergent, Literal[1]], Literal[1]]` is replaced with `tuple[Divergent, Literal[1]]` and the query converges.
     #[must_use]
-    pub(crate) fn recursive_type_normalized(self, db: &'db dyn Db, cycle: &salsa::Cycle) -> Self {
-        cycle.head_ids().fold(self, |ty, id| {
+    pub(crate) fn recursive_type_normalized(
+        self,
+        db: &'db dyn Db,
+        cycle_heads: impl Iterator<Item = salsa::Id>,
+    ) -> Self {
+        cycle_heads.fold(self, |ty, id| {
             ty.recursive_type_normalized_impl(db, Type::divergent(id), false)
                 .unwrap_or(Type::divergent(id))
         })
@@ -6174,7 +6196,7 @@ impl<'db> Type<'db> {
                     }
                 });
 
-                let is_recursive = any_over_type(db, alias.raw_value_type(db).expand_eagerly(db), &|ty| ty.is_divergent(), false);
+                let is_recursive = any_over_type(db, alias.raw_value_type(db).expand_eagerly(db), false, &|ty| ty.is_divergent());
 
                 // If the type mapping does not result in any change to this (non-recursive) type alias, do not expand it.
                 //
@@ -8095,18 +8117,13 @@ impl<'db> TypeVarInstance<'db> {
 
     fn type_is_self_referential(self, db: &'db dyn Db, ty: Type<'db>) -> bool {
         let identity = self.identity(db);
-        any_over_type(
-            db,
-            ty,
-            &|ty| match ty {
-                Type::TypeVar(bound_typevar) => identity == bound_typevar.typevar(db).identity(db),
-                Type::KnownInstance(KnownInstanceType::TypeVar(typevar)) => {
-                    identity == typevar.identity(db)
-                }
-                _ => false,
-            },
-            false,
-        )
+        any_over_type(db, ty, false, &|ty| match ty {
+            Type::TypeVar(bound_typevar) => identity == bound_typevar.typevar(db).identity(db),
+            Type::KnownInstance(KnownInstanceType::TypeVar(typevar)) => {
+                identity == typevar.identity(db)
+            }
+            _ => false,
+        })
     }
 
     #[salsa::tracked(
@@ -8323,7 +8340,7 @@ fn lazy_default_cycle_recover<'db>(
     // Normalize the default to ensure cycle convergence.
     match (previous_default, default) {
         (Some(prev), Some(default)) => Some(default.cycle_normalized(db, *prev, cycle)),
-        (None, Some(default)) => Some(default.recursive_type_normalized(db, cycle)),
+        (None, Some(default)) => Some(default.recursive_type_normalized(db, cycle.head_ids())),
         (_, None) => None,
     }
 }
@@ -8965,12 +8982,12 @@ impl<'db> TypeVarBoundOrConstraints<'db> {
     /// See [`Type::recursive_type_normalized`] for more details.
     fn recursive_type_normalized(self, db: &'db dyn Db, cycle: &salsa::Cycle) -> Self {
         match self {
-            TypeVarBoundOrConstraints::UpperBound(bound) => {
-                TypeVarBoundOrConstraints::UpperBound(bound.recursive_type_normalized(db, cycle))
-            }
+            TypeVarBoundOrConstraints::UpperBound(bound) => TypeVarBoundOrConstraints::UpperBound(
+                bound.recursive_type_normalized(db, cycle.head_ids()),
+            ),
             TypeVarBoundOrConstraints::Constraints(constraints) => {
                 TypeVarBoundOrConstraints::Constraints(
-                    constraints.map(db, |ty| ty.recursive_type_normalized(db, cycle)),
+                    constraints.map(db, |ty| ty.recursive_type_normalized(db, cycle.head_ids())),
                 )
             }
         }

crates/ty_python_semantic/src/types/constraints.rs

@@ -2012,10 +2012,10 @@ impl<'db> InteriorNode<'db> {
                 if constraint.typevar(db).identity(db) == bound_typevar {
                     return true;
                 }
-                if any_over_type(db, constraint.lower(db), &mentions_typevar, false) {
+                if any_over_type(db, constraint.lower(db), false, &mentions_typevar) {
                     return true;
                 }
-                if any_over_type(db, constraint.upper(db), &mentions_typevar, false) {
+                if any_over_type(db, constraint.upper(db), false, &mentions_typevar) {
                     return true;
                 }
                 false

crates/ty_python_semantic/src/types/function.rs

@@ -1712,8 +1712,8 @@ impl KnownFunction {
                 let contains_unknown_or_todo =
                     |ty| matches!(ty, Type::Dynamic(dynamic) if dynamic != DynamicType::Any);
                 if source_type.is_equivalent_to(db, *casted_type)
-                    && !any_over_type(db, *source_type, &contains_unknown_or_todo, true)
-                    && !any_over_type(db, *casted_type, &contains_unknown_or_todo, true)
+                    && !any_over_type(db, *source_type, true, &contains_unknown_or_todo)
+                    && !any_over_type(db, *casted_type, true, &contains_unknown_or_todo)
                 {
                     if let Some(builder) = context.report_lint(&REDUNDANT_CAST, call_expression) {
                         let source_display = source_type.display(db).to_string();

crates/ty_python_semantic/src/types/infer.rs

@@ -522,6 +522,7 @@ fn unpack_cycle_recover<'db>(
     result: UnpackResult<'db>,
     _unpack: Unpack<'db>,
 ) -> UnpackResult<'db> {
+    // NOT OSCILLATING
     result.cycle_normalized(db, previous_cycle_result, cycle)
 }
 
@@ -761,7 +762,7 @@ impl<'db> DefinitionInference<'db> {
             {
                 *binding_ty = binding_ty.cycle_normalized(db, *previous_binding, cycle);
             } else {
-                *binding_ty = binding_ty.recursive_type_normalized(db, cycle);
+                *binding_ty = binding_ty.recursive_type_normalized(db, cycle.head_ids());
             }
         }
         for (declaration, declaration_ty) in &mut self.declarations {
@@ -774,8 +775,8 @@ impl<'db> DefinitionInference<'db> {
                     decl_ty.cycle_normalized(db, previous_declaration.inner_type(), cycle)
                 });
             } else {
-                *declaration_ty =
-                    declaration_ty.map_type(|decl_ty| decl_ty.recursive_type_normalized(db, cycle));
+                *declaration_ty = declaration_ty
+                    .map_type(|decl_ty| decl_ty.recursive_type_normalized(db, cycle.head_ids()));
             }
         }
 
@@ -919,7 +920,7 @@ impl<'db> ExpressionInference<'db> {
                 }) {
                     *binding_ty = binding_ty.cycle_normalized(db, *previous_binding, cycle);
                 } else {
-                    *binding_ty = binding_ty.recursive_type_normalized(db, cycle);
+                    *binding_ty = binding_ty.recursive_type_normalized(db, cycle.head_ids());
                 }
             }
         }

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -15193,20 +15193,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 } else {
                     return Err(GenericContextError::InvalidArgument);
                 }
-            } else if any_over_type(
-                db,
-                *typevar,
-                &|ty| match ty {
-                    Type::Dynamic(DynamicType::TodoUnpack | DynamicType::TodoStarredExpression) => {
-                        true
-                    }
-                    Type::NominalInstance(nominal) => {
-                        nominal.has_known_class(db, KnownClass::TypeVarTuple)
-                    }
-                    _ => false,
-                },
-                true,
-            ) {
+            } else if any_over_type(db, *typevar, true, &|ty| match ty {
+                Type::Dynamic(DynamicType::TodoUnpack | DynamicType::TodoStarredExpression) => true,
+                Type::NominalInstance(nominal) => {
+                    nominal.has_known_class(db, KnownClass::TypeVarTuple)
+                }
+                _ => false,
+            }) {
                 return Err(GenericContextError::NotYetSupported);
             } else {
                 if let Some(builder) = self.context.report_lint(&INVALID_ARGUMENT_TYPE, subscript) {

crates/ty_python_semantic/src/types/infer/builder/type_expression.rs

@@ -818,7 +818,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
         // instead of two. So until we properly support these, specialize all remaining type
         // variables with a `@Todo` type (since we don't know which of the type arguments
         // belongs to the remaining type variables).
-        if any_over_type(self.db(), value_ty, &|ty| ty.is_divergent(), true) {
+        if any_over_type(self.db(), value_ty, true, &|ty| ty.is_divergent()) {
             let value_ty = value_ty.apply_specialization(
                 db,
                 generic_context.specialize(

crates/ty_python_semantic/src/types/visitor.rs

@@ -298,11 +298,22 @@ impl<'db> TypeCollector<'db> {
 pub(super) fn any_over_type<'db>(
     db: &'db dyn Db,
     ty: Type<'db>,
+    should_visit_lazy_type_attributes: bool,
     query: &dyn Fn(Type<'db>) -> bool,
+) -> bool {
+    any_over_type_mut(db, ty, should_visit_lazy_type_attributes, &mut |ty| {
+        query(ty)
+    })
+}
+
+pub(super) fn any_over_type_mut<'db>(
+    db: &'db dyn Db,
+    ty: Type<'db>,
     should_visit_lazy_type_attributes: bool,
+    query: &mut dyn FnMut(Type<'db>) -> bool,
 ) -> bool {
     struct AnyOverTypeVisitor<'db, 'a> {
-        query: &'a dyn Fn(Type<'db>) -> bool,
+        query: RefCell<&'a mut dyn FnMut(Type<'db>) -> bool>,
         recursion_guard: TypeCollector<'db>,
         found_matching_type: Cell<bool>,
         should_visit_lazy_type_attributes: bool,
@@ -318,7 +329,7 @@ pub(super) fn any_over_type<'db>(
             if already_found {
                 return;
             }
-            let found = already_found | (self.query)(ty);
+            let found = already_found | (self.query.borrow_mut())(ty);
             self.found_matching_type.set(found);
             if found {
                 return;
@@ -328,7 +339,7 @@ pub(super) fn any_over_type<'db>(
     }
 
     let visitor = AnyOverTypeVisitor {
-        query,
+        query: RefCell::new(query),
         recursion_guard: TypeCollector::default(),
         found_matching_type: Cell::new(false),
         should_visit_lazy_type_attributes,