Added performance optimization for code flow analysis within loops.

microsoft · Aug 31, 2021 · ab0c1d6 · ab0c1d6
1 parent 6234fcd
commit ab0c1d6
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 3 deletions.
diff --git a/packages/pyright-internal/src/analyzer/binder.ts b/packages/pyright-internal/src/analyzer/binder.ts
@@ -90,6 +90,7 @@ import {
     FlowExhaustedMatch,
     FlowFlags,
     FlowLabel,
+    FlowLoopLabel,
     FlowNarrowForPattern,
     FlowNode,
     FlowPostContextManagerLabel,
@@ -2314,10 +2315,11 @@ export class Binder extends ParseTreeWalker {
     }
 
     private _createLoopLabel() {
-        const flowNode: FlowLabel = {
+        const flowNode: FlowLoopLabel = {
             flags: FlowFlags.LoopLabel,
             id: getUniqueFlowNodeId(),
             antecedents: [],
+            expressions: undefined,
         };
         return flowNode;
     }
@@ -2889,14 +2891,25 @@ export class Binder extends ParseTreeWalker {
         }
     }
 
-    private _bindLoopStatement(preLoopLabel: FlowLabel, postLoopLabel: FlowLabel, callback: () => void) {
+    private _bindLoopStatement(preLoopLabel: FlowLoopLabel, postLoopLabel: FlowLabel, callback: () => void) {
         const savedContinueTarget = this._currentContinueTarget;
         const savedBreakTarget = this._currentBreakTarget;
+        const savedExpressions = this._currentScopeCodeFlowExpressions;
         this._currentContinueTarget = preLoopLabel;
         this._currentBreakTarget = postLoopLabel;
+        this._currentScopeCodeFlowExpressions = new Set<string>();
 
         callback();
 
+        preLoopLabel.expressions = this._currentScopeCodeFlowExpressions;
+
+        if (savedExpressions) {
+            this._currentScopeCodeFlowExpressions.forEach((value) => {
+                savedExpressions.add(value);
+            });
+        }
+
+        this._currentScopeCodeFlowExpressions = savedExpressions;
         this._currentContinueTarget = savedContinueTarget;
         this._currentBreakTarget = savedBreakTarget;
     }

diff --git a/packages/pyright-internal/src/analyzer/codeFlow.ts b/packages/pyright-internal/src/analyzer/codeFlow.ts
@@ -71,6 +71,14 @@ export interface FlowLabel extends FlowNode {
     antecedents: FlowNode[];
 }
 
+export interface FlowLoopLabel extends FlowLabel {
+    // Set of all expressions that require code flow analysis
+    // through the loop to determine their types. If an expression
+    // is not within this map, loop analysis can be skipped and
+    // determined from the first antecedent only.
+    expressions: Set<string> | undefined;
+}
+
 // FlowAssignment represents a node that assigns a value.
 export interface FlowAssignment extends FlowNode {
     node: CodeFlowReferenceExpressionNode;

diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -93,6 +93,7 @@ import {
     FlowExhaustedMatch,
     FlowFlags,
     FlowLabel,
+    FlowLoopLabel,
     FlowNarrowForPattern,
     FlowNode,
     FlowPostContextManagerLabel,
@@ -17576,7 +17577,18 @@ export function createTypeEvaluator(
                     }
 
                     if (curFlowNode.flags & FlowFlags.LoopLabel) {
-                        const loopNode = curFlowNode as FlowLabel;
+                        const loopNode = curFlowNode as FlowLoopLabel;
+
+                        // Is the current symbol modified in any way within the
+                        // loop? If not, we can skip all processing within the loop
+                        // and assume that the type comes from the first antecedent,
+                        // which feeds the loop.
+                        if (referenceKey) {
+                            if (!loopNode.expressions || !loopNode.expressions.has(referenceKey)) {
+                                curFlowNode = loopNode.antecedents[0];
+                                continue;
+                            }
+                        }
 
                         let firstWasIncomplete = false;
                         let isFirstTimeInLoop = false;