Speed up caching of subtype checks (#12539)

This is very performance critical. Implement a few micro-optimizations
to speed caching a bit. In particular, we use dict.get to reduce the
number of dict lookups required, and avoid tuple concatenation which
tends to be a bit slow, as it has to construct temporary objects.

It would probably be even better to avoid using tuples as keys
altogether. This could be a reasonable follow-up improvement.

Avoid caching if last known value is set, since it reduces the
likelihood of cache hits a lot, because the space of literal values
is big (essentially infinite).

Also make the global strict_optional attribute an instance-level
attribute for faster access, as we might now use it more frequently.

I extracted the cached subtype check code into a microbenchmark
and the new implementation seems about twice as fast (in an
artificial setting, though).

Work on #12526 (but should generally make things a little better).

Author: JukkaL

mypy/checker.py

@@ -83,7 +83,8 @@
 from mypy.plugin import Plugin, CheckerPluginInterface
 from mypy.sharedparse import BINARY_MAGIC_METHODS
 from mypy.scope import Scope
-from mypy import state, errorcodes as codes
+from mypy import errorcodes as codes
+from mypy.state import state
 from mypy.traverser import has_return_statement, all_return_statements
 from mypy.errorcodes import ErrorCode, UNUSED_AWAITABLE, UNUSED_COROUTINE
 from mypy.util import is_typeshed_file, is_dunder, is_sunder

mypy/join.py

@@ -17,7 +17,7 @@
 )
 from mypy.nodes import INVARIANT, COVARIANT, CONTRAVARIANT
 import mypy.typeops
-from mypy import state
+from mypy.state import state
 
 
 class InstanceJoiner:

mypy/meet.py

@@ -12,7 +12,7 @@
 from mypy.erasetype import erase_type
 from mypy.maptype import map_instance_to_supertype
 from mypy.typeops import tuple_fallback, make_simplified_union, is_recursive_pair
-from mypy import state
+from mypy.state import state
 from mypy import join
 
 # TODO Describe this module.

mypy/semanal_main.py

@@ -32,7 +32,7 @@
     MypyFile, TypeInfo, FuncDef, Decorator, OverloadedFuncDef, Var
 )
 from mypy.semanal_typeargs import TypeArgumentAnalyzer
-from mypy.state import strict_optional_set
+import mypy.state
 from mypy.semanal import (
     SemanticAnalyzer, apply_semantic_analyzer_patches, remove_imported_names_from_symtable
 )
@@ -356,7 +356,7 @@ def check_type_arguments(graph: 'Graph', scc: List[str], errors: Errors) -> None
                                         state.options,
                                         is_typeshed_file(state.path or ''))
         with state.wrap_context():
-            with strict_optional_set(state.options.strict_optional):
+            with mypy.state.state.strict_optional_set(state.options.strict_optional):
                 state.tree.accept(analyzer)
 
 
@@ -371,7 +371,7 @@ def check_type_arguments_in_targets(targets: List[FineGrainedDeferredNode], stat
                                     state.options,
                                     is_typeshed_file(state.path or ''))
     with state.wrap_context():
-        with strict_optional_set(state.options.strict_optional):
+        with mypy.state.state.strict_optional_set(state.options.strict_optional):
             for target in targets:
                 func: Optional[Union[FuncDef, OverloadedFuncDef]] = None
                 if isinstance(target.node, (FuncDef, OverloadedFuncDef)):

mypy/state.py

@@ -1,20 +1,28 @@
 from contextlib import contextmanager
 from typing import Optional, Tuple, Iterator
 
+from typing_extensions import Final
+
 # These are global mutable state. Don't add anything here unless there's a very
 # good reason.
 
-# Value varies by file being processed
-strict_optional = False
-find_occurrences: Optional[Tuple[str, str]] = None
 
+class StrictOptionalState:
+    # Wrap this in a class since it's faster that using a module-level attribute.
+
+    def __init__(self, strict_optional: bool) -> None:
+        # Value varies by file being processed
+        self.strict_optional = strict_optional
 
-@contextmanager
-def strict_optional_set(value: bool) -> Iterator[None]:
-    global strict_optional
-    saved = strict_optional
-    strict_optional = value
-    try:
-        yield
-    finally:
-        strict_optional = saved
+    @contextmanager
+    def strict_optional_set(self, value: bool) -> Iterator[None]:
+        saved = self.strict_optional
+        self.strict_optional = value
+        try:
+            yield
+        finally:
+            self.strict_optional = saved
+
+
+state: Final = StrictOptionalState(strict_optional=False)
+find_occurrences: Optional[Tuple[str, str]] = None

mypy/stubtest.py

@@ -21,6 +21,7 @@
 
 import mypy.build
 import mypy.modulefinder
+import mypy.state
 import mypy.types
 from mypy import nodes
 from mypy.config_parser import parse_config_file
@@ -560,7 +561,7 @@ def get_position(arg_name: str) -> int:
             return max(index for _, index in all_args[arg_name])
 
         def get_type(arg_name: str) -> mypy.types.ProperType:
-            with mypy.state.strict_optional_set(True):
+            with mypy.state.state.strict_optional_set(True):
                 all_types = [
                     arg.variable.type or arg.type_annotation for arg, _ in all_args[arg_name]
                 ]
@@ -1099,7 +1100,7 @@ def is_subtype_helper(left: mypy.types.Type, right: mypy.types.Type) -> bool:
         # Special case checks against TypedDicts
         return True
 
-    with mypy.state.strict_optional_set(True):
+    with mypy.state.state.strict_optional_set(True):
         return mypy.subtypes.is_subtype(left, right)
 
 

mypy/subtypes.py

@@ -25,7 +25,7 @@
 from mypy.expandtype import expand_type_by_instance
 from mypy.typestate import TypeState, SubtypeKind
 from mypy.options import Options
-from mypy import state
+from mypy.state import state
 
 # Flags for detected protocol members
 IS_SETTABLE: Final = 1
@@ -206,7 +206,8 @@ def build_subtype_kind(*,
                            ignore_pos_arg_names: bool = False,
                            ignore_declared_variance: bool = False,
                            ignore_promotions: bool = False) -> SubtypeKind:
-        return (False,  # is proper subtype?
+        return (state.strict_optional,
+                False,  # is proper subtype?
                 ignore_type_params,
                 ignore_pos_arg_names,
                 ignore_declared_variance,
@@ -1316,7 +1317,11 @@ def build_subtype_kind(*,
                            ignore_promotions: bool = False,
                            erase_instances: bool = False,
                            keep_erased_types: bool = False) -> SubtypeKind:
-        return True, ignore_promotions, erase_instances, keep_erased_types
+        return (state.strict_optional,
+                True,
+                ignore_promotions,
+                erase_instances,
+                keep_erased_types)
 
     def _is_proper_subtype(self, left: Type, right: Type) -> bool:
         return is_proper_subtype(left, right,

mypy/suggestions.py

@@ -27,7 +27,7 @@
 )
 from typing_extensions import TypedDict
 
-from mypy.state import strict_optional_set
+from mypy.state import state
 from mypy.types import (
     Type, AnyType, TypeOfAny, CallableType, UnionType, NoneType, Instance, TupleType,
     TypeVarType, FunctionLike, UninhabitedType,
@@ -439,7 +439,7 @@ def get_suggestion(self, mod: str, node: FuncDef) -> PyAnnotateSignature:
 
         is_method = bool(node.info) and not node.is_static
 
-        with strict_optional_set(graph[mod].options.strict_optional):
+        with state.strict_optional_set(graph[mod].options.strict_optional):
             guesses = self.get_guesses(
                 is_method,
                 self.get_starting_type(node),
@@ -454,7 +454,7 @@ def get_suggestion(self, mod: str, node: FuncDef) -> PyAnnotateSignature:
         # Now try to find the return type!
         self.try_type(node, best)
         returns = get_return_types(self.manager.all_types, node)
-        with strict_optional_set(graph[mod].options.strict_optional):
+        with state.strict_optional_set(graph[mod].options.strict_optional):
             if returns:
                 ret_types = generate_type_combinations(returns)
             else:
@@ -988,7 +988,7 @@ def refine_union(t: UnionType, s: ProperType) -> Type:
 
     # Turn strict optional on when simplifying the union since we
     # don't want to drop Nones.
-    with strict_optional_set(True):
+    with state.strict_optional_set(True):
         return make_simplified_union(new_items)
 
 

mypy/test/testtypes.py

@@ -17,7 +17,7 @@
 from mypy.nodes import ARG_POS, ARG_OPT, ARG_STAR, ARG_STAR2, CONTRAVARIANT, INVARIANT, COVARIANT
 from mypy.subtypes import is_subtype, is_more_precise, is_proper_subtype
 from mypy.test.typefixture import TypeFixture, InterfaceTypeFixture
-from mypy.state import strict_optional_set
+from mypy.state import state
 from mypy.typeops import true_only, false_only, make_simplified_union
 
 
@@ -410,15 +410,15 @@ def test_true_only_of_union(self) -> None:
         assert to.items[1] is tup_type
 
     def test_false_only_of_true_type_is_uninhabited(self) -> None:
-        with strict_optional_set(True):
+        with state.strict_optional_set(True):
             fo = false_only(self.tuple(AnyType(TypeOfAny.special_form)))
             assert_type(UninhabitedType, fo)
 
     def test_false_only_tuple(self) -> None:
-        with strict_optional_set(False):
+        with state.strict_optional_set(False):
             fo = false_only(self.tuple(self.fx.a))
             assert_equal(fo, NoneType())
-        with strict_optional_set(True):
+        with state.strict_optional_set(True):
             fo = false_only(self.tuple(self.fx.a))
             assert_equal(fo, UninhabitedType())
 
@@ -437,7 +437,7 @@ def test_false_only_of_instance(self) -> None:
         assert self.fx.a.can_be_true
 
     def test_false_only_of_union(self) -> None:
-        with strict_optional_set(True):
+        with state.strict_optional_set(True):
             tup_type = self.tuple()
             # Union of something that is unknown, something that is always true, something
             # that is always false
@@ -1059,9 +1059,9 @@ def test_literal_type(self) -> None:
     # FIX generic interfaces + ranges
 
     def assert_meet_uninhabited(self, s: Type, t: Type) -> None:
-        with strict_optional_set(False):
+        with state.strict_optional_set(False):
             self.assert_meet(s, t, self.fx.nonet)
-        with strict_optional_set(True):
+        with state.strict_optional_set(True):
             self.assert_meet(s, t, self.fx.uninhabited)
 
     def assert_meet(self, s: Type, t: Type, meet: Type) -> None:

mypy/typeops.py

@@ -26,7 +26,7 @@
 
 from mypy.typevars import fill_typevars
 
-from mypy import state
+from mypy.state import state
 
 
 def is_recursive_pair(s: Type, t: Type) -> bool:

mypy/types.py

@@ -12,7 +12,7 @@
 
 from mypy.backports import OrderedDict
 import mypy.nodes
-from mypy import state
+from mypy.state import state
 from mypy.nodes import (
     INVARIANT, SymbolNode, FuncDef,
     ArgKind, ARG_POS, ARG_STAR, ARG_STAR2,

mypy/typestate.py

@@ -9,7 +9,6 @@
 from mypy.nodes import TypeInfo
 from mypy.types import Instance, TypeAliasType, get_proper_type, Type
 from mypy.server.trigger import make_trigger
-from mypy import state
 
 # Represents that the 'left' instance is a subtype of the 'right' instance
 SubtypeRelationship: _TypeAlias = Tuple[Instance, Instance]
@@ -124,20 +123,29 @@ def reset_all_subtype_caches_for(info: TypeInfo) -> None:
 
     @staticmethod
     def is_cached_subtype_check(kind: SubtypeKind, left: Instance, right: Instance) -> bool:
+        if left.last_known_value is not None or right.last_known_value is not None:
+            # If there is a literal last known value, give up. There
+            # will be an unbounded number of potential types to cache,
+            # making caching less effective.
+            return False
         info = right.type
-        if info not in TypeState._subtype_caches:
+        cache = TypeState._subtype_caches.get(info)
+        if cache is None:
             return False
-        cache = TypeState._subtype_caches[info]
-        key = (state.strict_optional,) + kind
-        if key not in cache:
+        subcache = cache.get(kind)
+        if subcache is None:
             return False
-        return (left, right) in cache[key]
+        return (left, right) in subcache
 
     @staticmethod
     def record_subtype_cache_entry(kind: SubtypeKind,
                                    left: Instance, right: Instance) -> None:
+        if left.last_known_value is not None or right.last_known_value is not None:
+            # These are unlikely to match, due to the large space of
+            # possible values.  Avoid uselessly increasing cache sizes.
+            return
         cache = TypeState._subtype_caches.setdefault(right.type, dict())
-        cache.setdefault((state.strict_optional,) + kind, set()).add((left, right))
+        cache.setdefault(kind, set()).add((left, right))
 
     @staticmethod
     def reset_protocol_deps() -> None:

mypyc/irbuild/main.py

@@ -25,7 +25,7 @@ def f(x: int) -> int:
 
 from mypy.nodes import MypyFile, Expression, ClassDef
 from mypy.types import Type
-from mypy.state import strict_optional_set
+from mypy.state import state
 from mypy.build import Graph
 
 from mypyc.common import TOP_LEVEL_NAME
@@ -45,7 +45,7 @@ def f(x: int) -> int:
 # The stubs for callable contextmanagers are busted so cast it to the
 # right type...
 F = TypeVar('F', bound=Callable[..., Any])
-strict_optional_dec = cast(Callable[[F], F], strict_optional_set(True))
+strict_optional_dec = cast(Callable[[F], F], state.strict_optional_set(True))
 
 
 @strict_optional_dec  # Turn on strict optional for any type manipulations we do