[red-knot] Statically known branches

[sharkdp]

Summary

Rendered version of the test suite including a proper introduction to the topic / motivation: click.

This changeset adds support for precise type-inference and boundness-handling of definitions inside control-flow branches with statically-known conditions, i.e. test-expressions whose truthiness we can unambiguously infer as always false or always true. In code:

x = 1

if "z" in "haystack":  # Literal[False]
    x = 2

reveal_type(x)  # revealed: Literal[1]

and:

x = 1

if "y" in "haystack":  # Literal[True]
    x = 2

reveal_type(x)  # revealed: Literal[2]

One option to implement this would have been to add special handling for a limited set of test-expressions in semantic index-building. We would then analyze expressions like sys.version_info >= (x, y), typing.TYPE_CHECKING, True and False without any type inference and consequently close down (or unconditionally open) branches whose truthiness we can analyze in this way. This would simplify the implementation, but is much less general than the approach taken here.

Instead, we collect all necessary information during semantic index building, and then re-analyze the control flow during type-checking. The way this works is by recording the list of 'active' branching conditions for each introduced binding. Consider this example:

if test1:
    x = 1  # active conditions: [test1]

    if test2:
        x = 2  # active conditions: [test1, test2]

if test3:
    x = 3  # active conditions: [test3]

use(x)

In semantic-index building, when we reach the use(x) statement, we have three live bindings. But when we later analyze the test-conditions in type inference, we may potentially conclude that some of them are not present/visible.

One possibility to do so is to detect test-conditions to be always false. For example, if test3 is statically known to be false, we can infer that the x = 3 branch is never taken. This binding is then not considered in type inference and in boundness-handling. For the x = 2 binding, we need to consider the possibility that either test1 or test2 can be statically false. In summary, if any of the active branching conditions is statically false, the binding is cancelled.

Another possibility is that we can infer a test-condition to be always true. This can also hide bindings. For example, if we can infer test2 to always be true, we know that the x = 1 binding is never visible. This is slightly more complex to handle, as we haven't even seen the test2 condition when we analyze the x = 1 binding in semantic-index building.

The way we solve this in type-inference is by going through all live bindings in reverse order (starting at x = 3), stopping whenever we detect a binding that is unconditionally visible. A binding is unconditionally visible if all of its active conditions are statically known to be true. This is slightly complicated by the fact that we have unconditional branching as well. For example, consider:

if test1:
    for _ in iterable():
        x = 1

This binding is not unconditionally visible, even if test1 is always true, as the loop body may never be executed. This is currently handled by adding a BranchingCondition::Unconditional marker to the active conditions of the binding, which always evaluates to an ambiguous truthiness.

Some of this might sound similar to what we achieve with the may_be_unbound and may_be_undeclared flags. However, the new mechanism does not allow us to get rid of these flags, as there are also cases like this:

if test:  # type: bool
    x = 1
else:
    x = 2

The symbol x is always bound, but seeing the x = 1 or x = 2 definitions with a branching condition of test: bool is not enough to see that. We also need the global information that these two branches will be merged (where we logically OR the two may_be_unbound = false flags to a global may_be_unbound = false).

This branch also includes:

• statically-known branches handling for Boolean expressions and while loops
• support for typing.TYPE_CHECKING
• support for narrowing in while loops
• new target-version requirements in some Markdown tests which were now required due to the understanding of sys.version_info branches.

closes #12700
closes #14170
closes #14861

Test plan

• New Markdown tests for the main feature in statically-known-branches.md
• New Markdown tests for typing.TYPE_CHECKING in known_constants.md
• New Markdown tests for while-loop narrowing in narrow/while.md
• Regression test for boundness-handling in while loops
• Adapted tests for EllipsisType, Never, etc

Performance investigation

(thanks to @MichaReiser for some input!)

Measurements

I can reproduce the regression locally. It's not as pronounced as on codspeed, which could be attributed to the fact that the codspeed results do not involve any I/O overhead (they are based on CPU counters, not wall time). So the local results might be "diluted" a bit:

tomllib, -13%

Command	Mean [ms]	Min [ms]	Max [ms]	Relative
./red_knot_main	21.8 ± 1.5	18.8	27.0	1.00
./red_knot_feature	24.6 ± 1.6	21.2	29.1	1.13 ± 0.11

Black codebase, -7%

Command	Mean [ms]	Min [ms]	Max [ms]	Relative
./red_knot_main	123.4 ± 5.4	115.6	135.4	1.00
./red_knot_feature	131.6 ± 3.5	123.7	137.3	1.07 ± 0.05

Analysis (ongoing)

• With this feature enabled, we need to perform more work. For every binding/declaration, we need to infer the type of all active branching conditions and call Type::bool on it (which may involve looking up other symbols). This can entail loading entire new modules. For example: in the tomllib benchmark, we do not need to resolve sys on main, but we do on this branch. The reason for this are various sys.version_info >= (M, m) conditions in builtins.pyi which now need to be resolved. I tried to account for this by adding a import sys; sys.version_info line in a local copy of tomllib, so the version on main would also have to import sys. This reduces the difference from 13% to 11%, but the comparison is still not fair. There are a lot more symbol-lookups on this feature branch (mostly __bool__).
• In semantic-index building, we need to collect, store (and clone, when snapshotting) more information with the addition of the branching-condition data.

To do

• Investigate performance regression

[codspeed-hq]

CodSpeed Performance Report

Merging #14759 will degrade performances by 21.62%

_{Comparing david/statically-known-branches (0a0e5e6) with main (c3a64b4)}

Summary

❌ 2 regressions
✅ 30 untouched benchmarks

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Benchmarks breakdown

	Benchmark	main	david/statically-known-branches	Change
❌	red_knot_check_file[cold]	68 ms	86.8 ms	-21.62%
❌	red_knot_check_file[incremental]	4.2 ms	4.4 ms	-4.31%

[carljm]

Impressively comprehensive tests, thank you!!

Per request, I only looked at the mdtests.

Just a few comments, nothing significant.

[sharkdp]

We are "too smart" now and infer Literal[3] here, as we can see that the first pattern can never match. To keep the original intent of the test, I changed the match target to an unknown int.

[sharkdp]

I just realized that this is problematic in cases with multiple "always True" patterns. At runtime, we'll take the first of these branches. But in our current implementation here, since we always walk definitions backwards, we'll take the last instead:

x = 1

match "a":
    case "a" if 1 == 1:
        x = 2
    case "a":
        x = 3

reveal_type(x)  # we infer Literal[3], but it should be Literal[2]

Before I attempt to fix this, it would be good to get some feedback on whether or not we want static-truthiness checking in match statements at all. I thought it was a nice stretch-task, but maybe it causes more trouble than it's worth. The examples presented here are not particularly useful, but there might be some use cases involving enums? Not sure.

[carljm]

I don't think we should spend a lot of extra time on implementation work that is specific to statically-known branches with match patterns; it's nice to have, but not clear how often it will be relevant.

I'm curious why this seems to affect match statements with multiple always-matching cases, but not elif chains with multiple always-true branches? The two would seem to have a lot in common. In both cases the key fact is that an always-true condition in an elif or case does not mean "we definitely take this path" (we might have already taken a previous one instead). It only means "we definitely do not take any branch after this one."

[sharkdp]

I'm curious why this seems to affect match statements with multiple always-matching cases, but not elif chains with multiple always-true branches?

It currently works for elif chains but not for match patterns because of negative constraints. The way it currently works is that we record all constraints (including negative ones) when a binding is created:

if flag:
    x = 1  # conditional on flag
elif True:
    x = 2  # conditional on ~flag and True
elif True:
    x = 3  # conditional on ~flag, ~True and True

The x = 3 binding is therefore not unconditionally visible, which currently leads us to correctly infer Literal[1, 2] for this case.

But now that I write it out like this, I also realize that this reveals a problem with my approach. Consider this example:

x = 1

if flag:
    x = 2  # conditional on flag
elif True:
    x = 3  # conditional on ~flag, True

Here, we currently infer Literal[1, 2, 3], while it would be more correct to infer Literal[2, 3]. The problem is that we have the negative ~flag condition in the x = 3 binding, which prevents us from detecting it as "unconditionally visible". And if we would detect it as unconditionally visible, we would run into the problem you mentioned. We would block off any previous bindings and incorrectly infer Literal[3].

We could clearly do better.

I'm currently thinking that there are two problems:

• It seems to me like the set of constraints that we apply for static truthiness analysis needs to be different from the set of constraints that we apply for narrowing. For narrowing, we need that ~flag constraint. But for static truthiness analysis, we don't care if flag has ambiguous truthiness or not, we always take an elif True branch.
• When iterating the list of bindings in reverse, we should not stop at the first "unconditionally visible" binding, we need to go back to the earliest possible "unconditionally visible" binding.