From 23b4668508e0e31ea5ce09e2531b996e8919227b Mon Sep 17 00:00:00 2001
From: Eric Traut <eric@traut.com>
Date: Mon, 7 Oct 2024 11:37:07 -0700
Subject: [PATCH] Backed out recent performance optimization that resulted in a
 regression for protocol matching in certain circumstances. This addresses
 #9153, #9132, and #9130. (#9164)

---
 .../src/analyzer/constraintTracker.ts         |  8 ++++++
 .../src/analyzer/protocols.ts                 | 28 +++++++++++++++++--
 .../src/tests/samples/protocol51.py           |  6 ++++
 .../src/tests/typeEvaluator7.test.ts          |  6 ++++
 4 files changed, 45 insertions(+), 3 deletions(-)
 create mode 100644 packages/pyright-internal/src/tests/samples/protocol51.py

diff --git a/packages/pyright-internal/src/analyzer/constraintTracker.ts b/packages/pyright-internal/src/analyzer/constraintTracker.ts
index 55fd3a8382f8..a0b2bd9b8b27 100644
--- a/packages/pyright-internal/src/analyzer/constraintTracker.ts
+++ b/packages/pyright-internal/src/analyzer/constraintTracker.ts
@@ -237,6 +237,14 @@ export class ConstraintTracker {
         }
     }
 
+    isSame(other: ConstraintTracker) {
+        if (other._constraintSets.length !== this._constraintSets.length) {
+            return false;
+        }
+
+        return this._constraintSets.every((set, index) => set.isSame(other._constraintSets[index]));
+    }
+
     isEmpty() {
         return this._constraintSets.every((set) => set.isEmpty());
     }
diff --git a/packages/pyright-internal/src/analyzer/protocols.ts b/packages/pyright-internal/src/analyzer/protocols.ts
index 833039ddec58..1a7b93b24627 100644
--- a/packages/pyright-internal/src/analyzer/protocols.ts
+++ b/packages/pyright-internal/src/analyzer/protocols.ts
@@ -58,6 +58,7 @@ interface ProtocolCompatibility {
     srcType: Type;
     destType: Type;
     flags: AssignTypeFlags;
+    constraints: ConstraintTracker | undefined;
     isCompatible: boolean;
 }
 
@@ -125,6 +126,7 @@ export function assignClassToProtocol(
 
     protocolAssignmentStack.push({ srcType, destType });
     let isCompatible = true;
+    const clonedConstraints = constraints?.clone();
 
     try {
         isCompatible = assignToProtocolInternal(evaluator, destType, srcType, diag, constraints, flags, recursionCount);
@@ -138,7 +140,7 @@ export function assignClassToProtocol(
     protocolAssignmentStack.pop();
 
     // Cache the results for next time.
-    setProtocolCompatibility(destType, srcType, flags, isCompatible);
+    setProtocolCompatibility(destType, srcType, flags, clonedConstraints, isCompatible);
 
     return isCompatible;
 }
@@ -229,7 +231,12 @@ function getProtocolCompatibility(
     }
 
     const entry = entries.find((entry) => {
-        return isTypeSame(entry.destType, destType) && isTypeSame(entry.srcType, srcType) && entry.flags === flags;
+        return (
+            isTypeSame(entry.destType, destType) &&
+            isTypeSame(entry.srcType, srcType) &&
+            entry.flags === flags &&
+            isConstraintTrackerSame(constraints, entry.constraints)
+        );
     });
 
     return entry?.isCompatible;
@@ -239,6 +246,7 @@ function setProtocolCompatibility(
     destType: ClassType,
     srcType: ClassType,
     flags: AssignTypeFlags,
+    constraints: ConstraintTracker | undefined,
     isCompatible: boolean
 ) {
     let map = srcType.shared.protocolCompatibility as Map<string, ProtocolCompatibility[]> | undefined;
@@ -253,13 +261,27 @@ function setProtocolCompatibility(
         map.set(destType.shared.fullName, entries);
     }
 
-    entries.push({ destType, srcType, flags, isCompatible });
+    entries.push({
+        destType,
+        srcType,
+        flags,
+        constraints: constraints,
+        isCompatible,
+    });
 
     if (entries.length > maxProtocolCompatibilityCacheEntries) {
         entries.shift();
     }
 }
 
+function isConstraintTrackerSame(context1: ConstraintTracker | undefined, context2: ConstraintTracker | undefined) {
+    if (!context1 || !context2) {
+        return context1 === context2;
+    }
+
+    return context1.isSame(context2);
+}
+
 function assignToProtocolInternal(
     evaluator: TypeEvaluator,
     destType: ClassType,
diff --git a/packages/pyright-internal/src/tests/samples/protocol51.py b/packages/pyright-internal/src/tests/samples/protocol51.py
new file mode 100644
index 000000000000..f7acb6ef71e2
--- /dev/null
+++ b/packages/pyright-internal/src/tests/samples/protocol51.py
@@ -0,0 +1,6 @@
+# This sample tests a regression case related to the protocol compatibility
+# cache.
+
+
+int(round(1.2, 0))
+round(3.4, 2)
diff --git a/packages/pyright-internal/src/tests/typeEvaluator7.test.ts b/packages/pyright-internal/src/tests/typeEvaluator7.test.ts
index d28134584cc5..6e0f12c02b47 100644
--- a/packages/pyright-internal/src/tests/typeEvaluator7.test.ts
+++ b/packages/pyright-internal/src/tests/typeEvaluator7.test.ts
@@ -593,6 +593,12 @@ test('Protocol50', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('Protocol51', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['protocol51.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('ProtocolExplicit1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['protocolExplicit1.py']);