From 6a772de5d5d7d195009608b003e80b51857e47f9 Mon Sep 17 00:00:00 2001
From: "evgeny.bovykin" <evgeny.bovykin@jetbrains.com>
Date: Thu, 27 Nov 2025 15:23:37 +0100
Subject: [PATCH] PY-85638 PY-85893 Use soft-key soft-value map for
 TypeEvalContext caches

This prevents the caches from blowing up in size when analyzing a lot of files in a row

The most obvious example is running `Analyze code` on the whole project. Previously, all PSI elements from all files would pollute the cache, potentially consuming multiple Gb of memory. See PyDjangoSourcesPerformanceTest as an example. It used to require 16 Gb of heap, and it would still sometimes fail with OOM. Now, it runs smoothly with 2Gb of heap

GitOrigin-RevId: f9854f8abdd7086414df8f610b0f2d611ca0d455
---
 .../python/psi/types/TypeEvalContext.java     | 22 +++++++++++++++----
 .../resources/intellij.python.psi.impl.xml    |  2 ++
 2 files changed, 20 insertions(+), 4 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 cea68026e2fa..441b3ae3b12e 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
@@ -27,6 +27,7 @@ import org.jetbrains.annotations.Nullable;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentMap;
 import java.util.function.Function;
 
 
@@ -52,10 +53,23 @@ public sealed class TypeEvalContext {
 
   private final ThreadLocal<ProcessingContext> myProcessingContext = ThreadLocal.withInitial(ProcessingContext::new);
 
-  protected final Map<PyTypedElement, PyType> myEvaluated = CollectionFactory.createConcurrentSoftValueMap();
-  protected final Map<PyTypedElement, PyType> myExternalEvaluated = CollectionFactory.createConcurrentSoftValueMap();
-  protected final Map<PyCallable, PyType> myEvaluatedReturn = CollectionFactory.createConcurrentSoftValueMap();
-  protected final Map<Pair<PyExpression, Object>, PyType> contextTypeCache = CollectionFactory.createConcurrentSoftValueMap();
+  protected final Map<PyTypedElement, PyType> myEvaluated = getConcurrentMapForCaching();
+  protected final Map<PyTypedElement, PyType> myExternalEvaluated = getConcurrentMapForCaching();
+  protected final Map<PyCallable, PyType> myEvaluatedReturn = getConcurrentMapForCaching();
+  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;
+      return CollectionFactory.createConcurrentSoftKeySoftValueMap(10, 0.75f, concurrencyLevel);
+    }
+    else {
+      return CollectionFactory.createConcurrentSoftValueMap();
+    }
+  }
 
   protected static final Logger logger = Logger.getInstance(TypeEvalContext.class);
 
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 c4cf06f0c13c..1e930301fee4 100644
--- a/python/python-psi-impl/resources/intellij.python.psi.impl.xml
+++ b/python/python-psi-impl/resources/intellij.python.psi.impl.xml
@@ -553,6 +553,8 @@
                  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"
+                 description="Enable soft-key caching for TypeEvalContext"/>
 
     <virtualFileCustomDataProvider implementation="com.jetbrains.python.formatter.PyCodeStyleSettingsProvider"/>
   </extensions>