[red-knot] Support re-export conventions for stub files

[dhruvmanila]

This is an alternative implementation to #15848.

Summary

This PR adds support for re-export conventions for imports for stub files.

How does this work?

• Add a new flag on the Import and ImportFrom definitions to indicate whether they're being exported or not
• Add a new enum to indicate whether the symbol lookup is happening within the same file or is being queried from another file (e.g., an import statement)
• When a Symbol is being queried, we'll skip the definitions that are (a) coming from a stub file (b) external lookup and (c) check the re-export flag on the definition

This implementation does not yet support __all__ and * imports as both are features that needs to be implemented independently.

closes: #14099
closes: #15476

Test Plan

Add test cases, update existing ones if required.

[codspeed-hq]

CodSpeed Performance Report

Merging #16073 will degrade performances by 4.21%

_{Comparing dhruv/re-export-2 (929cf75) with main (3d0a58e)}

Summary

❌ 1 (👁 1) regressions
✅ 31 untouched benchmarks

Benchmarks breakdown

	Benchmark	BASE	HEAD	Change
👁	red_knot_check_file[incremental]	4.5 ms	4.7 ms	-4.21%

[MichaReiser]

Hmm, I must have overlooked this when this query was introduced. Salsa queries that have more than one argument (other than db) aren't very performant because salsa has to intern all arguments. It's the same as

#[salsa::interned]
struct SymbolByIdArgs<'db> {
	lookup: SymbolLookup,
	scope: ScopeId<'db>,
	is_dunder_slots: bool,
	symbol_id: ScopedSymbolId,
}

#[salsa::tracked]
fn symbol_by_id(db: &dyn Db, args: SymboLByIdArgs) { ... }

I wonder if this is the reason for the perf regression.

We could try if splitting the query into a internal_symbol_by_id and external_symbol_by_id query reduces the perf regression (and possibly external_dunder_slots). But we can wait with this until the PR is close to finish.

I don't there's a way to get to a single argumet function because we don't have a SymbolId salsa ingredient

[dhruvmanila]

Yeah, that could be. Thanks for the suggestions, I can try that out.

[carljm]

It does seem plausible that adding this argument increases incremental re-validation cost, since now there will be two cached memos for this query instead of one, for every global symbol that is both imported and referenced inside its file.

If this is the case, I suspect that splitting into multiple queries may not help? I would expect it to result in exactly the same increase in the total number of memos that need validating, those memos will just be split across two queries instead of being all instances of the same query.

We could reconsider whether this needs to be a query at all? I think we made it one because empirically it was a perf win to do so (despite the multiple-arguments issue), but we could check whether that is still the case after this PR. (I don't think we need it to be a query for isolation reasons, because the cross-module calls to this should all already have gone through an infer_definition_types query on the import statement.)

Another thing that could help is to check whether the target of an import is a stub file much earlier (before calling this query), and modify the semantics of the enum so that it's not "internal" vs "external" but rather "ImportFromStub" vs "Normal". That way we would only double the number of cached values for global symbols in stub files, not for global symbols in all modules. (The effect of this might be understated by our tomllib benchmark, since it's small and probably over-indexes on typeshed, relative to a checking a larger project.)

[dhruvmanila]

If this is the case, I suspect that splitting into multiple queries may not help? I would expect it to result in exactly the same increase in the total number of memos that need validating, those memos will just be split across two queries instead of being all instances of the same query.

For reference, I tried it in dafc135 (#16073) and it didn't improve the regression and I suspect it's due to the same reasons that you've outlined.

[MichaReiser]

I don't think we need it to be a query for isolation reasons, because the cross-module calls to this should all already have gone through an infer_definition_types query on the import statement.)

It does have to be a query because evaluating visibility constraints requires accessing the AST nodes but we could consider moving the query elsewhere. E.g. make evaluating the visibility constraints a query (and only if there are any non trivial constraints)

[dhruvmanila]

Performance notes (on-going)

Tried removing the lookup: SymbolLookup field and instead have two salsa query but that doesn’t improve the performance (at least locally) (e8ca883 - going to check if Codspeed yields the same results).

Patch:

diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
index 390675f32..a4c68ea63 100644
--- a/crates/red_knot_python_semantic/src/types.rs
+++ b/crates/red_knot_python_semantic/src/types.rs
@@ -126,6 +126,37 @@ fn symbol<'db>(
     name: &str,
 ) -> Symbol<'db> {
     #[salsa::tracked]
+    fn internal_symbol_by_id<'db>(
+        db: &'db dyn Db,
+        scope: ScopeId<'db>,
+        is_dunder_slots: bool,
+        symbol_id: ScopedSymbolId,
+    ) -> Symbol<'db> {
+        symbol_by_id(
+            db,
+            SymbolLookup::Internal,
+            scope,
+            is_dunder_slots,
+            symbol_id,
+        )
+    }
+
+    #[salsa::tracked]
+    fn external_symbol_by_id<'db>(
+        db: &'db dyn Db,
+        scope: ScopeId<'db>,
+        is_dunder_slots: bool,
+        symbol_id: ScopedSymbolId,
+    ) -> Symbol<'db> {
+        symbol_by_id(
+            db,
+            SymbolLookup::External,
+            scope,
+            is_dunder_slots,
+            symbol_id,
+        )
+    }
+
     fn symbol_by_id<'db>(
         db: &'db dyn Db,
         lookup: SymbolLookup,
@@ -233,7 +264,10 @@ fn symbol<'db>(
     let is_dunder_slots = name == "__slots__";
     table
         .symbol_id_by_name(name)
-        .map(|symbol| symbol_by_id(db, lookup, scope, is_dunder_slots, symbol))
+        .map(|symbol| match lookup {
+            SymbolLookup::Internal => internal_symbol_by_id(db, scope, is_dunder_slots, symbol),
+            SymbolLookup::External => external_symbol_by_id(db, scope, is_dunder_slots, symbol),
+        })
         .unwrap_or(Symbol::Unbound)
 }
 

---

I compared the trace logs between main and re-export-2 using the debug build on tomllib with the following command (includes the now merged coarse-grained commit):

RED_KNOT_PROFILE_LOG=1 RAYON_NUM_THREADS=1 red_knot check -vvv

Cold: https://www.diffchecker.com/JqHRR99N/

• There are more calls to symbol("object") compared to before
• The number of ingredients are the same (as expected)
• Reduced number of symbol("__getitem__") calls
• Increased number of symbol_by_id salsa query

Incremental: https://www.diffchecker.com/KqLUCdie/

• Similar observations as in the cold benchmark
• For the incremental part, I’m seeing a lot of DidValidateMemoizedValue salsa events. As per the docs (mentioned below), is it that now the incremental benchmark is going through more number of queries as compared to on main and that a lot of time is spent on validating whether salsa can re-use the cached value or not.

    /// Occurs when we found that all inputs to a memoized value are
    /// up-to-date and hence the value can be re-used without
    /// executing the closure.
    ///
    /// Executes before the "re-used" value is returned.

---

NOTE: There's still a 4% regression on the incremental benchmark which I want to look into.

[MichaReiser]

Thanks for the analysis. Calling more salsa queries is probably the culprit because we pay an overhead for every salsa that requires "checking" if it's still green during the incremental benchmark.

[dhruvmanila]

There's still a 4% regression which I want to look into but I'd find it useful to get some initial feedback on the implementation itself which is why I'm marking this as ready for review.

[carljm]

There's still a 4% regression

Bummer, I thought there was a reasonable chance this approach would eliminate the regression. I guess it is coming from somewhere that both approaches have in common...

[carljm]

Nice, I think this approach looks really good.

[carljm]

It does seem plausible that adding this argument increases incremental re-validation cost, since now there will be two cached memos for this query instead of one, for every global symbol that is both imported and referenced inside its file.

If this is the case, I suspect that splitting into multiple queries may not help? I would expect it to result in exactly the same increase in the total number of memos that need validating, those memos will just be split across two queries instead of being all instances of the same query.

We could reconsider whether this needs to be a query at all? I think we made it one because empirically it was a perf win to do so (despite the multiple-arguments issue), but we could check whether that is still the case after this PR. (I don't think we need it to be a query for isolation reasons, because the cross-module calls to this should all already have gone through an infer_definition_types query on the import statement.)

Another thing that could help is to check whether the target of an import is a stub file much earlier (before calling this query), and modify the semantics of the enum so that it's not "internal" vs "external" but rather "ImportFromStub" vs "Normal". That way we would only double the number of cached values for global symbols in stub files, not for global symbols in all modules. (The effect of this might be understated by our tomllib benchmark, since it's small and probably over-indexes on typeshed, relative to a checking a larger project.)

[AlexWaygood]

🥳 Great to see this get over the line — thanks @dhruvmanila! 😃