[ty] Optimize TDD atom ordering

[sharkdp]

Summary

While looking at some logging output that I added to ReachabilityConstraintBuilder::add_and_constraint in order to debug astral-sh/ty#1091, I noticed that it seemed to suggest that the TDD was built in an imbalanced way for code like the following, where we have a sequence of non-nested if conditions:

def f(t1, t2, t3, t4, …):
    x = 0
    if t1:
        x = 1
    if t2:
        x = 2
    if t3:
        x = 3
    if t4:
        x = 4
    …

To understand this a bit better, I added some code to the ReachabilityConstraintBuilder to render the resulting TDD. On main, we get a tree that looks like the following, where you can see a pattern of N sub-trees that grow linearly with N (number of if statements). This results in an overall tree structure that has N² nodes (see graph below):

If we zoom in to one of these subgraphs, we can see what the problem is. When we add new constraints that represent combinations like t1 AND ~t2 AND ~t3 AND t4 AND …, they start with the evaluation of "early" conditions (t1, t2, …). This means that we have to create new subgraphs for each new if condition because there is little sharing with the previous structure. We evaluate the Boolean condition in a right-associative way: t1 AND (~t2 AND (~t3 AND t4))):

If we change the ordering of TDD atoms, we can change that to a left-associative evaluation: (((t1 AND ~t2) AND ~t3) AND t4) …. This means that we can re-use previous subgraphs (t1 AND ~t2), which results in a much more compact graph structure overall (note how "late" conditions are now at the top, and "early" conditions are further down in the graph):

If we count the number of TDD nodes for a growing number if if statements, we can see that this change results in a slower growth. It's worth noting that the growth is still superlinear, though:

On the actual code from the referenced ticket (the t_main.py file reduced to its main function, with the main function limited to 2000 lines instead of 11000 to allow the version on main to run to completion), the effect is much more dramatic. Instead of 26 million TDD nodes (main), we now only create 250 thousand (this branch), which is slightly less than 1%.

The change in this PR allows us to build the semantic index and type-check the problematic t_main.py file in astral-sh/ty#1091 in 9 seconds. This is still not great, but an obvious improvement compared to running out of memory after minutes of execution.

An open question remains whether this change is beneficial for all kinds of code patterns, or just this linear sequence of if statements. It does not seem unreasonable to think that referring to "earlier" conditions is generally a good idea, but I learned from Doug that it's generally not possible to find a TDD-construction heuristic that is non-pathological for all kinds of inputs. Fortunately, it seems like this change here results in performance improvements across all of our benchmarks, which should increase the confidence in this change:

Benchmark	Improvement
hydra-zen	+13%
DateType	+5%
sympy (walltime)	+4%
attrs	+4%
pydantic (walltime)	+2%
pandas (walltime)	+2%
altair (walltime)	+2%
static-frame	+2%
anyio	+1%
freqtrade	+1%
colour-science	+1%
tanjun	+1%

closes astral-sh/ty#1091

[github-actions]

Diagnostic diff on typing conformance tests

No changes detected when running ty on typing conformance tests ✅

[github-actions]

mypy_primer results

No ecosystem changes detected ✅

Memory usage changes were detected when running on open source projects

sphinx (https://github.com/sphinx-doc/sphinx)
- TOTAL MEMORY USAGE: ~273MB
+ TOTAL MEMORY USAGE: ~260MB

[codspeed-hq]

CodSpeed Instrumentation Performance Report

Merging #20098 will improve performances by 13.81%

_{Comparing david/fix-1091 (8e988ae) with main (5663426)}

Summary

⚡ 2 improvements
✅ 40 untouched benchmarks

Benchmarks breakdown

	Benchmark	BASE	HEAD	Change
⚡	attrs	375.1 ms	360.4 ms	+4.09%
⚡	hydra-zen	803.6 ms	706.1 ms	+13.81%

[carljm]

Awesome investigation and results!