add statement-level inference to avoid scope cycles

Author: ibraheemdev

crates/ty_python_semantic/resources/mdtest/bidirectional.md

@@ -428,6 +428,9 @@ from typing import Callable, TypedDict
 class Bar(TypedDict):
     bar: int
 
+def id[T](x: T) -> T:
+    return x
+
 f1 = lambda x: {"bar": 1}
 reveal_type(f1)  # revealed: (x) -> dict[str, int]
 
@@ -442,34 +445,46 @@ reveal_type(f3)  # revealed: (int, /) -> Bar
 f4: Callable[[str], str] = lambda x: reveal_type(x)  # revealed: str
 reveal_type(f4)  # revealed: (x: str) -> str
 
+f5: Callable[[str], str] = id(lambda x: reveal_type(x))  # revealed: str
+reveal_type(f5)  # revealed: (x: str) -> str
+
 # TODO: This should not error once we support `Unpack`.
 # error: [invalid-assignment]
-f5: Callable[[*tuple[int, ...]], None] = lambda x, y, z: None
-reveal_type(f5)  # revealed: (tuple[int, ...], /) -> None
+f6: Callable[[*tuple[int, ...]], None] = lambda x, y, z: None
+reveal_type(f6)  # revealed: (tuple[int, ...], /) -> None
 
-f6: Callable[[int, str], None] = lambda *args: None
-reveal_type(f6)  # revealed: (*args) -> None
+f7: Callable[[int, str], None] = lambda *args: None
+reveal_type(f7)  # revealed: (*args) -> None
 
 # N.B. `Callable` annotations only support positional parameters.
 # error: [invalid-assignment]
-f7: Callable[[int], None] = lambda *, x=1: None
-reveal_type(f7)  # revealed: (int, /) -> None
+f8: Callable[[int], None] = lambda *, x=1: None
+reveal_type(f8)  # revealed: (int, /) -> None
 
 # TODO: This should reveal `(*args: int, *, x=1) -> None` once we support `Unpack`.
-f8: Callable[[*tuple[int, ...], int], None] = lambda *args, x=1: None
-reveal_type(f8)  # revealed: (*args, *, x=1) -> None
+f9: Callable[[*tuple[int, ...], int], None] = lambda *args, x=1: None
+reveal_type(f9)  # revealed: (*args, *, x=1) -> None
 
-f9: Callable[[str, int, str], tuple[str, int, str]] = lambda x, y, z: reveal_type((x, y, z))  # revealed: tuple[str, int, str]
-reveal_type(f9)  # revealed: (x: str, y: int, z: str) -> tuple[str, int, str]
+f10: Callable[[str, int, str], tuple[str, int, str]] = lambda x, y, z: reveal_type((x, y, z))  # revealed: tuple[str, int, str]
+reveal_type(f10)  # revealed: (x: str, y: int, z: str) -> tuple[str, int, str]
 
 # TODO: This should reveal `tuple[int, ...]` once we support `Unpack`.
-f10: Callable[[*tuple[int, ...]], tuple[int, ...]] = lambda *args: reveal_type(args)  # revealed: tuple[Unknown, ...]
-reveal_type(f10)  # revealed: (*args) -> tuple[Unknown, ...]
+f11: Callable[[*tuple[int, ...]], tuple[int, ...]] = lambda *args: reveal_type(args)  # revealed: tuple[Unknown, ...]
+reveal_type(f11)  # revealed: (*args) -> tuple[Unknown, ...]
 
 # TODO: Better generic call inference.
 def _(x: list[int]):
-    f11 = list(map(lambda y: y + 1, x))
-    reveal_type(f11)  # revealed: list[Unknown]
+    f12 = list(map(lambda y: y + 1, x))
+    reveal_type(f12)  # revealed: list[Unknown]
+
+def _() -> Callable[[int], int]:
+    return id(lambda x: reveal_type(x))  # revealed: int
+
+def _():
+    def takes_callable(_: Callable[[int], int]): ...
+
+    takes_callable(lambda x: reveal_type(x))  # revealed: int
+    takes_callable(id(id(lambda x: reveal_type(x))))  # revealed: int
 ```
 
 We do not currently account for type annotations present later in the scope:

crates/ty_python_semantic/src/semantic_index.rs

@@ -27,6 +27,7 @@ pub use crate::semantic_index::scope::FileScopeId;
 use crate::semantic_index::scope::{
     NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopeKind, ScopeLaziness,
 };
+use crate::semantic_index::statement::{Statement, StatementNodeKey};
 use crate::semantic_index::symbol::ScopedSymbolId;
 use crate::semantic_index::use_def::{EnclosingSnapshotKey, ScopedEnclosingSnapshotId, UseDefMap};
 use crate::semantic_model::HasTrackedScope;
@@ -42,6 +43,7 @@ pub(crate) mod predicate;
 mod re_exports;
 mod reachability_constraints;
 pub(crate) mod scope;
+pub mod statement;
 pub(crate) mod symbol;
 mod use_def;
 
@@ -321,9 +323,15 @@ pub(crate) struct SemanticIndex<'db> {
     /// Map from a standalone expression to its [`Expression`] ingredient.
     expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
 
+    /// Map from a standalone statemetn to its [`Statement`] ingredient.
+    statements_by_node: FxHashMap<StatementNodeKey, Statement<'db>>,
+
     /// Map from nodes that create a scope to the scope they create.
     scopes_by_node: FxHashMap<NodeWithScopeKey, FileScopeId>,
 
+    /// Maps from lambda expressions to their containing statement.
+    enclosing_lambda_statements: FxHashMap<ExpressionNodeKey, Statement<'db>>,
+
     /// Map from the file-local [`FileScopeId`] to the salsa-ingredient [`ScopeId`].
     scope_ids_by_scope: IndexVec<FileScopeId, ScopeId<'db>>,
 
@@ -473,6 +481,13 @@ impl<'db> SemanticIndex<'db> {
             .map(|node_ref| self.expect_single_definition(node_ref))
     }
 
+    pub(crate) fn enclosing_lambda_statement(
+        &self,
+        lambda: ExpressionNodeKey,
+    ) -> Option<Statement<'db>> {
+        self.enclosing_lambda_statements.get(&lambda).copied()
+    }
+
     pub(crate) fn is_scope_reachable(&self, db: &'db dyn Db, scope_id: FileScopeId) -> bool {
         self.parent_scope_id(scope_id)
             .is_none_or(|parent_scope_id| {
@@ -601,6 +616,13 @@ impl<'db> SemanticIndex<'db> {
             .contains_key(&expression_key.into())
     }
 
+    pub(crate) fn try_statement(
+        &self,
+        statement_key: impl Into<StatementNodeKey>,
+    ) -> Option<Statement<'db>> {
+        self.statements_by_node.get(&statement_key.into()).copied()
+    }
+
     /// Returns the id of the scope that `node` creates.
     /// This is different from [`definition::Definition::scope`] which
     /// returns the scope in which that definition is defined in.

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -47,6 +47,7 @@ use crate::semantic_index::scope::{
     FileScopeId, NodeWithScopeKey, NodeWithScopeKind, NodeWithScopeRef,
 };
 use crate::semantic_index::scope::{Scope, ScopeId, ScopeKind, ScopeLaziness};
+use crate::semantic_index::statement::{Statement, StatementNodeKey};
 use crate::semantic_index::symbol::{ScopedSymbolId, Symbol};
 use crate::semantic_index::use_def::{
     EnclosingSnapshotKey, FlowSnapshot, PreviousDefinitions, ScopedDefinitionId,
@@ -100,6 +101,8 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
     /// The assignments we're currently visiting, with
     /// the most recent visit at the end of the Vec
     current_assignments: Vec<CurrentAssignment<'ast, 'db>>,
+    /// The statement we're currently visiting.
+    current_statement: Option<(&'ast ast::Stmt, FileScopeId)>,
     /// The match case we're currently visiting.
     current_match_case: Option<CurrentMatchCase<'ast>>,
     /// The name of the first function parameter of the innermost function that we're currently visiting.
@@ -129,6 +132,8 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
     scopes_by_expression: ExpressionsScopeMapBuilder,
     definitions_by_node: FxHashMap<DefinitionNodeKey, Definitions<'db>>,
     expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
+    statements_by_node: FxHashMap<StatementNodeKey, Statement<'db>>,
+    enclosing_lambda_statements: FxHashMap<ExpressionNodeKey, Statement<'db>>,
     imported_modules: FxHashSet<ModuleName>,
     seen_submodule_imports: FxHashSet<String>,
     /// Hashset of all [`FileScopeId`]s that correspond to [generator functions].
@@ -150,6 +155,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             module: module_ref,
             scope_stack: Vec::new(),
             current_assignments: vec![],
+            current_statement: None,
             current_match_case: None,
             current_first_parameter_name: None,
             try_node_context_stack_manager: TryNodeContextStackManager::default(),
@@ -167,6 +173,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             scopes_by_node: FxHashMap::default(),
             definitions_by_node: FxHashMap::default(),
             expressions_by_node: FxHashMap::default(),
+            statements_by_node: FxHashMap::default(),
+            enclosing_lambda_statements: FxHashMap::default(),
 
             seen_submodule_imports: FxHashSet::default(),
             imported_modules: FxHashSet::default(),
@@ -1493,6 +1501,22 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         expression
     }
 
+    fn add_standalone_statement(
+        &mut self,
+        statement_node: &ast::Stmt,
+        scope: FileScopeId,
+    ) -> Statement<'db> {
+        let statement = Statement::new(
+            self.db,
+            self.file,
+            scope,
+            AstNodeRef::new(self.module, statement_node),
+        );
+        self.statements_by_node
+            .insert(statement_node.into(), statement);
+        statement
+    }
+
     fn with_type_params(
         &mut self,
         with_scope: NodeWithScopeRef,
@@ -1858,6 +1882,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         use_def_maps.shrink_to_fit();
         ast_ids.shrink_to_fit();
         self.definitions_by_node.shrink_to_fit();
+        self.statements_by_node.shrink_to_fit();
+        self.enclosing_lambda_statements.shrink_to_fit();
 
         self.scope_ids_by_scope.shrink_to_fit();
         self.scopes_by_node.shrink_to_fit();
@@ -1869,11 +1895,13 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             scopes: self.scopes,
             definitions_by_node: self.definitions_by_node,
             expressions_by_node: self.expressions_by_node,
+            statements_by_node: self.statements_by_node,
             scope_ids_by_scope: self.scope_ids_by_scope,
             ast_ids,
             scopes_by_expression: self.scopes_by_expression.build(),
             scopes_by_node: self.scopes_by_node,
             use_def_maps,
+            enclosing_lambda_statements: self.enclosing_lambda_statements,
             imported_modules: Arc::new(self.imported_modules),
             has_future_annotations: self.has_future_annotations,
             enclosing_snapshots: self.enclosing_snapshots,
@@ -1892,10 +1920,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         self.source_text
             .get_or_init(|| source_text(self.db, self.file))
     }
-}
 
-impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
-    fn visit_stmt(&mut self, stmt: &'ast ast::Stmt) {
+    fn visit_stmt_impl(&mut self, stmt: &'ast ast::Stmt) {
         self.with_semantic_checker(|semantic, context| semantic.visit_stmt(stmt, context));
 
         self.current_use_def_map_mut()
@@ -3107,6 +3133,14 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
             }
         }
     }
+}
+
+impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
+    fn visit_stmt(&mut self, stmt: &'ast ast::Stmt) {
+        self.current_statement = Some((stmt, self.current_scope()));
+        self.visit_stmt_impl(stmt);
+        self.current_statement = None;
+    }
 
     fn visit_keyword(&mut self, keyword: &'ast ast::Keyword) {
         walk_keyword(self, keyword);
@@ -3238,6 +3272,15 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 }
             }
             ast::Expr::Lambda(lambda) => {
+                // The body of a lambda expression needs access to the `Callable` type
+                // context the lambda is being inferred with, and so any statement
+                // containing a lambda must be inferable as a standalone statement.
+                if let Some((stmt, scope)) = self.current_statement {
+                    let standalone_stmt = self.add_standalone_statement(stmt, scope);
+                    self.enclosing_lambda_statements
+                        .insert(lambda.into(), standalone_stmt);
+                }
+
                 if let Some(parameters) = &lambda.parameters {
                     // The default value of the parameters needs to be evaluated in the
                     // enclosing scope.

crates/ty_python_semantic/src/semantic_index/statement.rs

@@ -0,0 +1,53 @@
+use crate::ast_node_ref::AstNodeRef;
+use crate::db::Db;
+use crate::node_key::NodeKey;
+use crate::semantic_index::scope::{FileScopeId, ScopeId};
+use ruff_db::files::File;
+use ruff_python_ast as ast;
+use salsa;
+
+/// An independently type-inferable statement.
+///
+/// ## Module-local type
+/// This type should not be used as part of any cross-module API because
+/// it holds a reference to the AST node. Range-offset changes
+/// then propagate through all usages, and deserialization requires
+/// reparsing the entire module.
+///
+/// E.g. don't use this type in:
+///
+/// * a return type of a cross-module query
+/// * a field of a type that is a return type of a cross-module query
+/// * an argument of a cross-module query
+#[salsa::tracked(debug, heap_size=ruff_memory_usage::heap_size)]
+pub(crate) struct Statement<'db> {
+    /// The file in which the statement occurs.
+    pub(crate) file: File,
+
+    /// The scope in which the statement occurs.
+    pub(crate) file_scope: FileScopeId,
+
+    /// The statement node.
+    #[no_eq]
+    #[tracked]
+    #[returns(ref)]
+    pub(crate) node_ref: AstNodeRef<ast::Stmt>,
+}
+
+// The Salsa heap is tracked separately.
+impl get_size2::GetSize for Statement<'_> {}
+
+impl<'db> Statement<'db> {
+    pub(crate) fn scope(self, db: &'db dyn Db) -> ScopeId<'db> {
+        self.file_scope(db).to_scope_id(db, self.file(db))
+    }
+}
+
+#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, salsa::Update, get_size2::GetSize)]
+pub(crate) struct StatementNodeKey(NodeKey);
+
+impl From<&ast::Stmt> for StatementNodeKey {
+    fn from(node: &ast::Stmt) -> Self {
+        Self(NodeKey::from_node(node))
+    }
+}