@@ -1351,35 +1351,55 @@ impl Type {
13511351 )
13521352 }
13531353
1354+ fn recurse_type_variable_positions < ' a > ( & ' a self , f : & mut dyn FnMut ( & ' a Type ) ) {
1355+ let mut recurse_targs = |targs : & ' a TArgs | {
1356+ for targ in targs. as_slice ( ) . iter ( ) {
1357+ f ( targ) ;
1358+ }
1359+ } ;
1360+ // IMPORTANT: keep this match in sync with `recurse_type_variable_positions_mut`
1361+ match self {
1362+ // In `A[X]`, we only check `X` for a couple reasons:
1363+ // * If we were to blindly visit the entire ClassType, we would find Quantifieds in
1364+ // the definition of the class, which is almost never what we want: we want to
1365+ // know if `X` contains any references to Quantifieds, not whether `A` is generic.
1366+ // See https://github.com/facebook/pyrefly/issues/1962.
1367+ // * Not checking the rest of the ClassType is a critical performance optimization
1368+ // when visiting Vars. See https://github.com/facebook/pyrefly/issues/2016.
1369+ Type :: ClassType ( cls) => recurse_targs ( cls. targs ( ) ) ,
1370+ Type :: TypedDict ( TypedDict :: TypedDict ( td) ) => recurse_targs ( td. targs ( ) ) ,
1371+ // `Self` is a keyword, not a user-written type variable reference, so we don't
1372+ // recurse into it when looking for type variable references.
1373+ Type :: SelfType ( _) => { }
1374+ // Enum literals contain `ClassType`s that we shouldn't visit.
1375+ Type :: Literal ( _) => { }
1376+ _ => self . recurse ( f) ,
1377+ }
1378+ }
1379+
1380+ /// The mutable form of [`Type::recurse_type_variable_positions`], visiting the same positions.
1381+ fn recurse_type_variable_positions_mut ( & mut self , f : & mut dyn FnMut ( & mut Type ) ) {
1382+ let mut recurse_targs = |targs : & mut TArgs | {
1383+ for targ in targs. as_mut ( ) . iter_mut ( ) {
1384+ f ( targ) ;
1385+ }
1386+ } ;
1387+ // IMPORTANT: keep this match in sync with `recurse_type_variable_positions`
1388+ match self {
1389+ Type :: ClassType ( cls) => recurse_targs ( cls. targs_mut ( ) ) ,
1390+ Type :: TypedDict ( TypedDict :: TypedDict ( td) ) => recurse_targs ( td. targs_mut ( ) ) ,
1391+ Type :: SelfType ( _) | Type :: Literal ( _) => { }
1392+ _ => self . recurse_mut ( f) ,
1393+ }
1394+ }
1395+
13541396 fn visit_type_variables < ' a > ( & ' a self , f : & mut dyn FnMut ( TypeVariable < ' a > ) ) {
13551397 fn visit < ' a > ( ty : & ' a Type , f : & mut dyn FnMut ( TypeVariable < ' a > ) ) {
13561398 if let Some ( tv) = TypeVariable :: new ( ty) {
13571399 f ( tv) ;
13581400 return ;
13591401 }
1360- let mut recurse_targs = |targs : & ' a TArgs | {
1361- for targ in targs. as_slice ( ) . iter ( ) {
1362- visit ( targ, f) ;
1363- }
1364- } ;
1365- // IMPORTANT: keep this match in sync with `transform_types_in_type_variable_positions`
1366- match ty {
1367- // In `A[X]`, we only check `X` for a couple reasons:
1368- // * If we were to blindly visit the entire ClassType, we would find Quantifieds in
1369- // the definition of the class, which is almost never what we want: we want to
1370- // know if `X` contains any references to Quantifieds, not whether `A` is generic.
1371- // See https://github.com/facebook/pyrefly/issues/1962.
1372- // * Not checking the rest of the ClassType is a critical performance optimization
1373- // when visiting Vars. See https://github.com/facebook/pyrefly/issues/2016.
1374- Type :: ClassType ( cls) => recurse_targs ( cls. targs ( ) ) ,
1375- Type :: TypedDict ( TypedDict :: TypedDict ( td) ) => recurse_targs ( td. targs ( ) ) ,
1376- // `Self` is a keyword, not a user-written type variable reference, so we don't
1377- // recurse into it when looking for type variable references.
1378- Type :: SelfType ( _) => { }
1379- // Enum literals contain `ClassType`s that we shouldn't visit.
1380- Type :: Literal ( _) => { }
1381- _ => ty. recurse ( & mut |ty| visit ( ty, f) ) ,
1382- }
1402+ ty. recurse_type_variable_positions ( & mut |inside| visit ( inside, f) ) ;
13831403 }
13841404 visit ( self , f)
13851405 }
@@ -1432,39 +1452,14 @@ impl Type {
14321452 self . visit_type_variables ( & mut f)
14331453 }
14341454
1435- fn transform_types_in_type_variable_positions ( & mut self , f : & mut dyn FnMut ( & mut Type ) ) {
1455+ pub fn transform_types_in_type_variable_positions ( & mut self , f : & mut dyn FnMut ( & mut Type ) ) {
14361456 fn visit ( ty : & mut Type , f : & mut dyn FnMut ( & mut Type ) ) {
14371457 f ( ty) ;
1438- let mut recurse_targs = |targs : & mut TArgs | {
1439- for targ in targs. as_mut ( ) . iter_mut ( ) {
1440- visit ( targ, f) ;
1441- }
1442- } ;
1443- // IMPORTANT: keep this match in sync with `visit_type_variables`
1444- match ty {
1445- Type :: ClassType ( cls) => recurse_targs ( cls. targs_mut ( ) ) ,
1446- Type :: TypedDict ( TypedDict :: TypedDict ( td) ) => recurse_targs ( td. targs_mut ( ) ) ,
1447- // `Self` is a keyword, not a user-written type variable reference.
1448- Type :: SelfType ( _) => { }
1449- // Enum literals contain `ClassType`s that we shouldn't visit.
1450- Type :: Literal ( _) => { }
1451- _ => ty. recurse_mut ( & mut |ty| visit ( ty, f) ) ,
1452- }
1458+ ty. recurse_type_variable_positions_mut ( & mut |inside| visit ( inside, f) ) ;
14531459 }
14541460 visit ( self , f)
14551461 }
14561462
1457- /// Transform unreplaced references to legacy type variables. Note that references to in-scope
1458- /// legacy type variables in functions and classes are replaced with Quantified, so unreplaced
1459- /// references only appear in cases like a TypeVar definition or an out-of-scope type variable.
1460- pub fn transform_raw_legacy_type_variables ( & mut self , f : & mut dyn FnMut ( & mut Type ) ) {
1461- self . transform_types_in_type_variable_positions ( & mut |ty| {
1462- if ty. is_raw_legacy_type_variable ( ) {
1463- f ( ty) ;
1464- }
1465- } )
1466- }
1467-
14681463 /// Check if the type contains a placeholder var. See `collect_maybe_placeholder_vars`.
14691464 pub fn may_contain_placeholder_var ( & self ) -> bool {
14701465 let mut seen = false ;
0 commit comments