From 0523c63c6c8dc01b18f83d832aedebfe2352dcbd Mon Sep 17 00:00:00 2001
From: Anders Hejlsberg <andersh@microsoft.com>
Date: Fri, 7 Jun 2024 13:55:23 -0700
Subject: [PATCH] Optimize lambda assignment analysis using isMutableFlowNode
 walk

---
 src/compiler/binder.ts  |  2 +-
 src/compiler/checker.ts | 70 +++++++++++++++++++++++++++++++++++++----
 2 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/src/compiler/binder.ts b/src/compiler/binder.ts
index 0a9fcdbb464..4a033cef52d 100644
--- a/src/compiler/binder.ts
+++ b/src/compiler/binder.ts
@@ -1050,7 +1050,7 @@ function createBinder(): (file: SourceFile, options: CompilerOptions) => void {
             if (currentReturnTarget) {
                 addAntecedent(currentReturnTarget, currentFlow);
                 currentFlow = finishFlowLabel(currentReturnTarget);
-                if (isLambdaArgument || node.kind === SyntaxKind.Constructor || node.kind === SyntaxKind.ClassStaticBlockDeclaration || (isInJSFile(node) && (node.kind === SyntaxKind.FunctionDeclaration || node.kind === SyntaxKind.FunctionExpression))) {
+                if (isLambdaArgument && hasFlowMutation || node.kind === SyntaxKind.Constructor || node.kind === SyntaxKind.ClassStaticBlockDeclaration || (isInJSFile(node) && (node.kind === SyntaxKind.FunctionDeclaration || node.kind === SyntaxKind.FunctionExpression))) {
                     (node as FunctionLikeDeclaration | ClassStaticBlockDeclaration).returnFlowNode = currentFlow;
                 }
             }
diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
index 844a6f63a31..389a6c77835 100644
--- a/src/compiler/checker.ts
+++ b/src/compiler/checker.ts
@@ -27928,6 +27928,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 hasMutation: boolean[] | undefined;
         let isKeySet = false;
         let inLambdaArg = false;
         let flowDepth = 0;
@@ -28163,21 +28164,24 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
         }
 
         function getTypeAtFlowLambdaArgs(flow: FlowCall): FlowType {
-            const signatures = getSignaturesOfType(getTypeOfExpression(flow.node.expression), SignatureKind.Call);
             const flowType = getTypeAtFlowNode(flow.antecedent);
             const saveInitialType = initialType;
             const saveInLambdaArg = inLambdaArg;
             initialType = getTypeFromFlowType(flowType);
             inLambdaArg = true;
             let lambdaTypes: Type[] | undefined;
+            let signatures: readonly Signature[] | undefined;
             const args = flow.node.arguments;
             for (let i = 0; i < args.length; i++) {
                 const lambda = getLambdaArgument(args[i]);
-                if (lambda && lambda.returnFlowNode && !some(signatures, sig => !!(getObjectFlags(getTypeAtPosition(sig, i)) & ObjectFlags.DeferredCallback))) {
-                    const lambdaType = getTypeFromFlowType(getTypeAtFlowNode(lambda.returnFlowNode));
-                    if (lambdaType !== initialType) {
-                        lambdaTypes ??= [initialType];
-                        lambdaTypes.push(lambdaType);
+                if (lambda && lambda.returnFlowNode && isMutationFlowNode(lambda.returnFlowNode, /*noCacheCheck*/ false)) {
+                    signatures ??= getSignaturesOfType(getTypeOfExpression(flow.node.expression), SignatureKind.Call);
+                    if (!some(signatures, sig => !!(getObjectFlags(getTypeAtPosition(sig, i)) & ObjectFlags.DeferredCallback))) {
+                        const lambdaType = getTypeFromFlowType(getTypeAtFlowNode(lambda.returnFlowNode));
+                        if (lambdaType !== initialType) {
+                            lambdaTypes ??= [initialType];
+                            lambdaTypes.push(lambdaType);
+                        }
                     }
                 }
             }
@@ -28186,6 +28190,60 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
             return lambdaTypes ? createFlowType(getUnionOrEvolvingArrayType(lambdaTypes, UnionReduction.Literal), isIncomplete(flowType)) : flowType;
         }
 
+        function isMutationFlowNode(flow: FlowNode, noCacheCheck: boolean): boolean {
+            while (true) {
+                const flags = flow.flags;
+                if (flags & FlowFlags.Shared) {
+                    if (!noCacheCheck) {
+                        const id = getFlowNodeId(flow);
+                        const cached = (hasMutation ??= [])[id];
+                        return cached !== undefined ? cached : (hasMutation[id] = isMutationFlowNode(flow, /*noCacheCheck*/ true));
+                    }
+                    noCacheCheck = false;
+                }
+                if (flags & (FlowFlags.Call | FlowFlags.Condition | FlowFlags.SwitchClause | FlowFlags.ReduceLabel)) {
+                    flow = (flow as FlowCall | FlowCondition | FlowSwitchClause | FlowReduceLabel).antecedent;
+                }
+                else if (flags & FlowFlags.Assignment) {
+                    if (isOrContainsMatchingReference(reference, (flow as FlowArrayMutation).node)) {
+                        return true;
+                    }
+                    flow = (flow as FlowAssignment).antecedent;
+                }
+                else if (flags & FlowFlags.BranchLabel) {
+                    return some((flow as FlowLabel).antecedent, f => isMutationFlowNode(f, /*noCacheCheck*/ false));
+                }
+                else if (flags & FlowFlags.LoopLabel) {
+                    const id = getFlowNodeId(flow);
+                    const cached = (hasMutation ??= [])[id];
+                    if (cached !== undefined) return cached;
+                    hasMutation[id] = false;
+                    return hasMutation[id] = some((flow as FlowLabel).antecedent, f => isMutationFlowNode(f, /*noCacheCheck*/ false));
+                }
+                else if (flags & FlowFlags.ArrayMutation) {
+                    if (declaredType === autoType || declaredType === autoArrayType) {
+                        const node = (flow as FlowArrayMutation).node;
+                        const expr = node.kind === SyntaxKind.CallExpression ?
+                            (node.expression as PropertyAccessExpression).expression :
+                            (node.left as ElementAccessExpression).expression;
+                        if (isMatchingReference(reference, getReferenceCandidate(expr))) {
+                            return true;
+                        }
+                    }
+                    flow = (flow as FlowArrayMutation).antecedent;
+                }
+                else if (flags & FlowFlags.LambdaArgs) {
+                    return some((flow as FlowCall).node.arguments, arg => {
+                        const lambda = getLambdaArgument(arg);
+                        return !!(lambda && lambda.returnFlowNode && isMutationFlowNode(lambda.returnFlowNode, /*noCacheCheck*/ false));
+                    });
+                }
+                else {
+                    return false;
+                }
+            }
+        }
+
         function getTypeAtFlowArrayMutation(flow: FlowArrayMutation): FlowType | undefined {
             if (declaredType === autoType || declaredType === autoArrayType) {
                 const node = flow.node;