use query to avoid double-inferring expressions

crates/red_knot_python_semantic/src/semantic_index.rs

@@ -11,6 +11,7 @@ use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
 use crate::semantic_index::ast_ids::AstIds;
 use crate::semantic_index::builder::SemanticIndexBuilder;
 use crate::semantic_index::definition::{Definition, DefinitionNodeKey};
+use crate::semantic_index::expression::Expression;
 use crate::semantic_index::symbol::{
     FileScopeId, NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopedSymbolId, SymbolTable,
 };
@@ -81,14 +82,17 @@ pub(crate) struct SemanticIndex<'db> {
     /// List of all scopes in this file.
     scopes: IndexVec<FileScopeId, Scope>,
 
-    /// Maps expressions to their corresponding scope.
+    /// Map expressions to their corresponding scope.
     /// We can't use [`ExpressionId`] here, because the challenge is how to get from
     /// an [`ast::Expr`] to an [`ExpressionId`] (which requires knowing the scope).
     scopes_by_expression: FxHashMap<ExpressionNodeKey, FileScopeId>,
 
-    /// Maps from a node creating a definition to its definition.
+    /// Map from a node creating a definition to its definition.
     definitions_by_node: FxHashMap<DefinitionNodeKey, Definition<'db>>,
 
+    /// Map from a standalone expression to its [`Expression`] ingredient.
+    expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
+
     /// Map from nodes that create a scope to the scope they create.
     scopes_by_node: FxHashMap<NodeWithScopeKey, FileScopeId>,
 
@@ -175,14 +179,22 @@ impl<'db> SemanticIndex<'db> {
         AncestorsIter::new(self, scope)
     }
 
-    /// Returns the [`Definition`] salsa ingredient for `definition_node`.
+    /// Returns the [`Definition`] salsa ingredient for `definition_key`.
     pub(crate) fn definition(
         &self,
         definition_key: impl Into<DefinitionNodeKey>,
     ) -> Definition<'db> {
         self.definitions_by_node[&definition_key.into()]
     }
 
+    /// Returns the [`Expression`] ingredient for an expression node.
+    pub(crate) fn expression(
+        &self,
+        expression_key: impl Into<ExpressionNodeKey>,
+    ) -> Expression<'db> {
+        self.expressions_by_node[&expression_key.into()]
+    }
+
     /// Returns the id of the scope that `node` creates. This is different from [`Definition::scope`] which
     /// returns the scope in which that definition is defined in.
     pub(crate) fn node_scope(&self, node: NodeWithScopeRef) -> FileScopeId {

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -9,12 +9,14 @@ use ruff_python_ast as ast;
 use ruff_python_ast::name::Name;
 use ruff_python_ast::visitor::{walk_expr, walk_stmt, Visitor};
 
+use crate::ast_node_ref::AstNodeRef;
 use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
 use crate::semantic_index::ast_ids::AstIdsBuilder;
 use crate::semantic_index::definition::{
     AssignmentDefinitionNodeRef, Definition, DefinitionNodeKey, DefinitionNodeRef,
     ImportFromDefinitionNodeRef,
 };
+use crate::semantic_index::expression::Expression;
 use crate::semantic_index::symbol::{
     FileScopeId, NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopedSymbolId, SymbolFlags,
     SymbolTableBuilder,
@@ -41,6 +43,7 @@ pub(super) struct SemanticIndexBuilder<'db> {
     scopes_by_node: FxHashMap<NodeWithScopeKey, FileScopeId>,
     scopes_by_expression: FxHashMap<ExpressionNodeKey, FileScopeId>,
     definitions_by_node: FxHashMap<DefinitionNodeKey, Definition<'db>>,
+    expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
 }
 
 impl<'db> SemanticIndexBuilder<'db> {
@@ -61,6 +64,7 @@ impl<'db> SemanticIndexBuilder<'db> {
             scopes_by_expression: FxHashMap::default(),
             scopes_by_node: FxHashMap::default(),
             definitions_by_node: FxHashMap::default(),
+            expressions_by_node: FxHashMap::default(),
         };
 
         builder.push_scope_with_parent(NodeWithScopeRef::Module, None);
@@ -165,6 +169,22 @@ impl<'db> SemanticIndexBuilder<'db> {
         definition
     }
 
+    /// Record an expression that needs to be a Salsa ingredient, because we need to infer its type
+    /// standalone (type narrowing tests, RHS of an assignment.)
+    fn add_standalone_expression(&mut self, expression_node: &ast::Expr) {
+        let expression = Expression::new(
+            self.db,
+            self.file,
+            self.current_scope(),
+            #[allow(unsafe_code)]
+            unsafe {
+                AstNodeRef::new(self.module.clone(), expression_node)
+            },
+        );
+        self.expressions_by_node
+            .insert(expression_node.into(), expression);
+    }
+
     fn with_type_params(
         &mut self,
         with_params: &WithTypeParams,
@@ -245,6 +265,7 @@ impl<'db> SemanticIndexBuilder<'db> {
             symbol_tables,
             scopes: self.scopes,
             definitions_by_node: self.definitions_by_node,
+            expressions_by_node: self.expressions_by_node,
             scope_ids_by_scope: self.scope_ids_by_scope,
             ast_ids,
             scopes_by_expression: self.scopes_by_expression,
@@ -331,6 +352,7 @@ where
             ast::Stmt::Assign(node) => {
                 debug_assert!(self.current_assignment.is_none());
                 self.visit_expr(&node.value);
+                self.add_standalone_expression(&node.value);
                 self.current_assignment = Some(node.into());
                 for target in &node.targets {
                     self.visit_expr(target);

crates/red_knot_python_semantic/src/semantic_index/expression.rs

@@ -18,33 +18,14 @@ pub(crate) struct Expression<'db> {
     #[id]
     pub(crate) file_scope: FileScopeId,
 
-    /// The statement node containing the expression.
-    /// We hold the statement node to clarify which specific contexts these
-    /// independently-inferable expressions can occur in, and limit the variants of this enum.
+    /// The expression node.
     #[no_eq]
     #[return_ref]
-    pub(crate) ast: ExpressionKind,
+    pub(crate) node: AstNodeRef<ast::Expr>,
 }
 
 impl<'db> Expression<'db> {
-    /// Return the actual expression.
-    #[allow(unused)]
-    pub(crate) fn node(self, db: &'db dyn Db) -> &ast::Expr {
-        match self.ast(db) {
-            ExpressionKind::AssignedValue(node_ref) => &node_ref.node().value,
-            ExpressionKind::IfTest(node_ref) => &node_ref.node().test,
-        }
-    }
-
     pub(crate) fn scope(self, db: &'db dyn Db) -> ScopeId<'db> {
         self.file_scope(db).to_scope_id(db, self.file(db))
     }
 }
-
-#[derive(Clone, Debug)]
-#[allow(dead_code)]
-pub(crate) enum ExpressionKind {
-    AssignedValue(AstNodeRef<ast::StmtAssign>),
-    IfTest(AstNodeRef<ast::StmtIf>),
-    // TODO if-expressions, loops, etc
-}

crates/red_knot_python_semantic/src/types/infer.rs

@@ -192,8 +192,8 @@ impl<'db> TypeInferenceBuilder<'db> {
         }
     }
 
-    fn infer_region_expression(&mut self, _expression: Expression<'db>) {
-        todo!("avoid inferring assignment RHS multiple times")
+    fn infer_region_expression(&mut self, expression: Expression<'db>) {
+        self.infer_expression(expression.node(self.db));
     }
 
     fn infer_module(&mut self, module: &ast::ModModule) {
@@ -370,8 +370,12 @@ impl<'db> TypeInferenceBuilder<'db> {
         assignment: &ast::StmtAssign,
         definition: Definition<'db>,
     ) {
-        // TODO use query to avoid inferring value expression multiple times
-        let value_ty = self.infer_expression(&assignment.value);
+        let expression = self.index.expression(assignment.value.as_ref());
+        let result = infer_expression_types(self.db, expression);
+        self.extend(result);
+        let value_ty = self
+            .types
+            .expression_ty(assignment.value.scoped_ast_id(self.db, self.scope));
         self.types.definitions.insert(definition, value_ty);
     }