From c588852451ec63f9a26fefe5eb82976aa03d611c Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Tue, 20 Dec 2022 15:35:38 +0000 Subject: [PATCH] Speed up freshening type variables (#14323) Only perform type variable freshening if it's needed, i.e. there is a nested generic callable, since it's fairly expensive. Make the check for generic callables fast by creating a specialized type query visitor base class for queries with bool results. The visitor tries hard to avoid memory allocation in typical cases, since allocation is slow. This addresses at least some of the performance regression in #14095. This improved self-check performance by about 3% when compiled with mypyc (-O2). The new visitor class can potentially help with other type queries as well. I'll explore it in follow-up PRs. --- mypy/expandtype.py | 26 +++++++- mypy/type_visitor.py | 156 +++++++++++++++++++++++++++++++++++++++++++ mypy/types.py | 5 +- 3 files changed, 183 insertions(+), 4 deletions(-) diff --git a/mypy/expandtype.py b/mypy/expandtype.py index 1458fb74ce94..ca562ede264f 100644 --- a/mypy/expandtype.py +++ b/mypy/expandtype.py @@ -1,11 +1,14 @@ from __future__ import annotations from typing import Iterable, Mapping, Sequence, TypeVar, cast, overload +from typing_extensions import Final from mypy.nodes import ARG_POS, ARG_STAR, Var from mypy.type_visitor import TypeTranslator from mypy.types import ( + ANY_STRATEGY, AnyType, + BoolTypeQuery, CallableType, DeletedType, ErasedType, @@ -138,13 +141,30 @@ def freshen_function_type_vars(callee: F) -> F: return cast(F, fresh_overload) +class HasGenericCallable(BoolTypeQuery): + def __init__(self) -> None: + super().__init__(ANY_STRATEGY) + + def visit_callable_type(self, t: CallableType) -> bool: + return t.is_generic() or super().visit_callable_type(t) + + +# Share a singleton since this is performance sensitive +has_generic_callable: Final = HasGenericCallable() + + T = TypeVar("T", bound=Type) def freshen_all_functions_type_vars(t: T) -> T: - result = t.accept(FreshenCallableVisitor()) - assert isinstance(result, type(t)) - return result + result: Type + has_generic_callable.reset() + if not t.accept(has_generic_callable): + return t # Fast path to avoid expensive freshening + else: + result = t.accept(FreshenCallableVisitor()) + assert isinstance(result, type(t)) + return result class FreshenCallableVisitor(TypeTranslator): diff --git a/mypy/type_visitor.py b/mypy/type_visitor.py index 0f5ac05e68ac..823e74e7e283 100644 --- a/mypy/type_visitor.py +++ b/mypy/type_visitor.py @@ -15,6 +15,7 @@ from abc import abstractmethod from typing import Any, Callable, Generic, Iterable, Sequence, TypeVar, cast +from typing_extensions import Final from mypy_extensions import mypyc_attr, trait @@ -417,3 +418,158 @@ def visit_type_alias_type(self, t: TypeAliasType) -> T: def query_types(self, types: Iterable[Type]) -> T: """Perform a query for a list of types using the strategy to combine the results.""" return self.strategy([t.accept(self) for t in types]) + + +# Return True if at least one type component returns True +ANY_STRATEGY: Final = 0 +# Return True if no type component returns False +ALL_STRATEGY: Final = 1 + + +class BoolTypeQuery(SyntheticTypeVisitor[bool]): + """Visitor for performing recursive queries of types with a bool result. + + Use TypeQuery if you need non-bool results. + + 'strategy' is used to combine results for a series of types. It must + be ANY_STRATEGY or ALL_STRATEGY. + + Note: This visitor keeps an internal state (tracks type aliases to avoid + recursion), so it should *never* be re-used for querying different types + unless you call reset() first. + """ + + def __init__(self, strategy: int) -> None: + self.strategy = strategy + if strategy == ANY_STRATEGY: + self.default = False + else: + assert strategy == ALL_STRATEGY + self.default = True + # Keep track of the type aliases already visited. This is needed to avoid + # infinite recursion on types like A = Union[int, List[A]]. An empty set is + # represented as None as a micro-optimization. + self.seen_aliases: set[TypeAliasType] | None = None + # By default, we eagerly expand type aliases, and query also types in the + # alias target. In most cases this is a desired behavior, but we may want + # to skip targets in some cases (e.g. when collecting type variables). + self.skip_alias_target = False + + def reset(self) -> None: + """Clear mutable state (but preserve strategy). + + This *must* be called if you want to reuse the visitor. + """ + self.seen_aliases = None + + def visit_unbound_type(self, t: UnboundType) -> bool: + return self.query_types(t.args) + + def visit_type_list(self, t: TypeList) -> bool: + return self.query_types(t.items) + + def visit_callable_argument(self, t: CallableArgument) -> bool: + return t.typ.accept(self) + + def visit_any(self, t: AnyType) -> bool: + return self.default + + def visit_uninhabited_type(self, t: UninhabitedType) -> bool: + return self.default + + def visit_none_type(self, t: NoneType) -> bool: + return self.default + + def visit_erased_type(self, t: ErasedType) -> bool: + return self.default + + def visit_deleted_type(self, t: DeletedType) -> bool: + return self.default + + def visit_type_var(self, t: TypeVarType) -> bool: + return self.query_types([t.upper_bound] + t.values) + + def visit_param_spec(self, t: ParamSpecType) -> bool: + return self.default + + def visit_type_var_tuple(self, t: TypeVarTupleType) -> bool: + return self.default + + def visit_unpack_type(self, t: UnpackType) -> bool: + return self.query_types([t.type]) + + def visit_parameters(self, t: Parameters) -> bool: + return self.query_types(t.arg_types) + + def visit_partial_type(self, t: PartialType) -> bool: + return self.default + + def visit_instance(self, t: Instance) -> bool: + return self.query_types(t.args) + + def visit_callable_type(self, t: CallableType) -> bool: + # FIX generics + # Avoid allocating any objects here as an optimization. + args = self.query_types(t.arg_types) + ret = t.ret_type.accept(self) + if self.strategy == ANY_STRATEGY: + return args or ret + else: + return args and ret + + def visit_tuple_type(self, t: TupleType) -> bool: + return self.query_types(t.items) + + def visit_typeddict_type(self, t: TypedDictType) -> bool: + return self.query_types(list(t.items.values())) + + def visit_raw_expression_type(self, t: RawExpressionType) -> bool: + return self.default + + def visit_literal_type(self, t: LiteralType) -> bool: + return self.default + + def visit_star_type(self, t: StarType) -> bool: + return t.type.accept(self) + + def visit_union_type(self, t: UnionType) -> bool: + return self.query_types(t.items) + + def visit_overloaded(self, t: Overloaded) -> bool: + return self.query_types(t.items) # type: ignore[arg-type] + + def visit_type_type(self, t: TypeType) -> bool: + return t.item.accept(self) + + def visit_ellipsis_type(self, t: EllipsisType) -> bool: + return self.default + + def visit_placeholder_type(self, t: PlaceholderType) -> bool: + return self.query_types(t.args) + + def visit_type_alias_type(self, t: TypeAliasType) -> bool: + # Skip type aliases already visited types to avoid infinite recursion. + # TODO: Ideally we should fire subvisitors here (or use caching) if we care + # about duplicates. + if self.seen_aliases is None: + self.seen_aliases = set() + elif t in self.seen_aliases: + return self.default + self.seen_aliases.add(t) + if self.skip_alias_target: + return self.query_types(t.args) + return get_proper_type(t).accept(self) + + def query_types(self, types: list[Type] | tuple[Type, ...]) -> bool: + """Perform a query for a sequence of types using the strategy to combine the results.""" + # Special-case for lists and tuples to allow mypyc to produce better code. + if isinstance(types, list): + if self.strategy == ANY_STRATEGY: + return any(t.accept(self) for t in types) + else: + return all(t.accept(self) for t in types) + else: + if self.strategy == ANY_STRATEGY: + return any(t.accept(self) for t in types) + else: + return all(t.accept(self) for t in types) diff --git a/mypy/types.py b/mypy/types.py index ab2caa96e535..5c1fe2a0e960 100644 --- a/mypy/types.py +++ b/mypy/types.py @@ -2879,7 +2879,10 @@ def get_proper_types(it: Iterable[Type | None]) -> list[ProperType] | list[Prope # to make it easier to gradually get modules working with mypyc. # Import them here, after the types are defined. # This is intended as a re-export also. -from mypy.type_visitor import ( # noqa: F811 +from mypy.type_visitor import ( # noqa: F811,F401 + ALL_STRATEGY as ALL_STRATEGY, + ANY_STRATEGY as ANY_STRATEGY, + BoolTypeQuery as BoolTypeQuery, SyntheticTypeVisitor as SyntheticTypeVisitor, TypeQuery as TypeQuery, TypeTranslator as TypeTranslator,