[red-knot] Store starred-expression annotation types (#14106)

## Summary

- Store the expression type for annotations that are starred expressions
(see [discussion
here](https://github.com/astral-sh/ruff/pull/14091#discussion_r1828332857))
- Use `self.store_expression_type(…)` consistently throughout, as it
makes sure that no double-insertion errors occur.

closes #14115

## Test Plan

Added an invalid-syntax example to the corpus which leads to a panic on
`main`. Also added a Markdown test with a valid-syntax example that
would lead to a panic once we implement function parameter inference.

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

crates/red_knot_python_semantic/resources/mdtest/annotations/starred_expressions.md

@@ -0,0 +1,18 @@
+# Starred expression annotations
+
+Type annotations for `*args` can be starred expressions themselves:
+
+```py
+from typing_extensions import TypeVarTuple
+
+Ts = TypeVarTuple("Ts")
+
+def append_int(*args: *Ts) -> tuple[*Ts, int]:
+    # TODO: should show some representation of the variadic generic type
+    reveal_type(args)  # revealed: @Todo
+
+    return (*args, 1)
+
+# TODO should be tuple[Literal[True], Literal["a"], int]
+reveal_type(append_int(True, "a"))  # revealed: @Todo
+```

crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs

@@ -87,6 +87,14 @@ pub trait HasScopedAstId {
     fn scoped_ast_id(&self, db: &dyn Db, scope: ScopeId) -> Self::Id;
 }
 
+impl<T: HasScopedAstId> HasScopedAstId for Box<T> {
+    type Id = <T as HasScopedAstId>::Id;
+
+    fn scoped_ast_id(&self, db: &dyn Db, scope: ScopeId) -> Self::Id {
+        self.as_ref().scoped_ast_id(db, scope)
+    }
+}
+
 /// Uniquely identifies an [`ast::Expr`] in a [`crate::semantic_index::symbol::FileScopeId`].
 #[newtype_index]
 pub struct ScopedExpressionId;

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1443,8 +1443,8 @@ impl<'db> TypeInferenceBuilder<'db> {
             TargetKind::Name => value_ty,
         };
 
+        self.store_expression_type(name, target_ty);
         self.add_binding(name.into(), definition, target_ty);
-        self.types.expressions.insert(name_ast_id, target_ty);
     }
 
     fn infer_annotated_assignment_statement(&mut self, assignment: &ast::StmtAnnAssign) {
@@ -1697,10 +1697,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                 .unwrap_with_diagnostic(iterable.into(), &mut self.diagnostics)
         };
 
-        self.types.expressions.insert(
+        self.store_expression_type(target, loop_var_value_ty);
-            target.scoped_ast_id(self.db, self.scope()),
-            loop_var_value_ty,
-        );
         self.add_binding(target.into(), definition, loop_var_value_ty);
     }
 
@@ -2027,7 +2024,11 @@ impl<'db> TypeInferenceBuilder<'db> {
         ty
     }
 
-    fn store_expression_type(&mut self, expression: &ast::Expr, ty: Type<'db>) {
+    fn store_expression_type(
+        &mut self,
+        expression: &impl HasScopedAstId<Id = ScopedExpressionId>,
+        ty: Type<'db>,
+    ) {
         let expr_id = expression.scoped_ast_id(self.db, self.scope());
         let previous = self.types.expressions.insert(expr_id, ty);
         assert_eq!(previous, None);
@@ -3803,31 +3804,31 @@ impl<'db> TypeInferenceBuilder<'db> {
 impl<'db> TypeInferenceBuilder<'db> {
     fn infer_annotation_expression(&mut self, expression: &ast::Expr) -> Type<'db> {
         // https://typing.readthedocs.io/en/latest/spec/annotations.html#grammar-token-expression-grammar-annotation_expression
-        match expression {
+        let annotation_ty = match expression {
             // TODO: parse the expression and check whether it is a string annotation, since they
             // can be annotation expressions distinct from type expressions.
             // https://typing.readthedocs.io/en/latest/spec/annotations.html#string-annotations
-            ast::Expr::StringLiteral(_literal) => {
+            ast::Expr::StringLiteral(_literal) => Type::Todo,
-                self.store_expression_type(expression, Type::Todo);
-                Type::Todo
-            }
 
             // Annotation expressions also get special handling for `*args` and `**kwargs`.
             ast::Expr::Starred(starred) => self.infer_starred_expression(starred),
 
             // All other annotation expressions are (possibly) valid type expressions, so handle
             // them there instead.
-            type_expr => self.infer_type_expression(type_expr),
+            type_expr => self.infer_type_expression_no_store(type_expr),
-        }
+        };
+
+        self.store_expression_type(expression, annotation_ty);
+        annotation_ty
     }
 }
 
 /// Type expressions
 impl<'db> TypeInferenceBuilder<'db> {
-    fn infer_type_expression(&mut self, expression: &ast::Expr) -> Type<'db> {
+    fn infer_type_expression_no_store(&mut self, expression: &ast::Expr) -> Type<'db> {
         // https://typing.readthedocs.io/en/latest/spec/annotations.html#grammar-token-expression-grammar-type_expression
 
-        let ty = match expression {
+        match expression {
             ast::Expr::Name(name) => {
                 debug_assert_eq!(name.ctx, ast::ExprContext::Load);
                 self.infer_name_expression(name).to_instance(self.db)
@@ -3983,12 +3984,12 @@ impl<'db> TypeInferenceBuilder<'db> {
             }
 
             ast::Expr::IpyEscapeCommand(_) => todo!("Implement Ipy escape command support"),
-        };
+        }
-
+    }
-        let expr_id = expression.scoped_ast_id(self.db, self.scope());
-        let previous = self.types.expressions.insert(expr_id, ty);
-        assert!(previous.is_none());
 
+    fn infer_type_expression(&mut self, expression: &ast::Expr) -> Type<'db> {
+        let ty = self.infer_type_expression_no_store(expression);
+        self.store_expression_type(expression, ty);
         ty
     }
 

crates/red_knot_workspace/resources/test/corpus/25_func_annotations_starred.py

@@ -0,0 +1,8 @@
+# Regression test for https://github.com/astral-sh/ruff/issues/14115
+#
+# This is invalid syntax, but should not lead to a crash.
+
+def f() -> *int: ...
+
+
+f()