Skip to content

Commit ec07736

Browse files
ibraheemdevibraheemdev
authored
Merge 9b59806 into 61ab8ef 
2 parents 61ab8ef + 9b59806 commit ec07736

23 files changed

Lines changed: 977 additions & 177 deletions

File tree

‎crates/ruff_benchmark/benches/ty_walltime.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -202,7 +202,7 @@ static SYMPY: Benchmark = Benchmark::new(
202202
max_dep_date: "2025-06-17",
203203
python_version: SupportedPythonVersion::Py312,
204204
},
205-
14150,
205+
14250,
206206
);
207207

208208
static TANJUN: Benchmark = Benchmark::new(

‎crates/ruff_db/src/parsed.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ impl ParsedModule {
106106
)))),
107107
}
108108
}
109+
109110
/// Loads a reference to the parsed module.
110111
///
111112
/// Note that holding on to the reference will prevent garbage collection

‎crates/ty_python_core/src/ast_ids.rs‎

Lines changed: 32 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -32,67 +32,80 @@ pub(crate) struct AstIds {
3232
}
3333

3434
impl AstIds {
35-
fn use_id(&self, key: impl Into<ExpressionNodeKey>) -> ScopedUseId {
36-
self.uses_map[&key.into()]
35+
fn try_use_id(&self, key: impl Into<ExpressionNodeKey>) -> Option<ScopedUseId> {
36+
self.uses_map.get(&key.into()).copied()
3737
}
3838
}
3939

40-
fn ast_ids<'db>(db: &'db dyn Db, scope: ScopeId) -> &'db AstIds {
41-
semantic_index(db, scope.file(db)).ast_ids(scope.file_scope_id(db))
42-
}
43-
4440
/// Uniquely identifies a use of a name in a [`crate::FileScopeId`].
4541
#[newtype_index]
4642
#[derive(get_size2::GetSize)]
4743
pub struct ScopedUseId;
4844

4945
pub trait HasScopedUseId {
5046
/// Returns the ID that uniquely identifies the use in `scope`.
51-
fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId;
47+
fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
48+
self.try_scoped_use_id(db, scope).expect("expected use")
49+
}
50+
51+
/// Returns the ID that uniquely identifies the use in `scope`, if the current
52+
/// AST node represents a use.
53+
fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option<ScopedUseId>;
5254
}
5355

5456
impl HasScopedUseId for ast::Identifier {
55-
fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
57+
fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option<ScopedUseId> {
5658
let ast_ids = ast_ids(db, scope);
57-
ast_ids.use_id(self)
59+
ast_ids.try_use_id(self)
5860
}
5961
}
6062

6163
impl HasScopedUseId for ast::ExprName {
62-
fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
64+
fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option<ScopedUseId> {
6365
let expression_ref = ExprRef::from(self);
64-
expression_ref.scoped_use_id(db, scope)
66+
expression_ref.try_scoped_use_id(db, scope)
6567
}
6668
}
6769

6870
impl HasScopedUseId for ast::ExprAttribute {
69-
fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
71+
fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option<ScopedUseId> {
7072
let expression_ref = ExprRef::from(self);
71-
expression_ref.scoped_use_id(db, scope)
73+
expression_ref.try_scoped_use_id(db, scope)
7274
}
7375
}
7476

7577
impl HasScopedUseId for ast::ExprSubscript {
76-
fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
78+
fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option<ScopedUseId> {
7779
let expression_ref = ExprRef::from(self);
78-
expression_ref.scoped_use_id(db, scope)
80+
expression_ref.try_scoped_use_id(db, scope)
7981
}
8082
}
8183

8284
impl HasScopedUseId for ast::Keyword {
83-
fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
85+
fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option<ScopedUseId> {
86+
let ast_ids = ast_ids(db, scope);
87+
ast_ids.try_use_id(self)
88+
}
89+
}
90+
91+
impl HasScopedUseId for ast::Expr {
92+
fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option<ScopedUseId> {
8493
let ast_ids = ast_ids(db, scope);
85-
ast_ids.use_id(self)
94+
ast_ids.try_use_id(self)
8695
}
8796
}
8897

8998
impl HasScopedUseId for ast::ExprRef<'_> {
90-
fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
99+
fn try_scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> Option<ScopedUseId> {
91100
let ast_ids = ast_ids(db, scope);
92-
ast_ids.use_id(*self)
101+
ast_ids.try_use_id(*self)
93102
}
94103
}
95104

105+
fn ast_ids<'db>(db: &'db dyn Db, scope: ScopeId) -> &'db AstIds {
106+
semantic_index(db, scope.file(db)).ast_ids(scope.file_scope_id(db))
107+
}
108+
96109
#[derive(Debug, Default)]
97110
pub(super) struct AstIdsBuilder {
98111
uses_map: FxHashMap<ExpressionNodeKey, ScopedUseId>,

‎crates/ty_python_core/src/ast_node_ref.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ pub struct AstNodeRef<T> {
5252
}
5353

5454
impl<T> AstNodeRef<T> {
55-
pub(crate) fn index(&self) -> NodeIndex {
55+
pub fn index(&self) -> NodeIndex {
5656
self.index
5757
}
5858
}

‎crates/ty_python_core/src/builder.rs‎

Lines changed: 84 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@ use ruff_text_size::{Ranged, TextRange};
2121
use ty_module_resolver::{ModuleName, resolve_module};
2222

2323
use crate::HasTrackedScope;
24-
use crate::ast_ids::AstIdsBuilder;
2524
use crate::ast_ids::node_key::ExpressionNodeKey;
25+
use crate::ast_ids::{AstIdsBuilder, ScopedUseId};
2626
use crate::ast_node_ref::AstNodeRef;
2727
use crate::definition::{
2828
AnnotatedAssignmentDefinitionNodeRef, AssignmentDefinitionNodeRef,
@@ -103,7 +103,7 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
103103
current_assignments: Vec<CurrentAssignment<'ast, 'db>>,
104104
/// The statements we're currently visiting, with
105105
/// the most recent visit at the end of the Vec.
106-
current_statements: Vec<CurrentStatement<'ast>>,
106+
current_statements: Vec<CurrentStatement<'ast, 'db>>,
107107
/// The match case we're currently visiting.
108108
current_match_case: Option<CurrentMatchCase<'ast>>,
109109
/// The name of the first function parameter of the innermost function that we're currently visiting.
@@ -134,9 +134,13 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
134134
definitions_by_node: FxHashMap<DefinitionNodeKey, Definitions<'db>>,
135135
expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
136136
statements_by_node: FxHashMap<StatementNodeKey, Statement<'db>>,
137-
enclosing_lambda_statements: FxHashMap<ExpressionNodeKey, Statement<'db>>,
138137
imported_modules: FxHashSet<ModuleName>,
139138
seen_submodule_imports: FxHashSet<String>,
139+
// A map from a lambda expression to its enclosing statement.
140+
enclosing_lambda_statements: FxHashMap<ExpressionNodeKey, Statement<'db>>,
141+
// A map from a collection literal definition to a statement containing the use of that
142+
// collection.
143+
uses_by_collection: FxHashMap<Definition<'db>, Vec<(Statement<'db>, ExpressionNodeKey)>>,
140144
/// Hashset of all [`FileScopeId`]s that correspond to [generator functions].
141145
///
142146
/// [generator functions]: https://docs.python.org/3/glossary.html#term-generator
@@ -176,6 +180,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
176180
expressions_by_node: FxHashMap::default(),
177181
statements_by_node: FxHashMap::default(),
178182
enclosing_lambda_statements: FxHashMap::default(),
183+
uses_by_collection: FxHashMap::default(),
179184

180185
seen_submodule_imports: FxHashSet::default(),
181186
imported_modules: FxHashSet::default(),
@@ -752,12 +757,13 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
752757
self.current_place_table_mut().symbol_mut(id).mark_used();
753758
}
754759

755-
fn record_place_use(&mut self, place_id: ScopedPlaceId, expr: &'ast ast::Expr) {
760+
fn record_place_use(&mut self, place_id: ScopedPlaceId, expr: &'ast ast::Expr) -> ScopedUseId {
756761
if let ScopedPlaceId::Symbol(symbol_id) = place_id {
757762
self.mark_symbol_used(symbol_id);
758763
}
759764
let use_id = self.current_ast_ids().record_use(expr);
760765
self.current_use_def_map_mut().record_use(place_id, use_id);
766+
use_id
761767
}
762768

763769
fn record_place_definition(&mut self, place_id: ScopedPlaceId, expr: &'ast ast::Expr) {
@@ -1336,15 +1342,15 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
13361342
self.current_assignments.last_mut()
13371343
}
13381344

1339-
fn push_statement(&mut self, statement: CurrentStatement<'ast>) {
1345+
fn push_statement(&mut self, statement: CurrentStatement<'ast, 'db>) {
13401346
self.current_statements.push(statement);
13411347
}
13421348

1343-
fn pop_statement(&mut self) -> CurrentStatement<'ast> {
1349+
fn pop_statement(&mut self) -> CurrentStatement<'ast, 'db> {
13441350
self.current_statements.pop().unwrap()
13451351
}
13461352

1347-
fn current_statement_mut(&mut self) -> Option<&mut CurrentStatement<'ast>> {
1353+
fn current_statement_mut(&mut self) -> Option<&mut CurrentStatement<'ast, 'db>> {
13481354
self.current_statements.last_mut()
13491355
}
13501356

@@ -1518,11 +1524,9 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
15181524
ast::Stmt::ClassDef(class) => {
15191525
Some(Statement::Definition(self.expect_single_definition(class)))
15201526
}
1521-
ast::Stmt::Expr(expr) => self
1522-
.expressions_by_node
1523-
.get(&(&expr.value).into())
1524-
.copied()
1525-
.map(Statement::Expression),
1527+
ast::Stmt::Expr(ast::StmtExpr { value, .. }) => {
1528+
Some(Statement::Expression(self.add_standalone_expression(value)))
1529+
}
15261530
ast::Stmt::Assign(assign) => {
15271531
if let [ast::Expr::Name(name)] = &assign.targets[..] {
15281532
Some(Statement::Definition(self.expect_single_definition(name)))
@@ -1917,6 +1921,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
19171921
self.definitions_by_node.shrink_to_fit();
19181922
self.statements_by_node.shrink_to_fit();
19191923
self.enclosing_lambda_statements.shrink_to_fit();
1924+
self.uses_by_collection.shrink_to_fit();
19201925

19211926
self.scope_ids_by_scope.shrink_to_fit();
19221927
self.scopes_by_node.shrink_to_fit();
@@ -1935,6 +1940,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
19351940
scopes_by_node: self.scopes_by_node,
19361941
use_def_maps,
19371942
enclosing_lambda_statements: self.enclosing_lambda_statements,
1943+
uses_by_collection: self.uses_by_collection,
19381944
imported_modules: Arc::new(self.imported_modules),
19391945
has_future_annotations: self.has_future_annotations,
19401946
enclosing_snapshots: self.enclosing_snapshots,
@@ -2353,6 +2359,16 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
23532359

23542360
self.visit_expr(&node.value);
23552361

2362+
// Unannotated collection literals must be standalone expressions to participate
2363+
// in full-scope bidirectional inference.
2364+
if node.targets.len() == 1
2365+
&& (node.value.is_list_expr()
2366+
|| node.value.is_set_expr()
2367+
|| node.value.is_dict_expr())
2368+
{
2369+
self.add_standalone_assigned_expression(&node.value, node);
2370+
}
2371+
23562372
// Optimization for the common case: if there's just one target, and it's not an
23572373
// unpacking, and the target is a simple name, we don't need the RHS to be a
23582374
// standalone expression at all.
@@ -3174,22 +3190,38 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
31743190
impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
31753191
fn visit_stmt(&mut self, stmt: &'ast ast::Stmt) {
31763192
self.push_statement(CurrentStatement {
3177-
lambda_exprs: Vec::new(),
3193+
lambda_expressions: Vec::new(),
3194+
collection_uses: FxHashMap::default(),
31783195
});
31793196

31803197
self.visit_stmt_impl(stmt);
31813198

31823199
let current_statement = self.pop_statement();
3183-
if !current_statement.lambda_exprs.is_empty() {
3200+
3201+
if current_statement.lambda_expressions.is_empty()
3202+
&& current_statement.collection_uses.is_empty()
3203+
{
3204+
return;
3205+
}
3206+
3207+
let standalone_statement = self.add_standalone_statement(stmt);
3208+
for lambda in current_statement.lambda_expressions {
31843209
// The body of a lambda expression needs access to the `Callable` type
31853210
// context the lambda is being inferred with, and so any statement
31863211
// containing a lambda must be inferable as a standalone statement
31873212
// to avoid large scope-level cycles.
3188-
let standalone_stmt = self.add_standalone_statement(stmt);
3189-
for lambda in current_statement.lambda_exprs {
3190-
self.enclosing_lambda_statements
3191-
.insert(lambda.into(), standalone_stmt);
3192-
}
3213+
self.enclosing_lambda_statements
3214+
.insert(lambda.into(), standalone_statement);
3215+
}
3216+
for (definition, use_expression) in current_statement.collection_uses {
3217+
// The inferred element type of collection literal depends on uses
3218+
// of the collection in its containing scope, and so each use
3219+
// must be part of an standalone inferable statement to avoid
3220+
// large scope-level cycles.
3221+
self.uses_by_collection
3222+
.entry(definition)
3223+
.or_default()
3224+
.push((standalone_statement, use_expression));
31933225
}
31943226
}
31953227

@@ -3295,12 +3327,38 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
32953327

32963328
if let Some((place_expr, is_use, is_definition)) = deferred_effects {
32973329
let place_id = self.add_place(place_expr);
3330+
32983331
if is_use {
3299-
self.record_place_use(place_id, expr);
3332+
let use_id = self.record_place_use(place_id, expr);
3333+
3334+
// Keep track of any uses of collection literals.
3335+
let use_def = &self.use_def_maps[self.current_scope()];
3336+
for binding in use_def.bindings_at_use(use_id) {
3337+
let Some(definition) =
3338+
use_def.definition(binding.binding()).definition()
3339+
else {
3340+
continue;
3341+
};
3342+
3343+
if let Some(current_statement) = self.current_statements.last_mut()
3344+
&& definition
3345+
.kind(self.db)
3346+
.is_unannotated_collection_literal(self.module)
3347+
{
3348+
// Note that if the same collection is referenced multiple times
3349+
// in the same statement, we only store the first occurrence.
3350+
current_statement
3351+
.collection_uses
3352+
.entry(definition)
3353+
.or_insert(expr.into());
3354+
}
3355+
}
33003356
}
3357+
33013358
if is_definition {
33023359
self.record_place_definition(place_id, expr);
33033360
}
3361+
33043362
if let Some(unpack_position) = self
33053363
.current_assignment_mut()
33063364
.and_then(CurrentAssignment::unpack_position_mut)
@@ -3324,7 +3382,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
33243382
}
33253383
ast::Expr::Lambda(lambda) => {
33263384
if let Some(current_statement) = self.current_statement_mut() {
3327-
current_statement.lambda_exprs.push(lambda);
3385+
current_statement.lambda_expressions.push(lambda);
33283386
}
33293387

33303388
if let Some(parameters) = &lambda.parameters {
@@ -3768,9 +3826,11 @@ impl<'ast> From<&'ast ast::ExprNamed> for CurrentAssignment<'ast, '_> {
37683826
}
37693827
}
37703828

3771-
struct CurrentStatement<'ast> {
3772-
/// The lambda expressions part of this statement.
3773-
lambda_exprs: Vec<&'ast ast::ExprLambda>,
3829+
struct CurrentStatement<'ast, 'db> {
3830+
/// A list of lambda expressions contained in this statement.
3831+
lambda_expressions: Vec<&'ast ast::ExprLambda>,
3832+
/// A list of collection definitions whose uses are contained in this statement.
3833+
collection_uses: FxHashMap<Definition<'db>, ExpressionNodeKey>,
37743834
}
37753835

37763836
#[derive(Debug, PartialEq)]

‎crates/ty_python_core/src/definition.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -917,6 +917,15 @@ impl DefinitionKind<'_> {
917917
matches!(self, DefinitionKind::LoopHeader(_))
918918
}
919919

920+
pub fn is_unannotated_collection_literal(&self, module: &ParsedModuleRef) -> bool {
921+
let value = match self {
922+
DefinitionKind::Assignment(assignment) => assignment.value(module),
923+
_ => return false,
924+
};
925+
926+
value.is_dict_expr() || value.is_set_expr() || value.is_list_expr()
927+
}
928+
920929
/// Returns `true` if this definition is user-visible (i.e., not an internal
921930
/// control-flow construct like a loop header definition).
922931
pub const fn is_user_visible(&self) -> bool {
@@ -1491,6 +1500,14 @@ impl<'db> LoopHeaderDefinitionKind<'db> {
14911500
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, salsa::Update, get_size2::GetSize)]
14921501
pub struct DefinitionNodeKey(NodeKey);
14931502

1503+
impl DefinitionNodeKey {
1504+
pub fn from_assignment(node: &ast::StmtAssign) -> impl Iterator<Item = DefinitionNodeKey> {
1505+
node.targets
1506+
.iter()
1507+
.map(|target| Self(NodeKey::from_node(target)))
1508+
}
1509+
}
1510+
14941511
impl From<&ast::Alias> for DefinitionNodeKey {
14951512
fn from(node: &ast::Alias) -> Self {
14961513
Self(NodeKey::from_node(node))

0 commit comments

Comments
 (0)