Fixed bug in code flow engine that led to incorrect type evaluation o…

…f a variable in a nested loop. This addresses https://github.com/microsoft/pylance-release/issues/4509.

packages/pyright-internal/src/analyzer/codeFlowEngine.ts

@@ -606,7 +606,7 @@ export function getCodeFlowEngine(
                                     let isIncomplete = flowTypeResult.isIncomplete;
 
                                     if (flowType) {
-                                        const flowTypeResult = typeNarrowingCallback(flowType);
+                                        const flowTypeResult = typeNarrowingCallback(flowType, isIncomplete);
 
                                         if (flowTypeResult) {
                                             flowType = flowTypeResult.type;
@@ -667,7 +667,11 @@ export function getCodeFlowEngine(
                                             let narrowedType = refTypeInfo.type;
                                             let isIncomplete = !!refTypeInfo.isIncomplete;
 
-                                            const narrowedTypeResult = typeNarrowingCallback(refTypeInfo.type);
+                                            const narrowedTypeResult = typeNarrowingCallback(
+                                                refTypeInfo.type,
+                                                isIncomplete
+                                            );
+
                                             if (narrowedTypeResult) {
                                                 narrowedType = narrowedTypeResult.type;
                                                 if (narrowedTypeResult.isIncomplete) {
@@ -864,6 +868,7 @@ export function getCodeFlowEngine(
                         reference === undefined &&
                         cacheEntry.incompleteSubtypes?.some((subtype) => subtype.type !== undefined);
                     let firstAntecedentTypeIsIncomplete = false;
+                    let firstAntecedentTypeIsPending = false;
 
                     loopNode.antecedents.forEach((antecedent, index) => {
                         // If we've trying to determine reachability and we've already proven
@@ -872,6 +877,10 @@ export function getCodeFlowEngine(
                             return;
                         }
 
+                        if (firstAntecedentTypeIsPending && index > 0) {
+                            return;
+                        }
+
                         cacheEntry = getCacheEntry(loopNode)!;
 
                         // Is this entry marked "pending"? If so, we have recursed and there
@@ -883,9 +892,21 @@ export function getCodeFlowEngine(
                             index < cacheEntry.incompleteSubtypes.length &&
                             cacheEntry.incompleteSubtypes[index].isPending
                         ) {
-                            sawIncomplete = true;
-                            sawPending = true;
-                            return;
+                            // In rare circumstances, it's possible for a code flow graph with
+                            // nested loops to hit the case where the first antecedent is marked
+                            // as pending. In this case, we'll evaluate only the first antecedent
+                            // again even though it's pending. We're guaranteed to make forward
+                            // progress with the first antecedent, and that will allow us to establish
+                            // an initial type for this expression, but we don't want to evaluate
+                            // any other antecedents in this case because this could result in
+                            // infinite recursion.
+                            if (index === 0) {
+                                firstAntecedentTypeIsPending = true;
+                            } else {
+                                sawIncomplete = true;
+                                sawPending = true;
+                                return;
+                            }
                         }
 
                         // Have we already been here (i.e. does the entry exist and is
@@ -895,7 +916,11 @@ export function getCodeFlowEngine(
                             cacheEntry.incompleteSubtypes !== undefined && index < cacheEntry.incompleteSubtypes.length
                                 ? cacheEntry.incompleteSubtypes[index]
                                 : undefined;
-                        if (subtypeEntry === undefined || (!subtypeEntry?.isPending && subtypeEntry?.isIncomplete)) {
+                        if (
+                            subtypeEntry === undefined ||
+                            (!subtypeEntry?.isPending && subtypeEntry?.isIncomplete) ||
+                            index === 0
+                        ) {
                             const entryEvaluationCount = subtypeEntry === undefined ? 0 : subtypeEntry.evaluationCount;
 
                             // Set this entry to "pending" to prevent infinite recursion.

packages/pyright-internal/src/analyzer/typeGuards.ts

@@ -89,7 +89,7 @@ export interface TypeNarrowingResult {
     isIncomplete: boolean;
 }
 
-export type TypeNarrowingCallback = (type: Type) => TypeNarrowingResult | undefined;
+export type TypeNarrowingCallback = (type: Type, isIncomplete: boolean) => TypeNarrowingResult | undefined;
 
 // Given a reference expression and a test expression, returns a callback that
 // can be used to narrow the type described by the reference expression.
@@ -233,16 +233,17 @@ export function getTypeNarrowingCallback(
                         (ClassType.isEnumClass(rightType) || ClassType.isBuiltIn(rightType, 'bool')) &&
                         rightType.literalValue !== undefined
                     ) {
-                        return (type: Type) => {
+                        return (type: Type, isTypeIncomplete: boolean) => {
                             return {
                                 type: narrowTypeForLiteralComparison(
                                     evaluator,
                                     type,
+                                    isTypeIncomplete,
                                     rightType,
                                     adjIsPositiveTest,
                                     /* isIsOperator */ true
                                 ),
-                                isIncomplete: !!rightTypeResult.isIncomplete,
+                                isIncomplete: isTypeIncomplete || !!rightTypeResult.isIncomplete,
                             };
                         };
                     }
@@ -321,16 +322,17 @@ export function getTypeNarrowingCallback(
                     const rightType = rightTypeResult.type;
 
                     if (isClassInstance(rightType) && rightType.literalValue !== undefined) {
-                        return (type: Type) => {
+                        return (type: Type, isTypeIncomplete: boolean) => {
                             return {
                                 type: narrowTypeForLiteralComparison(
                                     evaluator,
                                     type,
+                                    isTypeIncomplete,
                                     rightType,
                                     adjIsPositiveTest,
                                     /* isIsOperator */ false
                                 ),
-                                isIncomplete: !!rightTypeResult.isIncomplete,
+                                isIncomplete: isTypeIncomplete || !!rightTypeResult.isIncomplete,
                             };
                         };
                     }
@@ -2053,10 +2055,17 @@ function narrowTypeForTypeIs(evaluator: TypeEvaluator, type: Type, classType: Cl
 function narrowTypeForLiteralComparison(
     evaluator: TypeEvaluator,
     referenceType: Type,
+    isTypeIncomplete: boolean,
     literalType: ClassType,
     isPositiveTest: boolean,
     isIsOperator: boolean
 ): Type {
+    // If the reference type is incomplete, don't attempt to narrow it. This can
+    // lead to incorrect analysis.
+    if (isTypeIncomplete) {
+        return referenceType;
+    }
+
     return mapSubtypes(referenceType, (subtype) => {
         subtype = evaluator.makeTopLevelTypeVarsConcrete(subtype);
         if (isClassInstance(subtype) && ClassType.isSameGenericClass(literalType, subtype)) {

packages/pyright-internal/src/tests/samples/loops38.py

@@ -0,0 +1,22 @@
+# This sample tests a code flow graph that includes a nested loop
+# and a variable that is assigned only in the outer loop.
+
+
+# pyright: strict
+
+
+# * Code flow graph for func1:
+# Assign[step+=1] ── True[step==0] ── Assign[_=] ── Loop ┬─ Loop ┬─ Assign[step=1] ── Start
+#                                                        │       ╰ Circular(Assign[step+=1])
+#                                                        ╰ FalseNever ─ False ─ Circular(Assign[_])
+def func1():
+    step = 1
+    while True:
+        for _ in [1]:
+            if step == 0:
+                step += 1
+                break
+        else:
+            return
+
+

packages/pyright-internal/src/tests/typeEvaluator3.test.ts

@@ -388,6 +388,12 @@ test('Loops37', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('Loops38', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['loops38.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('ForLoop1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['forLoop1.py']);