From d7b180e15d1cea1a6cf5c7a7c87f4d6c6fb3151d Mon Sep 17 00:00:00 2001
From: "evgeny.bovykin" <evgeny.bovykin@jetbrains.com>
Date: Wed, 10 Dec 2025 14:57:32 +0100
Subject: [PATCH] PY-86219 Use weak keys for TypeEvalContext caches by default

Space-RevId: 5bd8593af2f394e5dffaf8d120aad81c1735ed22

GitOrigin-RevId: e61a25688688e216eeddfd0e2212a6a9796cf2f7
---
 .../python/psi/types/TypeEvalContext.java     | 23 +++++++++++--------
 .../resources/intellij.python.psi.impl.xml    |  4 +++-
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java
index ed6c91398446..89ca20d0e660 100644
--- a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java
+++ b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java
@@ -14,6 +14,7 @@ import com.intellij.psi.PsiFile;
 import com.intellij.util.ArrayUtil;
 import com.intellij.util.ProcessingContext;
 import com.intellij.util.containers.CollectionFactory;
+import com.intellij.util.containers.HashingStrategy;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyTypeProvider;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
@@ -60,11 +61,14 @@ public sealed class TypeEvalContext {
   protected final Map<Pair<PyExpression, Object>, PyType> contextTypeCache = getConcurrentMapForCaching();
 
   private static <T> @NotNull ConcurrentMap<@NotNull T, @NotNull PyType> getConcurrentMapForCaching() {
-    if (Registry.is("python.typing.soft.keys.type.eval.context")) {
-      // In the current implementation, this value is only used to initialize the map and is basically ignored
-      // Just in case, set it to a reasonable value
-      // `Runtime.availableProcessors` shouldn't be called here, as that is a potentially expensive operation
-      int concurrencyLevel = 4;
+    // In the current implementation, this value is only used to initialize the map and is basically ignored
+    // Just in case, set it to a reasonable value
+    // `Runtime.availableProcessors` shouldn't be called here, as that is a potentially expensive operation
+    int concurrencyLevel = 4;
+    if (Registry.is("python.typing.weak.keys.type.eval.context")) {
+      return CollectionFactory.createConcurrentWeakKeySoftValueMap(10, 0.75f, concurrencyLevel, HashingStrategy.canonical());
+    }
+    else if (Registry.is("python.typing.soft.keys.type.eval.context")) {
       return CollectionFactory.createConcurrentSoftKeySoftValueMap(10, 0.75f, concurrencyLevel);
     }
     else {
@@ -273,13 +277,14 @@ public sealed class TypeEvalContext {
   }
 
   private @NotNull TypeEvalContext getLibraryContext(@NotNull Project project) {
+    // code completion will always have a new PsiFile, use the original file instead
+    PsiFile origin = myConstraints.myOrigin != null
+                     ? myConstraints.myOrigin.getOriginalFile()
+                     : null;
     TypeEvalConstraints constraints = new TypeEvalConstraints(myConstraints.myAllowDataFlow,
                                                               myConstraints.myAllowStubToAST,
                                                               myConstraints.myAllowCallContext,
-                                                              // code completion will always have a new PsiFile, use original file instead
-                                                              myConstraints.myOrigin != null
-                                                              ? myConstraints.myOrigin.getOriginalFile()
-                                                              : null);
+                                                              origin);
     return project.getService(TypeEvalContextCache.class).getLibraryContext(new LibraryTypeEvalContext(constraints));
   }
 
diff --git a/python/python-psi-impl/resources/intellij.python.psi.impl.xml b/python/python-psi-impl/resources/intellij.python.psi.impl.xml
index 4b53abe84d31..b7c38eaebf72 100644
--- a/python/python-psi-impl/resources/intellij.python.psi.impl.xml
+++ b/python/python-psi-impl/resources/intellij.python.psi.impl.xml
@@ -560,8 +560,10 @@
                  description="Allows narrowing types exhaustively to Never"/>
     <registryKey key="python.lsp.tools.enabled" defaultValue="true" restartRequired="true"
                  description="Enables lsp tools"/>
-    <registryKey key="python.typing.soft.keys.type.eval.context" defaultValue="true"
+    <registryKey key="python.typing.soft.keys.type.eval.context" defaultValue="false"
                  description="Enable soft-key caching for TypeEvalContext"/>
+    <registryKey key="python.typing.weak.keys.type.eval.context" defaultValue="true"
+                 description="Enable weak-key caching for TypeEvalContext"/>
 
     <virtualFileCustomDataProvider implementation="com.jetbrains.python.formatter.PyCodeStyleSettingsProvider"/>
   </extensions>