Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[red-knot] per-definition inference, use-def maps #12269

Merged
merged 3 commits into from
Jul 16, 2024
Merged

[red-knot] per-definition inference, use-def maps #12269

merged 3 commits into from
Jul 16, 2024

Conversation

carljm
Copy link
Contributor

@carljm carljm commented Jul 10, 2024

Implements definition-level type inference, with basic control flow (only if statements and if expressions so far) in Salsa.

There are a couple key ideas here:

  1. We can do type inference queries at any of three region granularities: an entire scope, a single definition, or a single expression. These are represented by the InferenceRegion enum, and the entry points are the salsa queries infer_scope_types, infer_definition_types, and infer_expression_types. Generally per-scope will be used for scopes that we are directly checking and per-definition will be used anytime we are looking up symbol types from another module/scope. Per-expression should be uncommon: used only for the RHS of an unpacking or multi-target assignment (to avoid re-inferring the RHS once per symbol defined in the assignment) and for test nodes in type narrowing (e.g. the test of an If node). All three queries return a TypeInference with a map of types for all definitions and expressions within their region. If you do e.g. scope-level inference, when it hits a definition, or an independently-inferable expression, it should use the relevant query (which may already be cached) to get all types within the smaller region. This avoids double-inferring smaller regions, even though larger regions encompass smaller ones.

  2. Instead of building a control-flow graph and lazily traversing it to find definitions which reach a use of a name (which is O(n^2) in the worst case), instead semantic indexing builds a use-def map, where every use of a name knows which definitions can reach that use. We also no longer track all definitions of a symbol in the symbol itself; instead the use-def map also records which defs remain visible at the end of the scope, and considers these the publicly-visible definitions of the symbol (see below).

Major items left as TODOs in this PR, to be done in follow-up PRs:

  1. Free/global references aren't supported yet (only lookup based on definitions in current scope), which means the override-check example doesn't currently work. This is the first thing I'll fix as follow-up to this PR.

  2. Control flow outside of if statements and expressions.

  3. Type narrowing.

There are also some smaller relevant changes here:

  1. Eliminate Option in the return type of member lookups; instead always return Type::Unbound for a name we can't find. Also use Type::Unbound for modules we can't resolve (not 100% sure about this one yet.)

  2. Eliminate the use of the terms "public" and "root" to refer to module-global scope or symbols. Instead consistently use the term "module-global". It's longer, but it's the clearest, and the most consistent with typical Python terminology. In particular I don't like "public" for this use because it has other implications around author intent (is an underscore-prefixed module-global symbol "public"?). And "root" is just not commonly used for this in Python.

  3. Eliminate the PublicSymbol Salsa ingredient. Many non-module-global symbols can also be seen from other scopes (e.g. by a free var in a nested scope, or by class attribute access), and thus need to have a "public type" (that is, the type not as seen from a particular use in the control flow of the same scope, but the type as seen from some other scope.) So all symbols need to have a "public type" (here I want to keep the use of the term "public", unless someone has a better term to suggest -- since it's "public type of a symbol" and not "public symbol" the confusion with e.g. initial underscores is less of an issue.) At least initially, I would like to try not having special handling for module-global symbols vs other symbols.

  4. Switch to using "definitions that reach end of scope" rather than "all definitions" in determining the public type of a symbol. I'm convinced that in general this is the right way to go. We may want to refine this further in future for some free-variable cases, but it can be changed purely by making changes to the building of the use-def map (the public_definitions index in it), without affecting any other code. One consequence of combining this with no control-flow support (just last-definition-wins) is that some inference tests now give more wrong-looking results; I left TODO comments on these tests to fix them when control flow is added.

And some potential areas for consideration in the future:

  1. Should symbol_ty be a Salsa query? This would require making all symbols a Salsa ingredient, and tracking even more dependencies. But it would save some repeated reconstruction of unions, for symbols with multiple public definitions. For now I'm not making it a query, but open to changing this in future with actual perf evidence that it's better.

@carljm carljm added the red-knot Multi-file analysis & type inference label Jul 10, 2024
Copy link
Member

@MichaReiser MichaReiser left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't fully understand all of it yet, but I like what I'm seeing.

crates/red_knot_python_semantic/src/types.rs Outdated Show resolved Hide resolved
Comment on lines 28 to 30
let Some(first) = def_types.next() else {
return Type::Unbound;
};
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
let Some(first) = def_types.next() else {
return Type::Unbound;
};
let first = def_types.next().unwrap_or(Type::Unbound);

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this better even though it means we'll make an entirely unnecessary second call to def_types.next() that we already know will return None? It seems like we may as well short-circuit here and just return Unbound directly.

crates/red_knot_python_semantic/src/types.rs Show resolved Hide resolved
@carljm carljm force-pushed the cjm/defs branch 2 times, most recently from 524f359 to 3cbbd50 Compare July 11, 2024 20:09
Copy link
Contributor

github-actions bot commented Jul 11, 2024

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

Formatter (stable)

✅ ecosystem check detected no format changes.

Formatter (preview)

✅ ecosystem check detected no format changes.

Copy link

codspeed-hq bot commented Jul 12, 2024

CodSpeed Performance Report

Merging #12269 will improve performances by 6.95%

Comparing cjm/defs (c7e4ef9) with main (30cef67)

Summary

⚡ 2 improvements
✅ 31 untouched benchmarks

Benchmarks breakdown

Benchmark main cjm/defs Change
red_knot_check_file[without_parse] 238.7 µs 228.9 µs +4.28%
red_knot_check_file[incremental] 90.2 µs 84.4 µs +6.95%

@carljm carljm marked this pull request as ready for review July 12, 2024 06:22
@carljm carljm requested a review from AlexWaygood as a code owner July 12, 2024 06:22
@carljm
Copy link
Contributor Author

carljm commented Jul 12, 2024

This is ready for review. I've updated the PR description with a fuller narrative of what's in here and why.

Benchmarks seem to say this PR has reclaimed most of the perf that was lost in the Salsa-type-interning PR, but I'm not sure why. It may be due to the removal of the ancestor-scope-walking name lookup code in favor of the use-def map.

Copy link
Member

@MichaReiser MichaReiser left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like what I'm seeing here.

My only concern is that I don't understand how the different type inference scopes are merged. See my inline comment.

It will be interesting to see what the cost of storing the definition -> Ty and expression -> Ty hash maps on every level. That's a lot of hash maps and a lot of redundant information that we need to keep around.

@carljm
Copy link
Contributor Author

carljm commented Jul 12, 2024

It will be interesting to see what the cost of storing the definition -> Ty and expression -> Ty hash maps on every level. That's a lot of hash maps and a lot of redundant information that we need to keep around.

Yeah :/ That's the part of this approach that I'm least happy about.

We could reduce the number of hashmaps, and the redundancy, by using some kind of chaining (where an outer TypeInference will re-request the sub-query -- should be cached -- when needed to get sub-results instead of actually merging them into its own results), but that means more cost on lookup. I'm not sure how the tradeoffs will shake out.

I'll think more about this.

@carljm
Copy link
Contributor Author

carljm commented Jul 12, 2024

Thanks for the great review!

@MichaReiser
Copy link
Member

MichaReiser commented Jul 13, 2024

We could reduce the number of hashmaps, and the redundancy, by using some kind of chaining (where an outer TypeInference will re-request the sub-query -- should be cached -- when needed to get sub-results instead of actually merging them into its own results), but that means more cost on lookup. I'm not sure how the tradeoffs will shake out.

Here's an idea that's not in the salsa spirit but upholds the important constraints that matter to salsa.

  • Add an expression_types field to SemanticIndex which is an IndexVec<ScopeId, ExpressionTypes>. Adding it to SemanticIndex causes it to get wiped every time the AST changes.
  • ExpressionTypes internally stores a RwLock<IndexVec<ScopedExpressionId, Option<Ty>>> (or Mutex, see below) and the semantic_index query either initialize it as empty, or with the right size, but all values set to None
  • All type inference stores the inferred expression types in ExpressionTypes (and not locally).
  • ExpressionTypes allows lookups and writes by ScopedExpressionId (which you can get from any expression by calling scoped_ast_id()). It uses internal mutability for writing. A RwLock does allow for race conditions where two threads both test if an expression has been inferred, before they start inferring that expression (and then writing it back). But I think this kind of race is okay. Both queries will infer the same type.

This approach would mitigate most concerns. The only extra cost is that we need to test for every expression if it has already been inferred.

I'm unclear on how it would work with any fixpoint iteration. How would we know which expression types need to be invalidated? Or would we know if we run a fixpoint and could then forcefully re-infer all expression types?

Regarding the review: Please re-request review when you think I should have another look

@carljm carljm force-pushed the cjm/defs branch 2 times, most recently from 2ab2e96 to f589b4b Compare July 16, 2024 04:12
@carljm carljm requested a review from MichaReiser July 16, 2024 04:19
@carljm
Copy link
Contributor Author

carljm commented Jul 16, 2024

Ok, I think this is ready for review again.

The use-def map building is now better (uses a single vec of definitions and ranges, instead of vecs of vecs) but there are probably ways it could be still better; open to comments. The main complexity is that avoiding vecs-of-vecs means sticking to a single range for the "active" defs of a given symbol, which means we always have to keep all the active defs for a given symbol adjacent. Sometimes at control-flow merge points this requires copying the same definition IDs from earlier to later in the all-definitions vector.

I ended up adding basic control flow (just for if statements and expressions), because it makes a big difference in how the range-based use-def map building works.

Here's an idea that's not in the salsa spirit but upholds the important constraints that matter to salsa. ...

That's a neat idea. I would like to hold off on it though; it's purely a performance optimization that shouldn't make a difference to the API for type inference. So I would rather leave it until we are at a point where we can evaluate the performance impact on real code. For now I think we can move forward with some duplicated stored types and keep things more salsa-native. It's just extra id -> id mappings; not a ton of wasted memory.

Copy link
Member

@MichaReiser MichaReiser left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The type inference part looks good to me.

I started (for 1.5h) reviewing the use-def implementation but I must admit that I don't have enough context information to efficiently review the changes and the data structures are non-trivial. I would have preferred if adding use-def and then refactoring to a more efficient data structure were its own PR with a summary explaining how the data structures work together and what considerations have been made. But it's a bit too late for that. So I'm just gonna skip this part. We can come back to it later or set up a session where you walk me through the code.

Cargo.toml Outdated Show resolved Hide resolved
crates/red_knot_python_semantic/src/types/infer.rs Outdated Show resolved Hide resolved
crates/red_knot_python_semantic/src/types/infer.rs Outdated Show resolved Hide resolved
crates/red_knot_python_semantic/src/types/infer.rs Outdated Show resolved Hide resolved
crates/red_knot_python_semantic/src/types/infer.rs Outdated Show resolved Hide resolved
crates/red_knot_python_semantic/src/types/infer.rs Outdated Show resolved Hide resolved
@carljm
Copy link
Contributor Author

carljm commented Jul 16, 2024

Planning to merge this as-is once tests pass. All unresolved comments above will be addressed in follow-up PRs.

@carljm carljm merged commit 595b1aa into main Jul 16, 2024
20 checks passed
@carljm carljm deleted the cjm/defs branch July 16, 2024 18:02
carljm added a commit that referenced this pull request Jul 18, 2024
Per comments in #12269, "module
global" is kind of long, and arguably redundant.

I tried just using "module" but there were too many cases where I felt
this was ambiguous. I like the way "global" works out better, though it
does require an understanding that in Python "global" generally means
"module global" not "globally global" (though in a sense module globals
are also globally global since modules are singletons).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
red-knot Multi-file analysis & type inference
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants