Fixed a bug that led to poor performance (effectively a hang) if an unannotated function is called through many (greater than 8) call expressions within a loop using different literal-value arguments each time. This addresses #7791. (#7792)

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

@@ -174,6 +174,7 @@ import {
     CallResult,
     CallSignature,
     CallSignatureInfo,
+    CallSiteEvaluationInfo,
     ClassMemberLookup,
     ClassTypeResult,
     DeclaredSymbolTypeInfo,
@@ -11484,7 +11485,7 @@ export function createTypeEvaluator(
         }
 
         // Calculate the return type.
-        let returnType = getFunctionEffectiveReturnType(type, matchResults.argParams);
+        let returnType = getFunctionEffectiveReturnType(type, { args: matchResults.argParams, errorNode });
 
         if (condition.length > 0) {
             returnType = TypeBase.cloneForCondition(returnType, condition);
@@ -21828,7 +21829,7 @@ export function createTypeEvaluator(
     // into account argument types to infer the return type.
     function getFunctionEffectiveReturnType(
         type: FunctionType,
-        args?: ValidateArgTypeParams[],
+        callSiteInfo?: CallSiteEvaluationInfo,
         inferTypeIfNeeded = true
     ) {
         const specializedReturnType = FunctionType.getSpecializedReturnType(type, /* includeInferred */ false);
@@ -21837,13 +21838,13 @@ export function createTypeEvaluator(
         }
 
         if (inferTypeIfNeeded) {
-            return getFunctionInferredReturnType(type, args);
+            return getFunctionInferredReturnType(type, callSiteInfo);
         }
 
         return UnknownType.create();
     }
 
-    function _getFunctionInferredReturnType(type: FunctionType, args?: ValidateArgTypeParams[]) {
+    function _getFunctionInferredReturnType(type: FunctionType, callSiteInfo?: CallSiteEvaluationInfo) {
         let returnType: Type | undefined;
         let isIncomplete = false;
         const analyzeUnannotatedFunctions = true;
@@ -21924,7 +21925,7 @@ export function createTypeEvaluator(
             FunctionType.hasUnannotatedParams(type) &&
             !FunctionType.isStubDefinition(type) &&
             !FunctionType.isPyTypedDefinition(type) &&
-            args
+            callSiteInfo
         ) {
             let hasDecorators = false;
             let isAsync = false;
@@ -21941,7 +21942,7 @@ export function createTypeEvaluator(
             // We can't use this technique if decorators or async are used because they
             // would need to be applied to the inferred return type.
             if (!hasDecorators && !isAsync) {
-                const contextualReturnType = getFunctionInferredReturnTypeUsingArguments(type, args);
+                const contextualReturnType = getFunctionInferredReturnTypeUsingArguments(type, callSiteInfo);
                 if (contextualReturnType) {
                     returnType = contextualReturnType;
                 }
@@ -21953,8 +21954,9 @@ export function createTypeEvaluator(
 
     function getFunctionInferredReturnTypeUsingArguments(
         type: FunctionType,
-        args: ValidateArgTypeParams[]
+        callSiteInfo: CallSiteEvaluationInfo
     ): Type | undefined {
+        const args = callSiteInfo.args;
         let contextualReturnType: Type | undefined;
 
         if (!type.details.declaration) {
@@ -22001,6 +22003,10 @@ export function createTypeEvaluator(
         const paramTypes: Type[] = [];
         let isResultFromCache = false;
 
+        // If the call is located in a loop, don't use literal argument types
+        // for the same reason we don't do literal math in loops.
+        const stripLiteralArgTypes = ParseTreeUtils.isWithinLoop(callSiteInfo.errorNode);
+
         // Suppress diagnostics because we don't want to generate errors.
         suppressDiagnostics(functionNode, () => {
             // Allocate a new temporary type cache for the context of just
@@ -22050,6 +22056,10 @@ export function createTypeEvaluator(
                             paramType = UnknownType.create();
                         }
 
+                        if (stripLiteralArgTypes) {
+                            paramType = stripLiteralValue(paramType);
+                        }
+
                         paramTypes.push(paramType);
                         writeTypeCache(param.name, { type: paramType }, EvaluatorFlags.None);
                     }

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

@@ -458,6 +458,11 @@ export interface MapSubtypesOptions {
     expandCallback?: (type: Type) => Type;
 }
 
+export interface CallSiteEvaluationInfo {
+    errorNode: ExpressionNode;
+    args: ValidateArgTypeParams[];
+}
+
 export interface TypeEvaluator {
     runWithCancellationToken<T>(token: CancellationToken, callback: () => T): T;
 
@@ -548,7 +553,7 @@ export interface TypeEvaluator {
     getInferredTypeOfDeclaration: (symbol: Symbol, decl: Declaration) => Type | undefined;
     getDeclaredTypeForExpression: (expression: ExpressionNode, usage?: EvaluatorUsage) => Type | undefined;
     getFunctionDeclaredReturnType: (node: FunctionNode) => Type | undefined;
-    getFunctionInferredReturnType: (type: FunctionType, args?: ValidateArgTypeParams[]) => Type;
+    getFunctionInferredReturnType: (type: FunctionType, callSiteInfo?: CallSiteEvaluationInfo) => Type;
     getBestOverloadForArguments: (
         errorNode: ExpressionNode,
         typeResult: TypeResult<OverloadedFunctionType>,

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

@@ -0,0 +1,50 @@
+# This sample tests the case where a call-site return type evaluation
+# is invoked multiple times within a loop using different literal values
+# each time.
+
+
+def func1(h, ids):
+    for _ in ids:
+        h = func2(h, 1)
+        h = func2(h, 2)
+        h = func2(h, 3)
+        h = func2(h, 4)
+        h = func2(h, 5)
+        h = func2(h, 6)
+        h = func2(h, 7)
+        h = func2(h, 8)
+        h = func2(h, 9)
+        h = func2(h, 10)
+        h = func2(h, 11)
+        h = func2(h, 12)
+        h = func2(h, 13)
+        h = func2(h, 14)
+        h = func2(h, 15)
+        h = func2(h, 16)
+        h = func2(h, 17)
+        h = func2(h, 18)
+        h = func2(h, 19)
+        h = func2(h, 20)
+        h = func2(h, 21)
+        h = func2(h, 22)
+        h = func2(h, 23)
+        h = func2(h, 24)
+        h = func2(h, 25)
+        h = func2(h, 26)
+        h = func2(h, 27)
+        h = func2(h, 28)
+        h = func2(h, 29)
+        h = func2(h, 30)
+        h = func2(h, 31)
+        h = func2(h, 32)
+        h = func2(h, 33)
+        h = func2(h, 34)
+        h = func2(h, 35)
+        h = func2(h, 36)
+        h = func2(h, 37)
+        h = func2(h, 38)
+        h = func2(h, 39)
+
+
+def func2(a, unused):
+    return a

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

@@ -412,6 +412,11 @@ test('CallSite2', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('CallSite3', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['callSite3.py']);
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('FString1', () => {
     const configOptions = new ConfigOptions(Uri.empty());