Skip to content

Commit

Permalink
subtypes: fast path for Union/Union subtype check (#14277)
Browse files Browse the repository at this point in the history
Enums are exploded into Union of Literal when narrowed.

Conditional branches on enum values can result in multiple distinct
narrowing of the same enum which are later subject to subtype checks 
(most notably via `is_same_type`, when exiting frame context in the binder). 
Such checks would have quadratic complexity: `O(N*M)` where `N` and `M` 
are the number of entries in each narrowed enum variable, and led to drastic 
slowdown if any of the enums involved has a large number of values.

Implement a linear-time fast path where literals are quickly filtered, with
a fallback to the slow path for more complex values.

In our codebase there is one method with a chain of a dozen `if`
statements operating on instances of an enum with a hundreds of values. 
Prior to the regression it was typechecked in less than 1s. After the regression 
it takes over 13min to typecheck. This patch fully fixes the regression for us.

Fixes #13821.
  • Loading branch information
huguesb committed Dec 28, 2022
1 parent 61a21ba commit 1f8621c
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 0 deletions.
30 changes: 30 additions & 0 deletions mypy/subtypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
UninhabitedType,
UnionType,
UnpackType,
_flattened,
get_proper_type,
is_named_instance,
)
Expand Down Expand Up @@ -891,6 +892,35 @@ def visit_union_type(self, left: UnionType) -> bool:
if not self._is_subtype(item, self.orig_right):
return False
return True

elif isinstance(self.right, UnionType):
# prune literals early to avoid nasty quadratic behavior which would otherwise arise when checking
# subtype relationships between slightly different narrowings of an Enum
# we achieve O(N+M) instead of O(N*M)

fast_check: set[ProperType] = set()

for item in _flattened(self.right.relevant_items()):
p_item = get_proper_type(item)
if isinstance(p_item, LiteralType):
fast_check.add(p_item)
elif isinstance(p_item, Instance):
if p_item.last_known_value is None:
fast_check.add(p_item)
else:
fast_check.add(p_item.last_known_value)

for item in left.relevant_items():
p_item = get_proper_type(item)
if p_item in fast_check:
continue
lit_type = mypy.typeops.simple_literal_type(p_item)
if lit_type in fast_check:
continue
if not self._is_subtype(item, self.orig_right):
return False
return True

return all(self._is_subtype(item, self.orig_right) for item in left.items)

def visit_partial_type(self, left: PartialType) -> bool:
Expand Down
9 changes: 9 additions & 0 deletions mypy/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -3346,6 +3346,15 @@ def has_recursive_types(typ: Type) -> bool:
return typ.accept(_has_recursive_type)


def _flattened(types: Iterable[Type]) -> Iterable[Type]:
for t in types:
tp = get_proper_type(t)
if isinstance(tp, UnionType):
yield from _flattened(tp.items)
else:
yield t


def flatten_nested_unions(
types: Iterable[Type], handle_type_alias_type: bool = True
) -> list[Type]:
Expand Down

0 comments on commit 1f8621c

Please sign in to comment.