Cached inference of all definitions in an unpacking (#13979)

## Summary

This PR adds a new salsa query and an ingredient to resolve all the
variables involved in an unpacking assignment like `(a, b) = (1, 2)` at
once. Previously, we'd recursively try to match the correct type for
each definition individually which will result in creating duplicate
diagnostics.

This PR still doesn't solve the duplicate diagnostics issue because that
requires a different solution like using salsa accumulator or
de-duplicating the diagnostics manually.

Related: #13773 

## Test Plan

Make sure that all unpack assignment test cases pass, there are no
panics in the corpus tests.

## Todo

- [x] Look at the performance regression
This commit is contained in:
Dhruv Manilawala 2024-11-04 17:11:57 +05:30 committed by GitHub
parent 012f385f5d
commit e302c2de7c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 314 additions and 202 deletions

View file

@ -25,12 +25,13 @@ use crate::semantic_index::symbol::{
};
use crate::semantic_index::use_def::{FlowSnapshot, UseDefMapBuilder};
use crate::semantic_index::SemanticIndex;
use crate::unpack::Unpack;
use crate::Db;
use super::constraint::{Constraint, ConstraintNode, PatternConstraint};
use super::definition::{
AssignmentKind, DefinitionCategory, ExceptHandlerDefinitionNodeRef,
MatchPatternDefinitionNodeRef, WithItemDefinitionNodeRef,
DefinitionCategory, ExceptHandlerDefinitionNodeRef, MatchPatternDefinitionNodeRef,
WithItemDefinitionNodeRef,
};
mod except_handlers;
@ -46,6 +47,13 @@ pub(super) struct SemanticIndexBuilder<'db> {
current_assignments: Vec<CurrentAssignment<'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
@ -75,6 +83,7 @@ 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(),
@ -211,7 +220,7 @@ impl<'db> SemanticIndexBuilder<'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()) };
let kind = unsafe { definition_node.into_owned(self.module.clone(), self.current_unpack) };
let category = kind.category();
let definition = Definition::new(
self.db,
@ -619,25 +628,43 @@ where
}
ast::Stmt::Assign(node) => {
debug_assert_eq!(&self.current_assignments, &[]);
self.visit_expr(&node.value);
self.add_standalone_expression(&node.value);
for (target_index, target) in node.targets.iter().enumerate() {
let kind = match target {
ast::Expr::List(_) | ast::Expr::Tuple(_) => Some(AssignmentKind::Sequence),
ast::Expr::Name(_) => Some(AssignmentKind::Name),
_ => None,
let value = self.add_standalone_expression(&node.value);
for target in &node.targets {
// 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 {
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
}
ast::Expr::Name(_) => true,
_ => false,
};
if let Some(kind) = kind {
self.push_assignment(CurrentAssignment::Assign {
assignment: node,
target_index,
kind,
});
if is_assignment_target {
self.push_assignment(CurrentAssignment::Assign(node));
}
self.visit_expr(target);
if kind.is_some() {
// only need to pop in the case where we pushed something
if is_assignment_target {
// Only need to pop in the case where we pushed something
self.pop_assignment();
self.current_unpack = None;
}
}
}
@ -971,18 +998,12 @@ where
if is_definition {
match self.current_assignment().copied() {
Some(CurrentAssignment::Assign {
assignment,
target_index,
kind,
}) => {
Some(CurrentAssignment::Assign(assign)) => {
self.add_definition(
symbol,
AssignmentDefinitionNodeRef {
assignment,
target_index,
value: &assign.value,
name: name_node,
kind,
},
);
}
@ -1228,11 +1249,7 @@ where
#[derive(Copy, Clone, Debug, PartialEq)]
enum CurrentAssignment<'a> {
Assign {
assignment: &'a ast::StmtAssign,
target_index: usize,
kind: AssignmentKind,
},
Assign(&'a ast::StmtAssign),
AnnAssign(&'a ast::StmtAnnAssign),
AugAssign(&'a ast::StmtAugAssign),
For(&'a ast::StmtFor),