[ty] Add a materialization visitor

Author: charliermarsh

crates/ty_python_semantic/resources/mdtest/pep695_type_aliases.md

@@ -489,6 +489,54 @@ static_assert(is_subtype_of(Bottom[JsonDict], Bottom[JsonDict]))
 static_assert(is_subtype_of(Bottom[JsonDict], Top[JsonDict]))
 ```
 
+### Equivalence of top materializations of mutually recursive invariant aliases
+
+```py
+from typing import Callable
+from ty_extensions import static_assert, is_equivalent_to, is_subtype_of, Top
+
+class Box[T]:
+    pass
+
+type A = Callable[[B], None]
+type B = Callable[[A], None]
+
+static_assert(is_equivalent_to(Top[Box[A]], Top[Box[B]]))
+static_assert(is_subtype_of(Top[Box[A]], Top[Box[B]]))
+static_assert(is_subtype_of(Top[Box[B]], Top[Box[A]]))
+```
+
+### Assignment through recursive aliases
+
+```py
+from __future__ import annotations
+
+type JSON = str | int | float | bool | list[JSON] | list[JSON_OBJECT] | dict[str, JSON] | None
+type JSON_OBJECT = dict[str, JSON]
+
+x: JSON_OBJECT = {"hello": 23}
+
+def f() -> JSON_OBJECT:
+    return {"hello": 23}
+```
+
+### Recursive dict alias in method return
+
+```py
+from __future__ import annotations
+from dataclasses import dataclass
+
+type NodeDict = dict[str, str | list[NodeDict]]
+
+@dataclass
+class Node:
+    label: str
+    children: list[Node]
+
+    def to_dict(self) -> NodeDict:
+        return {"label": self.label, "children": [child.to_dict() for child in self.children]}
+```
+
 ### Cyclic defaults
 
 ```py

crates/ty_python_semantic/src/types.rs

@@ -4,7 +4,9 @@ use ruff_diagnostics::{Edit, Fix};
 use rustc_hash::FxHashMap;
 
 use std::borrow::Cow;
+use std::cell::RefCell;
 use std::iter;
+use std::rc::Rc;
 use std::time::Duration;
 
 use bitflags::bitflags;
@@ -235,8 +237,108 @@ fn definition_expression_type<'db>(
     }
 }
 
+struct ApplyDefaultTypeMapping;
+struct ApplyTopMaterialization;
+struct ApplyBottomMaterialization;
+struct ApplyMaterializationEquivalence;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+enum ActiveTypeMapping {
+    Default,
+    TopMaterialization,
+    BottomMaterialization,
+}
+
 /// A [`TypeTransformer`] that is used in `apply_type_mapping` methods.
-pub(crate) type ApplyTypeMappingVisitor<'db> = TypeTransformer<'db, TypeMapping<'db, 'db>>;
+///
+/// Materialization is the only mapping mode that needs to visit the same type under two different
+/// mappings within a single recursive call chain (`Top` and `Bottom`). Keep separate cycle caches
+/// for those modes so invariant checks can safely reuse one visitor.
+pub(crate) struct ApplyTypeMappingVisitor<'db> {
+    default: TypeTransformer<'db, ApplyDefaultTypeMapping>,
+    top_materialization: TypeTransformer<'db, ApplyTopMaterialization>,
+    bottom_materialization: TypeTransformer<'db, ApplyBottomMaterialization>,
+    materialization_equivalence:
+        Rc<CycleDetector<ApplyMaterializationEquivalence, (Type<'db>, Type<'db>), bool>>,
+    active_type_mappings: RefCell<Vec<ActiveTypeMapping>>,
+}
+
+impl<'db> ApplyTypeMappingVisitor<'db> {
+    pub(crate) fn visit(&self, ty: Type<'db>, func: impl FnOnce() -> Type<'db>) -> Type<'db> {
+        let active_type_mapping = self
+            .active_type_mappings
+            .borrow()
+            .last()
+            .copied()
+            .unwrap_or(ActiveTypeMapping::Default);
+
+        match active_type_mapping {
+            ActiveTypeMapping::Default => self.default.visit(ty, func),
+            ActiveTypeMapping::TopMaterialization => self.top_materialization.visit(ty, func),
+            ActiveTypeMapping::BottomMaterialization => self.bottom_materialization.visit(ty, func),
+        }
+    }
+
+    pub(crate) fn with_type_mapping<T>(
+        &self,
+        type_mapping: &TypeMapping<'_, 'db>,
+        func: impl FnOnce() -> T,
+    ) -> T {
+        let active_type_mapping = match type_mapping {
+            TypeMapping::Materialize(MaterializationKind::Top) => {
+                ActiveTypeMapping::TopMaterialization
+            }
+            TypeMapping::Materialize(MaterializationKind::Bottom) => {
+                ActiveTypeMapping::BottomMaterialization
+            }
+            _ => ActiveTypeMapping::Default,
+        };
+
+        self.active_type_mappings
+            .borrow_mut()
+            .push(active_type_mapping);
+
+        let result = func();
+
+        let previous = self.active_type_mappings.borrow_mut().pop();
+        debug_assert_eq!(previous, Some(active_type_mapping));
+
+        result
+    }
+
+    pub(crate) fn is_equivalent_to_materialization(
+        &self,
+        db: &'db dyn Db,
+        left: Type<'db>,
+        right: Type<'db>,
+    ) -> bool {
+        self.materialization_equivalence.visit((left, right), || {
+            left.is_equivalent_to_with_materialization_visitor(db, right, self)
+        })
+    }
+
+    pub(crate) fn for_new_materialization_root(&self) -> Self {
+        Self {
+            default: TypeTransformer::default(),
+            top_materialization: TypeTransformer::default(),
+            bottom_materialization: TypeTransformer::default(),
+            materialization_equivalence: Rc::clone(&self.materialization_equivalence),
+            active_type_mappings: RefCell::default(),
+        }
+    }
+}
+
+impl<'db> Default for ApplyTypeMappingVisitor<'db> {
+    fn default() -> Self {
+        Self {
+            default: TypeTransformer::default(),
+            top_materialization: TypeTransformer::default(),
+            bottom_materialization: TypeTransformer::default(),
+            materialization_equivalence: Rc::new(CycleDetector::new(true)),
+            active_type_mappings: RefCell::default(),
+        }
+    }
+}
 
 /// A [`CycleDetector`] that is used in `find_legacy_typevars` methods.
 pub(crate) type FindLegacyTypeVarsVisitor<'db> = CycleDetector<FindLegacyTypeVars, Type<'db>, ()>;
@@ -5516,7 +5618,7 @@ impl<'db> Type<'db> {
             return self;
         }
 
-        match self {
+        visitor.with_type_mapping(type_mapping, || match self {
             Type::TypeVar(bound_typevar) => bound_typevar.apply_type_mapping_impl(db, type_mapping, visitor),
             Type::KnownInstance(known_instance) => known_instance.apply_type_mapping_impl(db, type_mapping, tcx, visitor),
 
@@ -5769,7 +5871,7 @@ impl<'db> Type<'db> {
             | Type::ClassLiteral(_)
             | Type::BoundSuper(_)
             | Type::SpecialForm(_) => self,
-        }
+        })
     }
 
     /// Locates any legacy `TypeVar`s in this type, and adds them to a set. This is used to build

crates/ty_python_semantic/src/types/class.rs

@@ -1124,11 +1124,13 @@ impl<'db> ClassType<'db> {
         let constraints = ConstraintSetBuilder::new();
         let relation_visitor = HasRelationToVisitor::default(&constraints);
         let disjointness_visitor = IsDisjointVisitor::default(&constraints);
+        let materialization_visitor = ApplyTypeMappingVisitor::default();
         let checker = TypeRelationChecker::subtyping(
             &constraints,
             InferableTypeVars::None,
             &relation_visitor,
             &disjointness_visitor,
+            &materialization_visitor,
         );
         checker
             .check_class_pair(db, self, target)

crates/ty_python_semantic/src/types/generics.rs

@@ -1230,7 +1230,11 @@ impl<'db> Specialization<'db> {
                     TypeVarVariance::Invariant => {
                         let top_materialization =
                             vartype.materialize(db, MaterializationKind::Top, visitor);
-                        if !vartype.is_equivalent_to(db, top_materialization) {
+                        if !visitor.is_equivalent_to_materialization(
+                            db,
+                            *vartype,
+                            top_materialization,
+                        ) {
                             has_dynamic_invariant_typevar = true;
                         }
                         *vartype
@@ -1270,11 +1274,13 @@ impl<'db> Specialization<'db> {
     ) -> ConstraintSet<'db, 'c> {
         let relation_visitor = HasRelationToVisitor::default(constraints);
         let disjointness_visitor = IsDisjointVisitor::default(constraints);
+        let materialization_visitor = ApplyTypeMappingVisitor::default();
         let checker = DisjointnessChecker::new(
             constraints,
             inferable,
             &relation_visitor,
             &disjointness_visitor,
+            &materialization_visitor,
         );
         checker.check_specialization_pair(db, self, other)
     }
@@ -1455,10 +1461,20 @@ impl<'c, 'db> TypeRelationChecker<'_, 'c, 'db> {
         target_type: Type<'db>,
         target_materialization: MaterializationKind,
     ) -> ConstraintSet<'db, 'c> {
-        let source_top = source_type.top_materialization(db);
-        let source_bottom = source_type.bottom_materialization(db);
-        let target_top = target_type.top_materialization(db);
-        let target_bottom = target_type.bottom_materialization(db);
+        let source_top =
+            source_type.materialize(db, MaterializationKind::Top, self.materialization_visitor);
+        let source_bottom = source_type.materialize(
+            db,
+            MaterializationKind::Bottom,
+            self.materialization_visitor,
+        );
+        let target_top =
+            target_type.materialize(db, MaterializationKind::Top, self.materialization_visitor);
+        let target_bottom = target_type.materialize(
+            db,
+            MaterializationKind::Bottom,
+            self.materialization_visitor,
+        );
 
         let is_subtype_of = |source: Type<'db>, target: Type<'db>| {
             // TODO:

crates/ty_python_semantic/src/types/instance.rs

@@ -719,11 +719,13 @@ impl<'db> ProtocolInstanceType<'db> {
             let constraints = ConstraintSetBuilder::new();
             let relation_visitor = HasRelationToVisitor::default(&constraints);
             let disjointness_visitor = IsDisjointVisitor::default(&constraints);
+            let materialization_visitor = ApplyTypeMappingVisitor::default();
             let checker = TypeRelationChecker::subtyping(
                 &constraints,
                 InferableTypeVars::None,
                 &relation_visitor,
                 &disjointness_visitor,
+                &materialization_visitor,
             );
             checker
                 .check_type_satisfies_protocol(db, Type::object(), protocol)