Port pyright optimization that skips loop flow nodes that dont contai…

…n relevant references

Fixes #51525

This requied reworking how we compare references so we could equate them with a `set` - that change, in and of itself,
may actually also be perf positive (since we now get to cache the walk over the reference's AST), but both changes together
is somewhere between neutral and 4% gains in local testing, depending on the suite. I'll need to see if this bears out on
the much slower CI box, though, since 4% locally is < 0.1s, which is well in the realm of "OS scheduler variance" for some of these
tests on my box. Do note, unlike pyright, which in the linked issue's change, calculates the references used within a loop
during binding, we calculate it late on-demand and cache it in the checker. This is basically required for us, since we equate
references using the identity of the resolved symbol of the root node of the reference (which... may be a bit circular with
expression checking!).

src/compiler/binder.ts

@@ -1295,8 +1295,8 @@ function createBinder(): (file: SourceFile, options: CompilerOptions) => void {
         return initFlowNode({ flags: FlowFlags.BranchLabel, antecedents: undefined });
     }
 
-    function createLoopLabel(): FlowLabel {
-        return initFlowNode({ flags: FlowFlags.LoopLabel, antecedents: undefined });
+    function createLoopLabel(node: NonNullable<FlowLabel["node"]>): FlowLabel {
+        return initFlowNode({ flags: FlowFlags.LoopLabel, antecedents: undefined, node });
     }
 
     function createReduceLabel(target: FlowLabel, antecedents: FlowNode[], antecedent: FlowNode): FlowReduceLabel {
@@ -1445,7 +1445,7 @@ function createBinder(): (file: SourceFile, options: CompilerOptions) => void {
     }
 
     function bindWhileStatement(node: WhileStatement): void {
-        const preWhileLabel = setContinueTarget(node, createLoopLabel());
+        const preWhileLabel = setContinueTarget(node, createLoopLabel(node));
         const preBodyLabel = createBranchLabel();
         const postWhileLabel = createBranchLabel();
         addAntecedent(preWhileLabel, currentFlow);
@@ -1458,7 +1458,7 @@ function createBinder(): (file: SourceFile, options: CompilerOptions) => void {
     }
 
     function bindDoStatement(node: DoStatement): void {
-        const preDoLabel = createLoopLabel();
+        const preDoLabel = createLoopLabel(node);
         const preConditionLabel = setContinueTarget(node, createBranchLabel());
         const postDoLabel = createBranchLabel();
         addAntecedent(preDoLabel, currentFlow);
@@ -1471,7 +1471,7 @@ function createBinder(): (file: SourceFile, options: CompilerOptions) => void {
     }
 
     function bindForStatement(node: ForStatement): void {
-        const preLoopLabel = setContinueTarget(node, createLoopLabel());
+        const preLoopLabel = setContinueTarget(node, createLoopLabel(node));
         const preBodyLabel = createBranchLabel();
         const postLoopLabel = createBranchLabel();
         bind(node.initializer);
@@ -1486,7 +1486,7 @@ function createBinder(): (file: SourceFile, options: CompilerOptions) => void {
     }
 
     function bindForInOrForOfStatement(node: ForInOrOfStatement): void {
-        const preLoopLabel = setContinueTarget(node, createLoopLabel());
+        const preLoopLabel = setContinueTarget(node, createLoopLabel(node));
         const postLoopLabel = createBranchLabel();
         bind(node.expression);
         addAntecedent(preLoopLabel, currentFlow);

src/compiler/checker.ts

@@ -153,6 +153,7 @@ import {
     EmitResolver,
     EmitTextWriter,
     emptyArray,
+    emptySet,
     endsWith,
     EntityName,
     EntityNameExpression,
@@ -25375,47 +25376,101 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
     }
 
     function isMatchingReference(source: Node, target: Node): boolean {
-        switch (target.kind) {
-            case SyntaxKind.ParenthesizedExpression:
-            case SyntaxKind.NonNullExpression:
-                return isMatchingReference(source, (target as NonNullExpression | ParenthesizedExpression).expression);
-            case SyntaxKind.BinaryExpression:
-                return (isAssignmentExpression(target) && isMatchingReference(source, target.left)) ||
-                    (isBinaryExpression(target) && target.operatorToken.kind === SyntaxKind.CommaToken && isMatchingReference(source, target.right));
-        }
-        switch (source.kind) {
+        const keySource = getReferenceCacheKey(source);
+        const keyTarget = getReferenceCacheKey(target);
+        return keySource !== undefined && keyTarget !== undefined && keySource === keyTarget;
+    }
+
+    /**
+     * @returns A cache key that looks like `1234."field"."\"prop\".\"name\""` or `undefined` if the node can't be a reference
+     * `1234` can be any of `new`, `import`, `super`, or `this` if one of those special expressions is the root of the reference, rather than an identifier with a known symbol
+     */
+    function getReferenceCacheKeyWorker(reference: Node): string | undefined {
+        switch (reference.kind) {
             case SyntaxKind.MetaProperty:
-                return target.kind === SyntaxKind.MetaProperty
-                    && (source as MetaProperty).keywordToken === (target as MetaProperty).keywordToken
-                    && (source as MetaProperty).name.escapedText === (target as MetaProperty).name.escapedText;
+                return `${(reference as MetaProperty).keywordToken === SyntaxKind.NewKeyword ? "new" : "import"}."${escapeString(idText((reference as MetaProperty).name))}"`;
             case SyntaxKind.Identifier:
-            case SyntaxKind.PrivateIdentifier:
-                return isThisInTypeQuery(source) ?
-                    target.kind === SyntaxKind.ThisKeyword :
-                    target.kind === SyntaxKind.Identifier && getResolvedSymbol(source as Identifier) === getResolvedSymbol(target as Identifier) ||
-                        (isVariableDeclaration(target) || isBindingElement(target)) &&
-                        getExportSymbolOfValueSymbolIfExported(getResolvedSymbol(source as Identifier)) === getSymbolOfDeclaration(target);
+            case SyntaxKind.PrivateIdentifier: {
+                if (isThisInTypeQuery(reference)) {
+                    return "this";
+                }
+                // Sometimes this identifier is the RHS of a property name expression - in those cases, the symbol is only known if expression checking on the property access expression is already
+                // complete (or complete enough to have assigned the resolved symbol)
+                // (it will also trigger undesirable behavior to call `getResolvedSymbol` on such an identifier, since it'll try to lookup a nonsense name!)
+                const sym = getNodeLinks(reference as Identifier).resolvedSymbol || isIdentifier(reference) && !(isPropertyAccessExpression(reference.parent) && reference.parent.name === reference) && getResolvedSymbol(reference);
+                const exportSymbol = sym && getExportSymbolOfValueSymbolIfExported(sym);
+                return exportSymbol ? `${getSymbolId(exportSymbol)}` : undefined;
+            }
+            case SyntaxKind.VariableDeclaration:
+            case SyntaxKind.BindingElement: {
+                const declSymbol = getSymbolOfDeclaration(reference as VariableDeclaration | BindingElement);
+                const exportSymbol = declSymbol && getExportSymbolOfValueSymbolIfExported(declSymbol);
+                return `${getSymbolId(exportSymbol)}`;
+            }
             case SyntaxKind.ThisKeyword:
-                return target.kind === SyntaxKind.ThisKeyword;
+                return "this";
             case SyntaxKind.SuperKeyword:
-                return target.kind === SyntaxKind.SuperKeyword;
+                return "super";
             case SyntaxKind.NonNullExpression:
             case SyntaxKind.ParenthesizedExpression:
-                return isMatchingReference((source as NonNullExpression | ParenthesizedExpression).expression, target);
+                return getReferenceCacheKey((reference as NonNullExpression | ParenthesizedExpression).expression);
             case SyntaxKind.PropertyAccessExpression:
             case SyntaxKind.ElementAccessExpression:
-                const sourcePropertyName = getAccessedPropertyName(source as AccessExpression);
-                const targetPropertyName = isAccessExpression(target) ? getAccessedPropertyName(target) : undefined;
-                return sourcePropertyName !== undefined && targetPropertyName !== undefined && targetPropertyName === sourcePropertyName &&
-                    isMatchingReference((source as AccessExpression).expression, (target as AccessExpression).expression);
+                const propertyName = getAccessedPropertyName(reference as AccessExpression);
+                const lhsName = getReferenceCacheKey((reference as AccessExpression).expression);
+                return propertyName !== undefined && lhsName !== undefined ? `${lhsName}."${escapeString(unescapeLeadingUnderscores(propertyName))}"` : undefined;
             case SyntaxKind.QualifiedName:
-                return isAccessExpression(target) &&
-                    (source as QualifiedName).right.escapedText === getAccessedPropertyName(target) &&
-                    isMatchingReference((source as QualifiedName).left, target.expression);
-            case SyntaxKind.BinaryExpression:
-                return (isBinaryExpression(source) && source.operatorToken.kind === SyntaxKind.CommaToken && isMatchingReference(source.right, target));
+                return `${getReferenceCacheKey((reference as QualifiedName).left)!}."${escapeString(idText((reference as QualifiedName).right))}"`;
+            case SyntaxKind.BinaryExpression: {
+                Debug.assert(isBinaryExpression(reference));
+                return reference.operatorToken.kind === SyntaxKind.CommaToken ? getReferenceCacheKey(reference.right) :
+                    isAssignmentExpression(reference) ? getReferenceCacheKey(reference.left) :
+                    undefined;
+            }
+        }
+        return undefined;
+    }
+
+    function getReferenceCacheKey(node: Node) {
+        const links = getNodeLinks(node);
+        if (!links.referenceCacheKey) {
+            const result = getReferenceCacheKeyWorker(node);
+            links.referenceCacheKey = result !== undefined ? result : "<none>";
+        }
+        return links.referenceCacheKey === "<none>" ? undefined : links.referenceCacheKey;
+    }
+
+    function getContainedReferences(node: Node): ReadonlySet<string> {
+        const links = getNodeLinks(node);
+        if (!links.containedReferences) {
+            links.containedReferences = getContainedReferencesWorker(node);
+        }
+        return links.containedReferences;
+    }
+
+    function getContainedReferencesWorker(node: Node): ReadonlySet<string> {
+        let references: ReadonlySet<never> | Set<string> = emptySet;
+        if (isStatement(node) || isExpressionNode(node)) {
+            // don't descend into variable declarations or binding elements, since those contain identifiers we don't want to collect
+            forEachChild(node, visitor);
+            const key = getReferenceCacheKey(node);
+            if (key !== undefined) {
+                addReference(key);
+            }
+        }
+        return references;
+
+        function visitor(n: Node) {
+            const refs = getContainedReferences(n);
+            refs.forEach(r => addReference(r));
+        }
+
+        function addReference(ref: string) {
+            if (references === emptySet) {
+                references = new Set<string>();
+            }
+            (references as Set<String>).add(ref);
         }
-        return false;
     }
 
     function getAccessedPropertyName(access: AccessExpression | BindingElement | ParameterDeclaration): __String | undefined {
@@ -26866,6 +26921,7 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
             const antecedentTypes: Type[] = [];
             let subtypeReduction = false;
             let firstAntecedentType: FlowType | undefined;
+            let possiblyNarrowsRef: boolean | undefined;
             for (const antecedent of flow.antecedents!) {
                 let flowType;
                 if (!firstAntecedentType) {
@@ -26874,6 +26930,15 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
                     flowType = firstAntecedentType = getTypeAtFlowNode(antecedent);
                 }
                 else {
+                    if (possiblyNarrowsRef === undefined) {
+                        const key = getReferenceCacheKey(reference);
+                        if (key !== undefined) {
+                            possiblyNarrowsRef = getContainedReferences(flow.node!).has(key);
+                        }
+                    }
+                    if (!possiblyNarrowsRef) {
+                        break;
+                    }
                     // All but the first antecedent are the looping control flow paths that lead
                     // back to the loop junction. We track these on the flow loop stack.
                     flowLoopNodes[flowLoopCount] = flow;

src/compiler/types.ts

@@ -4152,6 +4152,7 @@ export interface FlowStart extends FlowNodeBase {
 // FlowLabel represents a junction with multiple possible preceding control flows.
 export interface FlowLabel extends FlowNodeBase {
     antecedents: FlowNode[] | undefined;
+    node?: ForInOrOfStatement | ForStatement | WhileStatement | DoStatement;
 }
 
 // FlowAssignment represents a node that assigns a value to a narrowable reference,
@@ -6061,6 +6062,8 @@ export interface NodeLinks {
     parameterInitializerContainsUndefined?: boolean; // True if this is a parameter declaration whose type annotation contains "undefined".
     fakeScopeForSignatureDeclaration?: boolean; // True if this is a fake scope injected into an enclosing declaration chain.
     assertionExpressionType?: Type;     // Cached type of the expression of a type assertion
+    referenceCacheKey?: string;         // Cached reference equivalence cache key
+    containedReferences?: ReadonlySet<string>;  // Set of references contained within the node, including itself
 }
 
 /** @internal */

tests/baselines/reference/api/tsserverlibrary.d.ts

@@ -6088,6 +6088,7 @@ declare namespace ts {
     }
     interface FlowLabel extends FlowNodeBase {
         antecedents: FlowNode[] | undefined;
+        node?: ForInOrOfStatement | ForStatement | WhileStatement | DoStatement;
     }
     interface FlowAssignment extends FlowNodeBase {
         node: Expression | VariableDeclaration | BindingElement;

tests/baselines/reference/api/typescript.d.ts

@@ -2035,6 +2035,7 @@ declare namespace ts {
     }
     interface FlowLabel extends FlowNodeBase {
         antecedents: FlowNode[] | undefined;
+        node?: ForInOrOfStatement | ForStatement | WhileStatement | DoStatement;
     }
     interface FlowAssignment extends FlowNodeBase {
         node: Expression | VariableDeclaration | BindingElement;