Skip to content

[red-knot] Further optimize *-import visibility constraints #17375

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

Merged
merged 3 commits into from
Apr 15, 2025
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
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
21 changes: 5 additions & 16 deletions crates/red_knot_python_semantic/src/semantic_index/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1235,32 +1235,21 @@ where
symbol_id,
referenced_module,
);
let pre_definition = self.flow_snapshot();
self.push_additional_definition(symbol_id, node_ref);

// 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.push_additional_definition(symbol_id, node_ref);
self.current_use_def_map_mut()
.negate_star_import_visibility_constraint(
.record_and_negate_star_import_visibility_constraint(
star_import,
symbol_id,
constraint_id,
);

self.flow_merge(post_definition);
} 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();
Expand Down
59 changes: 19 additions & 40 deletions crates/red_knot_python_semantic/src/semantic_index/use_def.rs
Original file line number Diff line number Diff line change
Expand Up @@ -714,32 +714,32 @@ impl<'db> UseDefMapBuilder<'db> {
.add_and_constraint(self.scope_start_visibility, constraint);
}

#[must_use = "A `*`-import visibility constraint must always be negated after it is added"]
pub(super) fn record_star_import_visibility_constraint(
pub(super) fn record_and_negate_star_import_visibility_constraint(
Copy link
Contributor

Choose a reason for hiding this comment

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

I think it would be useful to future-us to add a doc comment to this method explaining what it is doing and why it is safe to do it, in the limited cases where we use this function.

Copy link
Member Author

Choose a reason for hiding this comment

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

Added some docs! Sadly this PR is no longer net-negative in terms of whether it adds or removes lines of code, but so it goes 🥲

Would appreciate a quick review from you or @sharkdp to double-check everything in the doc-comment is accurate!

&mut self,
star_import: StarImportPlaceholderPredicate<'db>,
symbol: ScopedSymbolId,
) -> StarImportVisibilityConstraintId {
) {
let predicate_id = self.add_predicate(star_import.into());
let visibility_id = self.visibility_constraints.add_atom(predicate_id);
self.symbol_states[symbol]
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),
);
post_definition_state
.record_visibility_constraint(&mut self.visibility_constraints, visibility_id);
StarImportVisibilityConstraintId(visibility_id)
}

pub(super) fn negate_star_import_visibility_constraint(
&mut self,
symbol_id: ScopedSymbolId,
constraint: StarImportVisibilityConstraintId,
) {
let negated_constraint = self
.visibility_constraints
.add_not_constraint(constraint.into_scoped_constraint_id());
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);
self.symbol_states[symbol]
.record_visibility_constraint(&mut self.visibility_constraints, negated_visibility_id);

self.symbol_states[symbol].merge(
post_definition_state,
&mut self.narrowing_constraints,
&mut self.visibility_constraints,
);
}

/// This method resets the visibility constraints for all symbols to a previous state
Expand Down Expand Up @@ -951,24 +951,3 @@ impl<'db> UseDefMapBuilder<'db> {
}
}
}

/// Newtype wrapper over [`ScopedVisibilityConstraintId`] to improve type safety.
///
/// By returning this type from [`UseDefMapBuilder::record_star_import_visibility_constraint`]
/// rather than [`ScopedVisibilityConstraintId`] directly, we ensure that
/// [`UseDefMapBuilder::negate_star_import_visibility_constraint`] must be called after the
/// visibility constraint has been added, and we ensure that
/// [`super::SemanticIndexBuilder::record_negated_visibility_constraint`] *cannot* be called with
/// the narrowing constraint (which would lead to incorrect behaviour).
///
/// This type is defined here rather than in the [`super::visibility_constraints`] module
/// because it should only ever be constructed and deconstructed from methods in the
/// [`UseDefMapBuilder`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct StarImportVisibilityConstraintId(ScopedVisibilityConstraintId);

impl StarImportVisibilityConstraintId {
fn into_scoped_constraint_id(self) -> ScopedVisibilityConstraintId {
self.0
}
}
Loading