Skip to content

Commit 4e4f1c1

Browse files
rchen152meta-codesync[bot]
authored andcommitted
Make visiting type variable positions more ergonomic
Summary: No functional change. Extracts `recurse_type_variable_positions` and `recurse_type_variable_positions_mut` helpers out of `visit_type_variables` and `transform_types_in_type_variable_positions`, respectively. This makes the parallelism between the two more obvious and makes it easier to spot the `match`s that need to be kept in sync with each other. I also deleted `transform_raw_legacy_type_variables`, which was a thin wrapper around `transform_types_in_type_variable_positions`. With the addition of the new `recurse_` helpers, we're accumulating quite a few levels of indirection, so let's cut out this one whose usefulness was borderline. Reviewed By: grievejia Differential Revision: D117438505 fbshipit-source-id: 9fed132a3a22b61e8662108bbb17ae03732a9acf
1 parent b9f38ab commit 4e4f1c1

2 files changed

Lines changed: 53 additions & 56 deletions

File tree

‎crates/pyrefly_types/src/types.rs‎

Lines changed: 45 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -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;

‎pyrefly/lib/alt/solve.rs‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2235,14 +2235,12 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
22352235
}
22362236
if let Some(default) = &mut tparam.default {
22372237
let mut out_of_scope_names = Vec::new();
2238-
default.transform_raw_legacy_type_variables(&mut |ty| {
2238+
default.transform_types_in_type_variable_positions(&mut |ty| {
22392239
let (name, kind) = match &*ty {
22402240
Type::TypeVar(t) => (t.qname().id(), QuantifiedKind::TypeVar),
22412241
Type::TypeVarTuple(t) => (t.qname().id(), QuantifiedKind::TypeVarTuple),
22422242
Type::ParamSpec(p) => (p.qname().id(), QuantifiedKind::ParamSpec),
2243-
_ => unreachable!(
2244-
"transform_raw_legacy_type_variables only visits legacy type variables"
2245-
),
2243+
_ => return,
22462244
};
22472245
*ty = match seen.get(name) {
22482246
Some(q) => (**q).clone().to_type(self.heap),
@@ -5703,7 +5701,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
57035701
let heap = self.heap;
57045702
let module = self.module().name();
57055703
callable.visit_mut(&mut |ty| {
5706-
ty.transform_raw_legacy_type_variables(&mut |ty| {
5704+
ty.transform_types_in_type_variable_positions(&mut |ty| {
57075705
if let Type::TypeVar(tv) = ty {
57085706
let q = seen_type_vars
57095707
.entry(tv.dupe())
@@ -5753,7 +5751,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
57535751
);
57545752
}
57555753
*ty = wrapped;
5756-
ty.transform_raw_legacy_type_variables(&mut |t| *t = self.heap.mk_any_error());
5754+
ty.transform_types_in_type_variable_positions(&mut |t| {
5755+
if t.is_raw_legacy_type_variable() {
5756+
*t = self.heap.mk_any_error()
5757+
}
5758+
});
57575759
}
57585760

57595761
fn check_implicit_return_against_annotation(

0 commit comments

Comments
 (0)