Port return conditional expression checking improvements

Co-authored-by: Gabriela Araujo Britto <[email protected]>

Author: Andarist

internal/ast/ast.go

@@ -5984,6 +5984,10 @@ func (node *ConditionalExpression) computeSubtreeFacts() SubtreeFacts {
 		propagateSubtreeFacts(node.WhenFalse)
 }
 
+func IsConditionalExpression(node *Node) bool {
+	return node.Kind == KindConditionalExpression
+}
+
 // PropertyAccessExpression
 
 type PropertyAccessExpression struct {

internal/checker/checker.go

@@ -3883,23 +3883,41 @@ func (c *Checker) checkReturnStatement(node *ast.Node) {
 			}
 		} else if c.getReturnTypeFromAnnotation(container) != nil {
 			unwrappedReturnType := core.OrElse(c.unwrapReturnType(returnType, functionFlags), returnType)
-			unwrappedExprType := exprType
-			if functionFlags&FunctionFlagsAsync != 0 {
-				unwrappedExprType = c.checkAwaitedType(exprType, false /*withAlias*/, node, diagnostics.The_return_type_of_an_async_function_must_either_be_a_valid_promise_or_must_not_contain_a_callable_then_member)
-			}
-			if unwrappedReturnType != nil {
-				// If the function has a return type, but promisedType is
-				// undefined, an error will be reported in checkAsyncFunctionReturnType
-				// so we don't need to report one here.
-				c.checkTypeAssignableToAndOptionallyElaborate(unwrappedExprType, unwrappedReturnType, node, exprNode, nil, nil)
-			}
+			c.checkReturnExpression(container, unwrappedReturnType, node, node.Expression(), exprType, false)
 		}
 	} else if !ast.IsConstructorDeclaration(container) && c.compilerOptions.NoImplicitReturns.IsTrue() && !c.isUnwrappedReturnTypeUndefinedVoidOrAny(container, returnType) {
 		// The function has a return type, but the return statement doesn't have an expression.
 		c.error(node, diagnostics.Not_all_code_paths_return_a_value)
 	}
 }
 
+// When checking an arrow expression such as `(x) => exp`, then `node` is the expression `exp`.
+// Otherwise, `node` is a return statement.
+func (c *Checker) checkReturnExpression(container *ast.Node, unwrappedReturnType *Type, node *ast.Node, expr *ast.Node, exprType *Type, inConditionalExpression bool) {
+	unwrappedExprType := exprType
+	functionFlags := getFunctionFlags(container)
+	if expr != nil {
+		unwrappedExpr := ast.SkipParentheses(expr)
+		if ast.IsConditionalExpression(unwrappedExpr) {
+			whenTrue := unwrappedExpr.AsConditionalExpression().WhenTrue
+			whenFalse := unwrappedExpr.AsConditionalExpression().WhenFalse
+			c.checkReturnExpression(container, unwrappedReturnType, node, whenTrue, c.checkExpression(whenTrue), true /*inConditionalExpression*/)
+			c.checkReturnExpression(container, unwrappedReturnType, node, whenFalse, c.checkExpression(whenFalse), true /*inConditionalExpression*/)
+			return
+		}
+	}
+	inReturnStatement := node.Kind == ast.KindReturnStatement
+	if functionFlags&FunctionFlagsAsync != 0 {
+		unwrappedExprType = c.checkAwaitedType(exprType, false /*withAlias*/, node, diagnostics.The_return_type_of_an_async_function_must_either_be_a_valid_promise_or_must_not_contain_a_callable_then_member)
+	}
+	effectiveExpr := expr // The effective expression for diagnostics purposes.
+	if expr != nil {
+		effectiveExpr = c.getEffectiveCheckNode(expr)
+	}
+	errorNode := core.IfElse(inReturnStatement && !inConditionalExpression, node, effectiveExpr)
+	c.checkTypeAssignableToAndOptionallyElaborate(unwrappedExprType, unwrappedReturnType, errorNode, effectiveExpr, nil, nil)
+}
+
 func (c *Checker) checkWithStatement(node *ast.Node) {
 	if !c.checkGrammarStatementInAmbientContext(node) {
 		if node.Flags&ast.NodeFlagsAwaitContext != 0 {
@@ -9775,13 +9793,7 @@ func (c *Checker) checkFunctionExpressionOrObjectLiteralMethodDeferred(node *ast
 			if returnType != nil {
 				returnOrPromisedType := c.unwrapReturnType(returnType, functionFlags)
 				if returnOrPromisedType != nil {
-					effectiveCheckNode := c.getEffectiveCheckNode(body)
-					if (functionFlags & FunctionFlagsAsyncGenerator) == FunctionFlagsAsync {
-						awaitedType := c.checkAwaitedType(exprType, false /*withAlias*/, effectiveCheckNode, diagnostics.The_return_type_of_an_async_function_must_either_be_a_valid_promise_or_must_not_contain_a_callable_then_member)
-						c.checkTypeAssignableToAndOptionallyElaborate(awaitedType, returnOrPromisedType, effectiveCheckNode, effectiveCheckNode, nil, nil)
-					} else {
-						c.checkTypeAssignableToAndOptionallyElaborate(exprType, returnOrPromisedType, effectiveCheckNode, effectiveCheckNode, nil, nil)
-					}
+					c.checkReturnExpression(node, returnOrPromisedType, body, body, exprType, false)
 				}
 			}
 		}

testdata/baselines/reference/submodule/compiler/conditionalReturnExpression.errors.txt

@@ -1,19 +1,24 @@
-conditionalReturnExpression.ts(2,5): error TS2322: Type '1 | 2' is not assignable to type '3'.
-  Type '1' is not assignable to type '3'.
+conditionalReturnExpression.ts(2,18): error TS2322: Type '1' is not assignable to type '3'.
+conditionalReturnExpression.ts(2,23): error TS2322: Type '2' is not assignable to type '3'.
+conditionalReturnExpression.ts(8,43): error TS2322: Type 'number' is not assignable to type 'string'.
+conditionalReturnExpression.ts(19,71): error TS2322: Type 'number' is not assignable to type 'string'.
 
 
-==== conditionalReturnExpression.ts (1 errors) ====
+==== conditionalReturnExpression.ts (4 errors) ====
     function return1(x: boolean): 3 {
         return (x ? (1) : 2);
-        ~~~~~~
-!!! error TS2322: Type '1 | 2' is not assignable to type '3'.
-!!! error TS2322:   Type '1' is not assignable to type '3'.
+                     ~
+!!! error TS2322: Type '1' is not assignable to type '3'.
+                          ~
+!!! error TS2322: Type '2' is not assignable to type '3'.
     }
 
     declare function getAny(): any;
 
     function return2(x: string): string {
         return x.startsWith("a") ? getAny() : 1;
+                                              ~
+!!! error TS2322: Type 'number' is not assignable to type 'string'.
     }
 
     function return3(x: string): string {
@@ -25,5 +30,7 @@ conditionalReturnExpression.ts(2,5): error TS2322: Type '1 | 2' is not assignabl
     }
 
     const return5 = (x: string): string => x.startsWith("a") ? getAny() : 1;
+                                                                          ~
+!!! error TS2322: Type 'number' is not assignable to type 'string'.
 
     const return6 = (x: string): string => (x.startsWith("a") ? getAny() : 1) as string;