Add scope and definitions for comprehensions (#12748)

## Summary

This PR adds scope and definition for comprehension nodes. This includes
the following nodes:
* List comprehension
* Dictionary comprehension
* Set comprehension 
* Generator expression

### Scope

Each expression here adds it's own scope with one caveat - the `iter`
expression of the first generator is part of the parent scope. For
example, in the following code snippet the `iter1` variable is evaluated
in the outer scope.

```py
[x for x in iter1]
```

> The iterable expression in the leftmost for clause is evaluated
directly in the enclosing scope and then passed as an argument to the
implicitly nested scope.
>
> Reference:
https://docs.python.org/3/reference/expressions.html#displays-for-lists-sets-and-dictionaries

There's another special case for assignment expressions:

> There is one special case: an assignment expression occurring in a
list, set or dict comprehension or in a generator expression (below
collectively referred to as “comprehensions”) binds the target in the
containing scope, honoring a nonlocal or global declaration for the
target in that scope, if one exists.
>
> Reference: https://peps.python.org/pep-0572/#scope-of-the-target

For example, in the following code snippet, the variables `a` and `b`
are available after the comprehension while `x` isn't:
```py
[a := 1 for x in range(2) if (b := 2)]
```

### Definition

Each comprehension node adds a single definition, the "target" variable
(`[_ for target in iter]`). This has been accounted for and a new
variant has been added to `DefinitionKind`.

### Type Inference

Currently, type inference is limited to a single scope. It doesn't
_enter_ in another scope to infer the types of the remaining expressions
of a node. To accommodate this, the type inference for a **scope**
requires new methods which _doesn't_ infer the type of the `iter`
expression of the leftmost outer generator (that's defined in the
enclosing scope).

The type inference for the scope region is split into two parts:
* `infer_generator_expression` (similarly for comprehensions) infers the
type of the `iter` expression of the leftmost outer generator
* `infer_generator_expression_scope` (similarly for comprehension)
infers the type of the remaining expressions except for the one
mentioned in the previous point

The type inference for the **definition** also needs to account for this
special case of leftmost generator. This is done by defining a `first`
boolean parameter which indicates whether this comprehension definition
occurs first in the enclosing expression.

## Test Plan

New test cases were added to validate multiple scenarios. Refer to the
documentation for each test case which explains what is being tested.

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -13,8 +13,8 @@ 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,
+    AssignmentDefinitionNodeRef, ComprehensionDefinitionNodeRef, Definition, DefinitionNodeKey,
+    DefinitionNodeRef, ImportFromDefinitionNodeRef,
 };
 use crate::semantic_index::expression::Expression;
 use crate::semantic_index::symbol::{
@@ -174,7 +174,7 @@ impl<'db> SemanticIndexBuilder<'db> {
         symbol: ScopedSymbolId,
         definition_node: impl Into<DefinitionNodeRef<'a>>,
     ) -> Definition<'db> {
-        let definition_node = definition_node.into();
+        let definition_node: DefinitionNodeRef<'_> = definition_node.into();
         let definition = Definition::new(
             self.db,
             self.file,
@@ -258,6 +258,49 @@ impl<'db> SemanticIndexBuilder<'db> {
         nested_scope
     }
 
+    /// Visit a list of [`Comprehension`] nodes, assumed to be the "generators" that compose a
+    /// comprehension (that is, the `for x in y` and `for y in z` parts of `x for x in y for y in z`.)
+    ///
+    /// [`Comprehension`]: ast::Comprehension
+    fn visit_generators(&mut self, scope: NodeWithScopeRef, generators: &'db [ast::Comprehension]) {
+        let mut generators_iter = generators.iter();
+
+        let Some(generator) = generators_iter.next() else {
+            unreachable!("Expression must contain at least one generator");
+        };
+
+        // The `iter` of the first generator is evaluated in the outer scope, while all subsequent
+        // nodes are evaluated in the inner scope.
+        self.visit_expr(&generator.iter);
+        self.push_scope(scope);
+
+        self.current_assignment = Some(CurrentAssignment::Comprehension {
+            node: generator,
+            first: true,
+        });
+        self.visit_expr(&generator.target);
+        self.current_assignment = None;
+
+        for expr in &generator.ifs {
+            self.visit_expr(expr);
+        }
+
+        for generator in generators_iter {
+            self.visit_expr(&generator.iter);
+
+            self.current_assignment = Some(CurrentAssignment::Comprehension {
+                node: generator,
+                first: false,
+            });
+            self.visit_expr(&generator.target);
+            self.current_assignment = None;
+
+            for expr in &generator.ifs {
+                self.visit_expr(expr);
+            }
+        }
+    }
+
     pub(super) fn build(mut self) -> SemanticIndex<'db> {
         let module = self.module;
         self.visit_body(module.suite());
@@ -476,8 +519,7 @@ where
         self.current_ast_ids().record_expression(expr);
 
         match expr {
-            ast::Expr::Name(name_node) => {
-                let ast::ExprName { id, ctx, .. } = name_node;
+            ast::Expr::Name(name_node @ ast::ExprName { id, ctx, .. }) => {
                 let flags = match ctx {
                     ast::ExprContext::Load => SymbolFlags::IS_USED,
                     ast::ExprContext::Store => SymbolFlags::IS_DEFINED,
@@ -500,8 +542,17 @@ where
                             self.add_definition(symbol, ann_assign);
                         }
                         Some(CurrentAssignment::Named(named)) => {
+                            // TODO(dhruvmanila): If the current scope is a comprehension, then the
+                            // named expression is implicitly nonlocal. This is yet to be
+                            // implemented.
                             self.add_definition(symbol, named);
                         }
+                        Some(CurrentAssignment::Comprehension { node, first }) => {
+                            self.add_definition(
+                                symbol,
+                                ComprehensionDefinitionNodeRef { node, first },
+                            );
+                        }
                         None => {}
                     }
                 }
@@ -527,7 +578,6 @@ where
                 }
                 self.push_scope(NodeWithScopeRef::Lambda(lambda));
                 self.visit_expr(lambda.body.as_ref());
-                self.pop_scope();
             }
             ast::Expr::If(ast::ExprIf {
                 body, test, orelse, ..
@@ -543,10 +593,66 @@ where
                 self.visit_expr(orelse);
                 self.flow_merge(&post_body);
             }
+            ast::Expr::ListComp(
+                list_comprehension @ ast::ExprListComp {
+                    elt, generators, ..
+                },
+            ) => {
+                self.visit_generators(
+                    NodeWithScopeRef::ListComprehension(list_comprehension),
+                    generators,
+                );
+                self.visit_expr(elt);
+            }
+            ast::Expr::SetComp(
+                set_comprehension @ ast::ExprSetComp {
+                    elt, generators, ..
+                },
+            ) => {
+                self.visit_generators(
+                    NodeWithScopeRef::SetComprehension(set_comprehension),
+                    generators,
+                );
+                self.visit_expr(elt);
+            }
+            ast::Expr::Generator(
+                generator @ ast::ExprGenerator {
+                    elt, generators, ..
+                },
+            ) => {
+                self.visit_generators(NodeWithScopeRef::GeneratorExpression(generator), generators);
+                self.visit_expr(elt);
+            }
+            ast::Expr::DictComp(
+                dict_comprehension @ ast::ExprDictComp {
+                    key,
+                    value,
+                    generators,
+                    ..
+                },
+            ) => {
+                self.visit_generators(
+                    NodeWithScopeRef::DictComprehension(dict_comprehension),
+                    generators,
+                );
+                self.visit_expr(key);
+                self.visit_expr(value);
+            }
             _ => {
                 walk_expr(self, expr);
             }
         }
+
+        if matches!(
+            expr,
+            ast::Expr::Lambda(_)
+                | ast::Expr::ListComp(_)
+                | ast::Expr::SetComp(_)
+                | ast::Expr::Generator(_)
+                | ast::Expr::DictComp(_)
+        ) {
+            self.pop_scope();
+        }
     }
 }
 
@@ -555,6 +661,10 @@ enum CurrentAssignment<'a> {
     Assign(&'a ast::StmtAssign),
     AnnAssign(&'a ast::StmtAnnAssign),
     Named(&'a ast::ExprNamed),
+    Comprehension {
+        node: &'a ast::Comprehension,
+        first: bool,
+    },
 }
 
 impl<'a> From<&'a ast::StmtAssign> for CurrentAssignment<'a> {