Skip to content

[red-knot] Simplify visibility constraint handling for *-import definitions #17486

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 1 commit into from
Apr 21, 2025
Merged

[red-knot] Simplify visibility constraint handling for *-import definitions #17486

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
68 changes: 27 additions & 41 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 @@ -355,15 +355,14 @@ impl<'db> SemanticIndexBuilder<'db> {
self.current_use_def_map_mut().merge(state);
}

/// 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) {
/// Add a symbol to the symbol table and the use-def map.
/// Return the [`ScopedSymbolId`] that uniquely identifies the symbol in both.
fn add_symbol(&mut self, name: Name) -> ScopedSymbolId {
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, added)
symbol_id
}

fn add_attribute(&mut self, name: Name) -> ScopedSymbolId {
Expand Down Expand Up @@ -797,7 +796,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 @@ -895,20 +894,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 @@ -1139,7 +1138,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());

// Record a use of the function name in the scope that it is defined in, so that it
// can be used to find previously defined functions with the same name. This is
Expand Down Expand Up @@ -1174,11 +1173,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 @@ -1215,7 +1214,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 @@ -1286,7 +1285,7 @@ where
//
// For more details, see the doc-comment on `StarImportPlaceholderPredicate`.
for export in exported_names(self.db, referenced_module) {
let (symbol_id, newly_added) = self.add_symbol(export.clone());
let symbol_id = self.add_symbol(export.clone());
let node_ref = StarImportDefinitionNodeRef { node, symbol_id };
let star_import = StarImportPlaceholderPredicate::new(
self.db,
Expand All @@ -1295,28 +1294,15 @@ where
referenced_module,
);

// 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 {
self.push_additional_definition(symbol_id, node_ref);
self.current_use_def_map_mut()
.record_and_negate_star_import_visibility_constraint(
star_import,
symbol_id,
);
} else {
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);
}
let pre_definition =
self.current_use_def_map().single_symbol_snapshot(symbol_id);
self.push_additional_definition(symbol_id, node_ref);
self.current_use_def_map_mut()
.record_and_negate_star_import_visibility_constraint(
star_import,
symbol_id,
pre_definition,
);
}

continue;
Expand All @@ -1336,7 +1322,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 @@ -1754,7 +1740,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 @@ -1841,7 +1827,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 @@ -2245,7 +2231,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 @@ -2266,7 +2252,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
34 changes: 21 additions & 13 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 @@ -775,7 +775,16 @@ impl<'db> UseDefMapBuilder<'db> {
.add_and_constraint(self.scope_start_visibility, constraint);
}

/// This method exists solely as a fast path for handling `*`-import visibility constraints.
/// Snapshot the state of a single symbol at the current point in control flow.
///
/// This is only used for `*`-import visibility constraints, which are handled differently
/// to most other visibility constraints. See the doc-comment for
/// [`Self::record_and_negate_star_import_visibility_constraint`] for more details.
pub(super) fn single_symbol_snapshot(&self, symbol: ScopedSymbolId) -> SymbolState {
self.symbol_states[symbol].clone()
}

/// This method exists solely for handling `*`-import visibility constraints.
///
/// The reason why we add visibility constraints for [`Definition`]s created by `*` imports
/// is laid out in the doc-comment for [`StarImportPlaceholderPredicate`]. But treating these
Expand All @@ -784,12 +793,11 @@ impl<'db> UseDefMapBuilder<'db> {
/// dominates. (Although `*` imports are not common generally, they are used in several
/// important places by typeshed.)
///
/// To solve these regressions, it was observed that we could add a fast path for `*`-import
/// definitions which added a new symbol to the global scope (as opposed to `*`-import definitions
/// that provided redefinitions for *pre-existing* global-scope symbols). The fast path does a
/// number of things differently to our normal handling of visibility constraints:
/// To solve these regressions, it was observed that we could do significantly less work for
/// `*`-import definitions. We do a number of things differently here to our normal handling of
/// visibility constraints:
///
/// - It only applies and negates the visibility constraints to a single symbol, rather than to
/// - We only apply and negate the visibility constraints to a single symbol, rather than to
/// all symbols. This is possible here because, unlike most definitions, we know in advance that
/// exactly one definition occurs inside the "if-true" predicate branch, and we know exactly
/// which definition it is.
Expand All @@ -800,9 +808,9 @@ impl<'db> UseDefMapBuilder<'db> {
/// the visibility constraints is only important for symbols that did not have any new
/// definitions inside either the "if-predicate-true" branch or the "if-predicate-false" branch.
///
/// - It avoids multiple expensive calls to [`Self::snapshot`]. This is possible because we know
/// the symbol is newly added, so we know the prior state of the symbol was
/// [`SymbolState::undefined`].
/// - We only snapshot the state for a single symbol prior to the definition, rather than doing
/// expensive calls to [`Self::snapshot`]. Again, this is possible because we know
/// that only a single definition occurs inside the "if-predicate-true" predicate branch.
///
/// - Normally we take care to check whether an "if-predicate-true" branch or an
/// "if-predicate-false" branch contains a terminal statement: these can affect the visibility
Expand All @@ -815,17 +823,17 @@ impl<'db> UseDefMapBuilder<'db> {
&mut self,
star_import: StarImportPlaceholderPredicate<'db>,
symbol: ScopedSymbolId,
pre_definition_state: SymbolState,
) {
let predicate_id = self.add_predicate(star_import.into());
let visibility_id = self.visibility_constraints.add_atom(predicate_id);
let negated_visibility_id = self
.visibility_constraints
.add_not_constraint(visibility_id);

let mut post_definition_state = std::mem::replace(
&mut self.symbol_states[symbol],
SymbolState::undefined(self.scope_start_visibility),
);
let mut post_definition_state =
std::mem::replace(&mut self.symbol_states[symbol], pre_definition_state);

post_definition_state
.record_visibility_constraint(&mut self.visibility_constraints, visibility_id);

Expand Down
2 changes: 1 addition & 1 deletion crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,7 +314,7 @@ impl SymbolBindings {
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub(super) struct SymbolState {
pub(in crate::semantic_index) struct SymbolState {
Copy link
Member

@dcreager dcreager Apr 21, 2025

Choose a reason for hiding this comment

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

🤯 We do have pub(grandparent) as a possibility, I did not know that! I still worry that we might be over-thinking visibility within this crate, but that is not at all a blocker for this PR.

declarations: SymbolDeclarations,
bindings: SymbolBindings,
}
Expand Down
Loading