Skip to content

[red-knot] Optimise visibility constraints for *-import definitions #17317

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.

Sign up for GitHub

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

Merged
merged 2 commits into from
Apr 9, 2025
Merged

[red-knot] Optimise visibility constraints for *-import definitions #17317

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5caf5b4
[red-knot] Optimise visibility constraints for `*`-import definitions
AlexWaygood Apr 9, 2025
0bc6440
improve type safety
AlexWaygood Apr 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 51 additions & 23 deletions crates/red_knot_python_semantic/src/semantic_index/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,12 +331,15 @@ impl<'db> SemanticIndexBuilder<'db> {
self.current_use_def_map_mut().merge(state);
}

fn add_symbol(&mut self, name: Name) -> ScopedSymbolId {
/// Return a 2-element tuple, where the first element is the [`ScopedSymbolId`] of the
/// symbol added, and the second element is a boolean indicating whether the symbol was *newly*
/// added or not
fn add_symbol(&mut self, name: Name) -> (ScopedSymbolId, bool) {
let (symbol_id, added) = self.current_symbol_table().add_symbol(name);
if added {
self.current_use_def_map_mut().add_symbol(symbol_id);
}
symbol_id
(symbol_id, added)
}

fn mark_symbol_bound(&mut self, id: ScopedSymbolId) {
Expand Down Expand Up @@ -747,7 +750,7 @@ impl<'db> SemanticIndexBuilder<'db> {
..
}) => (name, &None, default),
};
let symbol = self.add_symbol(name.id.clone());
let (symbol, _) = self.add_symbol(name.id.clone());
// TODO create Definition for PEP 695 typevars
// note that the "bound" on the typevar is a totally different thing than whether
// or not a name is "bound" by a typevar declaration; the latter is always true.
Expand Down Expand Up @@ -841,20 +844,20 @@ impl<'db> SemanticIndexBuilder<'db> {
self.declare_parameter(parameter);
}
if let Some(vararg) = parameters.vararg.as_ref() {
let symbol = self.add_symbol(vararg.name.id().clone());
let (symbol, _) = self.add_symbol(vararg.name.id().clone());
self.add_definition(
symbol,
DefinitionNodeRef::VariadicPositionalParameter(vararg),
);
}
if let Some(kwarg) = parameters.kwarg.as_ref() {
let symbol = self.add_symbol(kwarg.name.id().clone());
let (symbol, _) = self.add_symbol(kwarg.name.id().clone());
self.add_definition(symbol, DefinitionNodeRef::VariadicKeywordParameter(kwarg));
}
}

fn declare_parameter(&mut self, parameter: &'db ast::ParameterWithDefault) {
let symbol = self.add_symbol(parameter.name().id().clone());
let (symbol, _) = self.add_symbol(parameter.name().id().clone());

let definition = self.add_definition(symbol, parameter);

Expand Down Expand Up @@ -1071,7 +1074,7 @@ where
// The symbol for the function name itself has to be evaluated
// at the end to match the runtime evaluation of parameter defaults
// and return-type annotations.
let symbol = self.add_symbol(name.id.clone());
let (symbol, _) = self.add_symbol(name.id.clone());
self.add_definition(symbol, function_def);
}
ast::Stmt::ClassDef(class) => {
Expand All @@ -1095,11 +1098,11 @@ where
);

// In Python runtime semantics, a class is registered after its scope is evaluated.
let symbol = self.add_symbol(class.name.id.clone());
let (symbol, _) = self.add_symbol(class.name.id.clone());
self.add_definition(symbol, class);
}
ast::Stmt::TypeAlias(type_alias) => {
let symbol = self.add_symbol(
let (symbol, _) = self.add_symbol(
type_alias
.name
.as_name_expr()
Expand Down Expand Up @@ -1133,7 +1136,7 @@ where
(Name::new(alias.name.id.split('.').next().unwrap()), false)
};

let symbol = self.add_symbol(symbol_name);
let (symbol, _) = self.add_symbol(symbol_name);
self.add_definition(
symbol,
ImportDefinitionNodeRef {
Expand Down Expand Up @@ -1200,7 +1203,7 @@ where
//
// For more details, see the doc-comment on `StarImportPlaceholderPredicate`.
for export in exported_names(self.db, referenced_module) {
let symbol_id = self.add_symbol(export.clone());
let (symbol_id, newly_added) = self.add_symbol(export.clone());
let node_ref = StarImportDefinitionNodeRef { node, symbol_id };
let star_import = StarImportPlaceholderPredicate::new(
self.db,
Expand All @@ -1210,13 +1213,38 @@ where
);
let pre_definition = self.flow_snapshot();
self.push_additional_definition(symbol_id, node_ref);
let constraint_id =
self.record_visibility_constraint(star_import.into());
let post_definition = self.flow_snapshot();
self.flow_restore(pre_definition.clone());
self.record_negated_visibility_constraint(constraint_id);
self.flow_merge(post_definition);
self.simplify_visibility_constraints(pre_definition);

// Fast path for if there were no previous definitions
// of the symbol defined through the `*` import:
// we can apply the visibility constraint to *only* the added definition,
// rather than all definitions
if newly_added {
let constraint_id = self
.current_use_def_map_mut()
.record_star_import_visibility_constraint(
star_import,
symbol_id,
);

let post_definition = self.flow_snapshot();
self.flow_restore(pre_definition);

self.current_use_def_map_mut()
.negate_star_import_visibility_constraint(
symbol_id,
constraint_id,
);

self.flow_merge(post_definition);
Comment on lines +1223 to +1239
Copy link
Contributor

@carljm carljm Apr 9, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

This PR reclaims almost all the lost perf, so it seems fine to just go ahead with it as-is. But it still seems like we do more work here than necessary? In the newly-added case, we shouldn't have to do any snapshotting or flow-state merging at all. We can just add the star-import definition, apply the (positive) visibility constraint to it, and move on.

Copy link
Member Author

Choose a reason for hiding this comment

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

As in, something like this (relative to my PR)? Quite a lot of tests start failing if I apply this diff...

diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs
index 486482f7f..406014831 100644
--- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs
+++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs
@@ -1219,23 +1219,12 @@ where
                             // we can apply the visibility constraint to *only* the added definition,
                             // rather than all definitions
                             if newly_added {
-                                let constraint_id = self
+                                self
                                     .current_use_def_map_mut()
                                     .record_star_import_visibility_constraint(
                                         star_import,
                                         symbol_id,
                                     );
-
-                                let post_definition = self.flow_snapshot();
-                                self.flow_restore(pre_definition);
-
-                                self.current_use_def_map_mut()
-                                    .negate_star_import_visibility_constraint(
-                                        symbol_id,
-                                        constraint_id,
-                                    );
-
-                                self.flow_merge(post_definition);
                             } else {
                                 let constraint_id =
                                     self.record_visibility_constraint(star_import.into());
diff --git a/crates/red_knot_python_semantic/src/semantic_index/use_def.rs b/crates/red_knot_python_semantic/src/semantic_index/use_def.rs
index edb27d26e..002faf319 100644
--- a/crates/red_knot_python_semantic/src/semantic_index/use_def.rs
+++ b/crates/red_knot_python_semantic/src/semantic_index/use_def.rs
@@ -713,19 +713,6 @@ impl<'db> UseDefMapBuilder<'db> {
         visibility_id
     }
 
-    pub(super) fn negate_star_import_visibility_constraint(
-        &mut self,
-        symbol_id: ScopedSymbolId,
-        constraint: ScopedVisibilityConstraintId,
-    ) {
-        let negated_constraint = self.visibility_constraints.add_not_constraint(constraint);
-        self.symbol_states[symbol_id]
-            .record_visibility_constraint(&mut self.visibility_constraints, negated_constraint);
-        self.scope_start_visibility = self
-            .visibility_constraints
-            .add_and_constraint(self.scope_start_visibility, negated_constraint);
-    }
-
     /// This method resets the visibility constraints for all symbols to a previous state
     /// *if* there have been no new declarations or bindings since then. Consider the
     /// following example:

Copy link
Member Author

Choose a reason for hiding this comment

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

I'll go ahead and merge for now, but feel free to file a followup PR if you can see a way of optimising this further without breaking all my tests ;)

Copy link
Contributor

Choose a reason for hiding this comment

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

This is something I also didn't think about when writing that message on Discord, but I think you would also need to apply the negated visibility constraint to the implicit symbol = <unbound> binding (it is stored as a None entry at the beginning of the vector). I think that should allow you to get rid of snapshotting entirely for the "fast path".

We would basically still model

if <placeholder>:
    from module import symbol
else:
    symbol = <unbound>

but with the advantage that the constraint does not need to be applied to anything else.

Copy link
Contributor

Choose a reason for hiding this comment

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

Ah yeah, that's because "unbound" is a binding also, and we do need to apply the negated visibility constraint to the "unbound" binding. (The "unbound" binding is visible only if the star import target symbol does not end up existing.) I think this should still be doable by just applying constraints to the right bindings, without snapshotting and merging, but it would require even more special-cased new APIs. Not sure it's worth it since this version already does so well, definitely no need to look into it now.

Copy link
Member Author

Choose a reason for hiding this comment

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

Okay -- I'm going to move on for now as I think both of you have a far better understanding of our visibility constraints machinery than I do. But I welcome further optimisations of this code from anybody who wants to work on them 😄

Copy link
Contributor

Choose a reason for hiding this comment

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

Not sure it's worth it since this version already does so well, definitely no need to look into it now.

I agree. Seems like the new isolated application of constraints helps enough to simplify the overall constraint structure.

} else {
let constraint_id =
self.record_visibility_constraint(star_import.into());
let post_definition = self.flow_snapshot();
self.flow_restore(pre_definition.clone());
self.record_negated_visibility_constraint(constraint_id);
self.flow_merge(post_definition);
self.simplify_visibility_constraints(pre_definition);
}
}

continue;
Expand All @@ -1236,7 +1264,7 @@ where
self.has_future_annotations |= alias.name.id == "annotations"
&& node.module.as_deref() == Some("__future__");

let symbol = self.add_symbol(symbol_name.clone());
let (symbol, _) = self.add_symbol(symbol_name.clone());

self.add_definition(
symbol,
Expand Down Expand Up @@ -1636,7 +1664,7 @@ where
// which is invalid syntax. However, it's still pretty obvious here that the user
// *wanted* `e` to be bound, so we should still create a definition here nonetheless.
if let Some(symbol_name) = symbol_name {
let symbol = self.add_symbol(symbol_name.id.clone());
let (symbol, _) = self.add_symbol(symbol_name.id.clone());

self.add_definition(
symbol,
Expand Down Expand Up @@ -1721,7 +1749,7 @@ where
(ast::ExprContext::Del, _) => (false, true),
(ast::ExprContext::Invalid, _) => (false, false),
};
let symbol = self.add_symbol(id.clone());
let (symbol, _) = self.add_symbol(id.clone());

if is_use {
self.mark_symbol_used(symbol);
Expand Down Expand Up @@ -2007,7 +2035,7 @@ where
range: _,
}) = pattern
{
let symbol = self.add_symbol(name.id().clone());
let (symbol, _) = self.add_symbol(name.id().clone());
let state = self.current_match_case.as_ref().unwrap();
self.add_definition(
symbol,
Expand All @@ -2028,7 +2056,7 @@ where
rest: Some(name), ..
}) = pattern
{
let symbol = self.add_symbol(name.id().clone());
let (symbol, _) = self.add_symbol(name.id().clone());
let state = self.current_match_case.as_ref().unwrap();
self.add_definition(
symbol,
Expand Down
29 changes: 27 additions & 2 deletions crates/red_knot_python_semantic/src/semantic_index/use_def.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,7 @@ use crate::semantic_index::narrowing_constraints::{
NarrowingConstraints, NarrowingConstraintsBuilder, NarrowingConstraintsIterator,
};
use crate::semantic_index::predicate::{
Predicate, Predicates, PredicatesBuilder, ScopedPredicateId,
Predicate, Predicates, PredicatesBuilder, ScopedPredicateId, StarImportPlaceholderPredicate,
};
use crate::semantic_index::symbol::{FileScopeId, ScopedSymbolId};
use crate::semantic_index::visibility_constraints::{
Expand Down Expand Up @@ -603,7 +603,7 @@ pub(super) struct UseDefMapBuilder<'db> {
/// x # we store a reachability constraint of [test] for this use of `x`
///
/// y = 2
///
///
/// # we record a visibility constraint of [test] here, which retroactively affects
/// # the `y = 1` and the `y = 2` binding.
/// else:
Expand Down Expand Up @@ -701,6 +701,31 @@ impl<'db> UseDefMapBuilder<'db> {
.add_and_constraint(self.scope_start_visibility, constraint);
}

pub(super) fn record_star_import_visibility_constraint(
&mut self,
star_import: StarImportPlaceholderPredicate<'db>,
symbol: ScopedSymbolId,
) -> ScopedVisibilityConstraintId {
let predicate_id = self.add_predicate(star_import.into());
let visibility_id = self.visibility_constraints.add_atom(predicate_id);
self.symbol_states[symbol]
.record_visibility_constraint(&mut self.visibility_constraints, visibility_id);
visibility_id
}

pub(super) fn negate_star_import_visibility_constraint(
&mut self,
symbol_id: ScopedSymbolId,
constraint: ScopedVisibilityConstraintId,
) {
let negated_constraint = self.visibility_constraints.add_not_constraint(constraint);
self.symbol_states[symbol_id]
.record_visibility_constraint(&mut self.visibility_constraints, negated_constraint);
self.scope_start_visibility = self
.visibility_constraints
.add_and_constraint(self.scope_start_visibility, negated_constraint);
}

/// This method resets the visibility constraints for all symbols to a previous state
/// *if* there have been no new declarations or bindings since then. Consider the
/// following example:
Expand Down
Loading