Analyze control flow effects of lambdas passed as arguments

src/compiler/binder.ts

@@ -101,6 +101,7 @@ import {
     getImmediatelyInvokedFunctionExpression,
     getJSDocHost,
     getJSDocTypeTag,
+    getLambdaArgument,
     getLeftmostAccessExpression,
     getNameOfDeclaration,
     getNameOrArgument,
@@ -318,6 +319,7 @@ import {
     unreachableCodeIsError,
     unusedLabelIsError,
     VariableDeclaration,
+    walkUpParenthesizedExpressions,
     WhileStatement,
     WithStatement,
 } from "./_namespaces/ts.js";
@@ -1015,6 +1017,8 @@ function createBinder(): (file: SourceFile, options: CompilerOptions) => void {
                 !(node as FunctionLikeDeclaration).asteriskToken &&
                 !!getImmediatelyInvokedFunctionExpression(node)
             ) || node.kind === SyntaxKind.ClassStaticBlockDeclaration;
+            const isLambdaArgument = (node.kind === SyntaxKind.FunctionExpression || node.kind === SyntaxKind.ArrowFunction) &&
+                walkUpParenthesizedExpressions(node.parent).kind === SyntaxKind.CallExpression;
             // A non-async, non-generator IIFE is considered part of the containing control flow. Return statements behave
             // similarly to break statements that exit to a label just past the statement body.
             if (!isImmediatelyInvoked) {
@@ -1025,7 +1029,7 @@ function createBinder(): (file: SourceFile, options: CompilerOptions) => void {
             }
             // We create a return control flow graph for IIFEs and constructors. For constructors
             // we use the return control flow graph in strict property initialization checks.
-            currentReturnTarget = isImmediatelyInvoked || node.kind === SyntaxKind.Constructor || (isInJSFile(node) && (node.kind === SyntaxKind.FunctionDeclaration || node.kind === SyntaxKind.FunctionExpression)) ? createBranchLabel() : undefined;
+            currentReturnTarget = isImmediatelyInvoked || isLambdaArgument || node.kind === SyntaxKind.Constructor || (isInJSFile(node) && (node.kind === SyntaxKind.FunctionDeclaration || node.kind === SyntaxKind.FunctionExpression)) ? createBranchLabel() : undefined;
             currentExceptionTarget = undefined;
             currentBreakTarget = undefined;
             currentContinueTarget = undefined;
@@ -1047,7 +1051,7 @@ function createBinder(): (file: SourceFile, options: CompilerOptions) => void {
             if (currentReturnTarget) {
                 addAntecedent(currentReturnTarget, currentFlow);
                 currentFlow = finishFlowLabel(currentReturnTarget);
-                if (node.kind === SyntaxKind.Constructor || node.kind === SyntaxKind.ClassStaticBlockDeclaration || (isInJSFile(node) && (node.kind === SyntaxKind.FunctionDeclaration || node.kind === SyntaxKind.FunctionExpression))) {
+                if (isLambdaArgument || node.kind === SyntaxKind.Constructor || node.kind === SyntaxKind.ClassStaticBlockDeclaration || (isInJSFile(node) && (node.kind === SyntaxKind.FunctionDeclaration || node.kind === SyntaxKind.FunctionExpression))) {
                     (node as FunctionLikeDeclaration | ClassStaticBlockDeclaration).returnFlowNode = currentFlow;
                 }
             }
@@ -2214,6 +2218,9 @@ function createBinder(): (file: SourceFile, options: CompilerOptions) => void {
                 }
             }
         }
+        if (some(node.arguments, arg => !!getLambdaArgument(arg))) {
+            currentFlow = createFlowNode(FlowFlags.LambdaArgs, node, currentFlow);
+        }
         if (node.expression.kind === SyntaxKind.PropertyAccessExpression) {
             const propertyAccess = node.expression as PropertyAccessExpression;
             if (isIdentifier(propertyAccess.name) && isNarrowableOperand(propertyAccess.expression) && isPushOrUnshiftIdentifier(propertyAccess.name)) {

src/compiler/checker.ts

@@ -1106,6 +1106,7 @@ import {
     WideningContext,
     WithStatement,
     YieldExpression,
+    getLambdaArgument,
 } from "./_namespaces/ts.js";
 import * as moduleSpecifiers from "./_namespaces/ts.moduleSpecifiers.js";
 import * as performance from "./_namespaces/ts.performance.js";
@@ -27773,8 +27774,8 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
                 }
                 noCacheCheck = false;
             }
-            if (flags & (FlowFlags.Assignment | FlowFlags.Condition | FlowFlags.ArrayMutation)) {
-                flow = (flow as FlowAssignment | FlowCondition | FlowArrayMutation).antecedent;
+            if (flags & (FlowFlags.Assignment | FlowFlags.Condition | FlowFlags.ArrayMutation | FlowFlags.LambdaArgs)) {
+                flow = (flow as FlowAssignment | FlowCondition | FlowArrayMutation | FlowCall).antecedent;
             }
             else if (flags & FlowFlags.Call) {
                 const signature = getEffectsSignature((flow as FlowCall).node);
@@ -27842,8 +27843,8 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
                 }
                 noCacheCheck = false;
             }
-            if (flags & (FlowFlags.Assignment | FlowFlags.Condition | FlowFlags.ArrayMutation | FlowFlags.SwitchClause)) {
-                flow = (flow as FlowAssignment | FlowCondition | FlowArrayMutation | FlowSwitchClause).antecedent;
+            if (flags & (FlowFlags.Assignment | FlowFlags.Condition | FlowFlags.ArrayMutation | FlowFlags.SwitchClause | FlowFlags.LambdaArgs)) {
+                flow = (flow as FlowAssignment | FlowCondition | FlowArrayMutation | FlowSwitchClause | FlowCall).antecedent;
             }
             else if (flags & FlowFlags.Call) {
                 if ((flow as FlowCall).node.expression.kind === SyntaxKind.SuperKeyword) {
@@ -27903,6 +27904,7 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
     function getFlowTypeOfReference(reference: Node, declaredType: Type, initialType = declaredType, flowContainer?: Node, flowNode = tryCast(reference, canHaveFlowNode)?.flowNode) {
         let key: string | undefined;
         let isKeySet = false;
+        let inLambdaArg = false;
         let flowDepth = 0;
         if (flowAnalysisDisabled) {
             return errorType;
@@ -27994,6 +27996,9 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
                         continue;
                     }
                 }
+                else if (flags & FlowFlags.LambdaArgs) {
+                    type = getTypeAtFlowLambdaArgs(flow as FlowCall);
+                }
                 else if (flags & FlowFlags.ReduceLabel) {
                     const target = (flow as FlowReduceLabel).node.target;
                     const saveAntecedents = target.antecedent;
@@ -28005,7 +28010,7 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
                     // Check if we should continue with the control flow of the containing function.
                     const container = (flow as FlowStart).node;
                     if (
-                        container && container !== flowContainer &&
+                        container && container !== flowContainer && !inLambdaArg &&
                         reference.kind !== SyntaxKind.PropertyAccessExpression &&
                         reference.kind !== SyntaxKind.ElementAccessExpression &&
                         !(reference.kind === SyntaxKind.ThisKeyword && container.kind !== SyntaxKind.ArrowFunction)
@@ -28132,6 +28137,28 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
             return undefined;
         }
 
+        function getTypeAtFlowLambdaArgs(flow: FlowCall): FlowType {
+            const flowType = getTypeAtFlowNode(flow.antecedent);
+            const saveInitialType = initialType;
+            const saveInLambdaArg = inLambdaArg;
+            initialType = getTypeFromFlowType(flowType);
+            inLambdaArg = true;
+            let lambdaTypes: Type[] | undefined;
+            for (const arg of flow.node.arguments) {
+                const lambda = getLambdaArgument(arg);
+                if (lambda && lambda.returnFlowNode) {
+                    const lambdaType = getTypeFromFlowType(getTypeAtFlowNode(lambda.returnFlowNode));
+                    if (lambdaType !== initialType) {
+                        lambdaTypes ??= [initialType];
+                        lambdaTypes.push(lambdaType);
+                    }
+                }
+            }
+            inLambdaArg = saveInLambdaArg;
+            initialType = saveInitialType;
+            return lambdaTypes ? createFlowType(getUnionOrEvolvingArrayType(lambdaTypes, UnionReduction.Literal), isIncomplete(flowType)) : flowType;
+        }
+
         function getTypeAtFlowArrayMutation(flow: FlowArrayMutation): FlowType | undefined {
             if (declaredType === autoType || declaredType === autoArrayType) {
                 const node = flow.node;

src/compiler/types.ts

@@ -4111,8 +4111,9 @@ export const enum FlowFlags {
     ArrayMutation  = 1 << 8,  // Potential array mutation
     Call           = 1 << 9,  // Potential assertion call
     ReduceLabel    = 1 << 10, // Temporarily reduce antecedents of label
-    Referenced     = 1 << 11, // Referenced as antecedent once
-    Shared         = 1 << 12, // Referenced as antecedent more than once
+    LambdaArgs     = 1 << 11, // Call expression with lambda arguments
+    Referenced     = 1 << 12, // Referenced as antecedent once
+    Shared         = 1 << 13, // Referenced as antecedent more than once
 
     Label = BranchLabel | LoopLabel,
     Condition = TrueCondition | FalseCondition,

src/compiler/utilities.ts

@@ -11632,3 +11632,15 @@ export function hasInferredType(node: Node): node is HasInferredType {
             return false;
     }
 }
+
+/** @internal */
+export function getLambdaArgument(node: Node): FunctionLikeDeclaration | undefined {
+    switch (node.kind) {
+        case SyntaxKind.FunctionExpression:
+        case SyntaxKind.ArrowFunction:
+            return node as FunctionLikeDeclaration;
+        case SyntaxKind.ParenthesizedExpression:
+            return getLambdaArgument((node as ParenthesizedExpression).expression);
+    }
+    return undefined;
+}