Skip to content

Commit

Permalink
Remove unpack field from SemanticIndexBuilder (#14101)
Browse files Browse the repository at this point in the history 
## Summary

Related to
#13979 (comment),
this PR removes the `current_unpack` state field from
`SemanticIndexBuilder` and passes the `Unpack` ingredient via the
`CurrentAssignment` -> `DefinitionNodeRef` conversion to finally store
it on `DefintionNodeKind`.

This involves updating the lifetime of `AnyParameterRef` (parameter to
`declare_parameter`) to use the `'db` lifetime. Currently, all AST nodes
stored on various enums are marked with `'a` lifetime but they're always
utilized using the `'db` lifetime.

This also removes the dedicated `'a` lifetime parameter on
`add_definition` which is currently being used in `DefinitionNodeRef`.
As mentioned, all AST nodes live through the `'db` lifetime so we can
remove the `'a` lifetime parameter from that method and use the `'db`
lifetime instead.
  • Loading branch information
@dhruvmanila
dhruvmanila authored Nov 6, 2024
1 parent eead549 commit 34b6a9b
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 47 deletions.
66 changes: 33 additions & 33 deletions crates/red_knot_python_semantic/src/semantic_index/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,12 +48,6 @@ pub(super) struct SemanticIndexBuilder<'db> {
/// The match case we're currently visiting.
current_match_case: Option<CurrentMatchCase<'db>>,

/// The [`Unpack`] ingredient for the current definition that belongs to an unpacking
/// assignment. This is used to correctly map multiple definitions to the *same* unpacking.
/// For example, in `a, b = 1, 2`, both `a` and `b` creates separate definitions but they both
/// belong to the same unpacking.
current_unpack: Option<Unpack<'db>>,

/// Flow states at each `break` in the current loop.
loop_break_states: Vec<FlowSnapshot>,
/// Per-scope contexts regarding nested `try`/`except` statements
Expand Down Expand Up @@ -83,7 +77,6 @@ impl<'db> SemanticIndexBuilder<'db> {
scope_stack: Vec::new(),
current_assignments: vec![],
current_match_case: None,
current_unpack: None,
loop_break_states: vec![],
try_node_context_stack_manager: TryNodeContextStackManager::default(),

Expand Down Expand Up @@ -206,15 +199,15 @@ impl<'db> SemanticIndexBuilder<'db> {
self.current_symbol_table().mark_symbol_used(id);
}

fn add_definition<'a>(
fn add_definition(
&mut self,
symbol: ScopedSymbolId,
definition_node: impl Into<DefinitionNodeRef<'a>>,
definition_node: impl Into<DefinitionNodeRef<'db>>,
) -> Definition<'db> {
let definition_node: DefinitionNodeRef<'_> = definition_node.into();
#[allow(unsafe_code)]
// SAFETY: `definition_node` is guaranteed to be a child of `self.module`
let kind = unsafe { definition_node.into_owned(self.module.clone(), self.current_unpack) };
let kind = unsafe { definition_node.into_owned(self.module.clone()) };
let category = kind.category();
let definition = Definition::new(
self.db,
Expand Down Expand Up @@ -448,7 +441,7 @@ impl<'db> SemanticIndexBuilder<'db> {
self.pop_scope();
}

fn declare_parameter(&mut self, parameter: AnyParameterRef) {
fn declare_parameter(&mut self, parameter: AnyParameterRef<'db>) {
let symbol = self.add_symbol(parameter.name().id().clone());

let definition = self.add_definition(symbol, parameter);
Expand Down Expand Up @@ -630,35 +623,38 @@ where
// We only handle assignments to names and unpackings here, other targets like
// attribute and subscript are handled separately as they don't create a new
// definition.
let is_assignment_target = match target {
let current_assignment = match target {
ast::Expr::List(_) | ast::Expr::Tuple(_) => {
self.current_unpack = Some(Unpack::new(
self.db,
self.file,
self.current_scope(),
#[allow(unsafe_code)]
unsafe {
AstNodeRef::new(self.module.clone(), target)
},
value,
countme::Count::default(),
));
true
Some(CurrentAssignment::Assign {
node,
unpack: Some(Unpack::new(
self.db,
self.file,
self.current_scope(),
#[allow(unsafe_code)]
unsafe {
AstNodeRef::new(self.module.clone(), target)
},
value,
countme::Count::default(),
)),
})
}
ast::Expr::Name(_) => true,
_ => false,
ast::Expr::Name(_) => {
Some(CurrentAssignment::Assign { node, unpack: None })
}
_ => None,
};

if is_assignment_target {
self.push_assignment(CurrentAssignment::Assign(node));
if let Some(current_assignment) = current_assignment {
self.push_assignment(current_assignment);
}

self.visit_expr(target);

if is_assignment_target {
if current_assignment.is_some() {
// Only need to pop in the case where we pushed something
self.pop_assignment();
self.current_unpack = None;
}
}
}
Expand Down Expand Up @@ -992,11 +988,12 @@ where

if is_definition {
match self.current_assignment().copied() {
Some(CurrentAssignment::Assign(assign)) => {
Some(CurrentAssignment::Assign { node, unpack }) => {
self.add_definition(
symbol,
AssignmentDefinitionNodeRef {
value: &assign.value,
unpack,
value: &node.value,
name: name_node,
},
);
Expand Down Expand Up @@ -1246,7 +1243,10 @@ where

#[derive(Copy, Clone, Debug, PartialEq)]
enum CurrentAssignment<'a> {
Assign(&'a ast::StmtAssign),
Assign {
node: &'a ast::StmtAssign,
unpack: Option<Unpack<'a>>,
},
AnnAssign(&'a ast::StmtAnnAssign),
AugAssign(&'a ast::StmtAugAssign),
For(&'a ast::StmtFor),
Expand Down
31 changes: 17 additions & 14 deletions crates/red_knot_python_semantic/src/semantic_index/definition.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,7 @@ pub(crate) struct ImportFromDefinitionNodeRef<'a> {

#[derive(Copy, Clone, Debug)]
pub(crate) struct AssignmentDefinitionNodeRef<'a> {
pub(crate) unpack: Option<Unpack<'a>>,
pub(crate) value: &'a ast::Expr,
pub(crate) name: &'a ast::ExprName,
}
Expand Down Expand Up @@ -223,13 +224,9 @@ pub(crate) struct MatchPatternDefinitionNodeRef<'a> {
pub(crate) index: u32,
}

impl DefinitionNodeRef<'_> {
impl<'db> DefinitionNodeRef<'db> {
#[allow(unsafe_code)]
pub(super) unsafe fn into_owned(
self,
parsed: ParsedModule,
unpack: Option<Unpack<'_>>,
) -> DefinitionKind<'_> {
pub(super) unsafe fn into_owned(self, parsed: ParsedModule) -> DefinitionKind<'db> {
match self {
DefinitionNodeRef::Import(alias) => {
DefinitionKind::Import(AstNodeRef::new(parsed, alias))
Expand All @@ -249,13 +246,15 @@ impl DefinitionNodeRef<'_> {
DefinitionNodeRef::NamedExpression(named) => {
DefinitionKind::NamedExpression(AstNodeRef::new(parsed, named))
}
DefinitionNodeRef::Assignment(AssignmentDefinitionNodeRef { value, name }) => {
DefinitionKind::Assignment(AssignmentDefinitionKind {
target: TargetKind::from(unpack),
value: AstNodeRef::new(parsed.clone(), value),
name: AstNodeRef::new(parsed, name),
})
}
DefinitionNodeRef::Assignment(AssignmentDefinitionNodeRef {
unpack,
value,
name,
}) => DefinitionKind::Assignment(AssignmentDefinitionKind {
target: TargetKind::from(unpack),
value: AstNodeRef::new(parsed.clone(), value),
name: AstNodeRef::new(parsed, name),
}),
DefinitionNodeRef::AnnotatedAssignment(assign) => {
DefinitionKind::AnnotatedAssignment(AstNodeRef::new(parsed, assign))
}
Expand Down Expand Up @@ -327,7 +326,11 @@ impl DefinitionNodeRef<'_> {
Self::Function(node) => node.into(),
Self::Class(node) => node.into(),
Self::NamedExpression(node) => node.into(),
Self::Assignment(AssignmentDefinitionNodeRef { value: _, name }) => name.into(),
Self::Assignment(AssignmentDefinitionNodeRef {
value: _,
unpack: _,
name,
}) => name.into(),
Self::AnnotatedAssignment(node) => node.into(),
Self::AugmentedAssignment(node) => node.into(),
Self::For(ForStmtDefinitionNodeRef {
Expand Down

0 comments on commit 34b6a9b

Please sign in to comment.