PY-75831 Split cache(resolve/type) into library and user part

GitOrigin-RevId: 8dfd0120379c9a34051d66e147ffdc2c69f0db66

python/python-psi-api/src/com/jetbrains/python/psi/types/PyLibraryModificationTracker.kt

@@ -0,0 +1,80 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.psi.types
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.components.service
+import com.intellij.openapi.diagnostic.Logger
+import com.intellij.openapi.diagnostic.ThrottledLogger
+import com.intellij.openapi.fileEditor.FileDocumentManagerListener
+import com.intellij.openapi.project.DumbService
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.roots.ProjectRootManager
+import com.intellij.openapi.util.ModificationTracker
+import com.intellij.openapi.util.SimpleModificationTracker
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiManager
+import com.intellij.psi.PsiTreeChangeAdapter
+import com.intellij.psi.PsiTreeChangeEvent
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.ProjectScope
+import com.intellij.util.ForcefulReparseModificationTracker
+import java.util.concurrent.TimeUnit
+
+@Service(Service.Level.PROJECT)
+class PyLibraryModificationTracker(project: Project) : ModificationTracker, Disposable {
+  private val myProjectRootManager: ModificationTracker = ProjectRootManager.getInstance(project)
+  private val myDumbServiceModificationTracker: ModificationTracker = DumbService.getInstance(project).modificationTracker
+  private val myForcefulReparseModificationTracker: ModificationTracker = ForcefulReparseModificationTracker.getInstance() // PsiClass from libraries may become invalid on reparse
+  private val myOnContentReloadModificationTracker: SimpleModificationTracker = SimpleModificationTracker()
+  private val creationStack = Throwable()
+  val projectLibraryScope: GlobalSearchScope = ProjectScope.getLibrariesScope(project)
+
+  init {
+    val connection = project.getMessageBus().connect(this)
+    PsiManager.getInstance(project).addPsiTreeChangeListener(object : PsiTreeChangeAdapter() {
+      override fun childrenChanged(event: PsiTreeChangeEvent) {
+        val file = event.file ?: return
+        val virtualFile = file.virtualFile ?: return
+        if (isLibraryFile(virtualFile)) {
+          myOnContentReloadModificationTracker.incModificationCount()
+        }
+      }
+    }, this)
+
+    connection.subscribe<FileDocumentManagerListener>(FileDocumentManagerListener.TOPIC, object : FileDocumentManagerListener {
+      override fun fileWithNoDocumentChanged(file: VirtualFile) {
+        if (!project.isInitialized()) {
+          THROTTLED_LOG.warn("SearchScope.contains(file) would log an error because WorkspaceFileIndex is not yet initialized. " +
+                             "Probably LibraryModificationTracker was created too early. " +
+                             "See LibraryModificationTracker creation stacktrace: ", creationStack)
+          return
+        }
+        if (isLibraryFile(file)) {
+          myOnContentReloadModificationTracker.incModificationCount()
+        }
+      }
+    })
+  }
+
+  private fun isLibraryFile(file: VirtualFile): Boolean {
+    return "pyi" == file.extension || projectLibraryScope.contains(file)
+  }
+
+  override fun getModificationCount(): Long {
+    return (myProjectRootManager.getModificationCount()
+            + myDumbServiceModificationTracker.getModificationCount()
+            + myForcefulReparseModificationTracker.getModificationCount()
+            + myOnContentReloadModificationTracker.getModificationCount())
+  }
+
+
+  override fun dispose() {
+  }
+
+  companion object {
+    private val THROTTLED_LOG = ThrottledLogger(Logger.getInstance(PyLibraryModificationTracker::class.java), TimeUnit.SECONDS.toMillis(30))
+
+    fun getInstance(project: Project): PyLibraryModificationTracker = project.service()
+  }
+}

python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java

@@ -2,10 +2,12 @@
 package com.jetbrains.python.psi.types;
 
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.roots.ProjectFileIndex;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.RecursionManager;
 import com.intellij.openapi.util.registry.Registry;
 import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.util.ArrayUtil;
@@ -233,7 +235,28 @@ public sealed class TypeEvalContext {
     return null;
   }
 
+  private static boolean isLibraryElement(@NotNull PsiElement element) {
+    VirtualFile vFile = element.getContainingFile().getOriginalFile().getVirtualFile();
+    return vFile != null && ("pyi".equals(vFile.getExtension()) || ProjectFileIndex.getInstance(element.getProject()).isInLibrary(vFile));
+  }
+
+  private @NotNull TypeEvalContext getLibraryContext(@NotNull Project project) {
+    return project.getService(TypeEvalContextCache.class).getLibraryContext(new LibraryTypeEvalContext(getConstraints()));
+  }
+
+  /**
+   * If true the element's type will be calculated and stored in the long-life context bounded to the PyLibraryModificationTracker.
+   */
+  protected boolean canDelegateToLibraryContext(PyTypedElement element) {
+    return Registry.is("python.use.separated.libraries.type.cache") && isLibraryElement(element);
+  }
+
   public @Nullable PyType getType(final @NotNull PyTypedElement element) {
+    if (canDelegateToLibraryContext(element)) {
+      var context = getLibraryContext(element.getProject());
+      return context.getType(element);
+    }
+
     final PyType knownType = getKnownType(element);
     if (knownType != null) {
       return knownType == PyNullType.INSTANCE ? null : knownType;
@@ -251,6 +274,11 @@ public sealed class TypeEvalContext {
   }
 
   public @Nullable PyType getReturnType(final @NotNull PyCallable callable) {
+    if (canDelegateToLibraryContext(callable)) {
+      var context = getLibraryContext(callable.getProject());
+      return context.getReturnType(callable);
+    }
+
     final PyType knownReturnType = getKnownReturnType(callable);
     if (knownReturnType != null) {
       return knownReturnType == PyNullType.INSTANCE ? null : knownReturnType;
@@ -416,4 +444,16 @@ public sealed class TypeEvalContext {
       return this == o;
     }
   }
+
+  final static class LibraryTypeEvalContext extends TypeEvalContext {
+    private LibraryTypeEvalContext(@NotNull TypeEvalConstraints constraints) {
+      super(constraints);
+    }
+
+    @Override
+    protected boolean canDelegateToLibraryContext(PyTypedElement element) {
+      // It's already the library-context.
+      return false;
+    }
+  }
 }

python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContextCache.java

@@ -36,4 +36,7 @@ public interface TypeEvalContextCache {
    */
   @NotNull
   TypeEvalContext getContext(@NotNull TypeEvalContext standard);
+
+  @NotNull
+  TypeEvalContext getLibraryContext(@NotNull TypeEvalContext standard);
 }

python/python-psi-impl/resources/intellij.python.psi.impl.xml

@@ -492,6 +492,8 @@
                  description="Require marking namespace packages explicitly, treat regular directories as implicit source roots"/>
     <registryKey key="python.type.hints.literal.string" defaultValue="true"
                  description="When enabled, activates LiteralString inference for Python string literals" />
+    <registryKey key="python.use.separated.libraries.type.cache" defaultValue="true"
+                 description="It enables the use of a library-level cache for PSI elements from packages."/>
     <registryKey key="python.statement.lists.incremental.reparse" defaultValue="false"
                  description="Enables incremental reparse for statement lists"/>
   </extensions>

python/python-psi-impl/src/com/jetbrains/python/psi/types/TypeEvalContextCacheImpl.java

@@ -23,6 +23,7 @@ import java.util.concurrent.ConcurrentMap;
  */
 final class TypeEvalContextCacheImpl implements TypeEvalContextCache, Disposable {
   private final @NotNull CachedValue<ConcurrentMap<TypeEvalConstraints, TypeEvalContext>> myCachedMapStorage;
+  private final @NotNull CachedValue<ConcurrentMap<TypeEvalConstraints, TypeEvalContext>> myLibrariesCachedMapStorage;
   private final LowMemoryWatcher myLowMemoryWatcher;
   private final SimpleModificationTracker myLowMemoryModificationTracker = new SimpleModificationTracker();
 
@@ -36,6 +37,15 @@ final class TypeEvalContextCacheImpl implements TypeEvalContextCache, Disposable
         return new CachedValueProvider.Result<>(map, PsiModificationTracker.MODIFICATION_COUNT, myLowMemoryModificationTracker);
       }
     });
+    myLibrariesCachedMapStorage = CachedValuesManager.getManager(project).createCachedValue(new CachedValueProvider<>() {
+      @Override
+      public @NotNull CachedValueProvider.Result<ConcurrentMap<TypeEvalConstraints, TypeEvalContext>> compute() {
+        // This method is called if cache is empty. Create new map for it.
+        // Concurrent map allows several threads to call get and put, so it is thread safe but not atomic
+        final ConcurrentMap<TypeEvalConstraints, TypeEvalContext> map = ContainerUtil.createConcurrentSoftValueMap();
+        return new CachedValueProvider.Result<>(map, PyLibraryModificationTracker.Companion.getInstance(project));
+      }
+    });
 
     myLowMemoryWatcher = LowMemoryWatcher.register(() -> {
       myLowMemoryModificationTracker.incModificationCount();
@@ -43,12 +53,12 @@ final class TypeEvalContextCacheImpl implements TypeEvalContextCache, Disposable
     });
   }
 
-  @Override
-  public @NotNull TypeEvalContext getContext(@NotNull TypeEvalContext standard) {
+  private static TypeEvalContext retrieveFromStorage(@NotNull TypeEvalContext standard,
+                                                     CachedValue<ConcurrentMap<TypeEvalConstraints, TypeEvalContext>> storage) {
     // map is thread safe but not atomic nor getValue() is, so in worst case several threads may produce same result
     // both explicit locking and computeIfAbsent leads to deadlock
-    final ConcurrentMap<TypeEvalConstraints, TypeEvalContext> map = myCachedMapStorage.getValue();
     final TypeEvalConstraints key = standard.getConstraints();
+    final ConcurrentMap<TypeEvalConstraints, TypeEvalContext> map = storage.getValue();
     final TypeEvalContext cachedContext = map.get(key);
     if (cachedContext != null) {
       return cachedContext;
@@ -58,6 +68,16 @@ final class TypeEvalContextCacheImpl implements TypeEvalContextCache, Disposable
     return oldValue == null ? standard : oldValue;
   }
 
+  @Override
+  public @NotNull TypeEvalContext getContext(@NotNull TypeEvalContext standard) {
+    return retrieveFromStorage(standard, myCachedMapStorage);
+  }
+
+  @Override
+  public @NotNull TypeEvalContext getLibraryContext(@NotNull TypeEvalContext standard) {
+    return retrieveFromStorage(standard, myLibrariesCachedMapStorage);
+  }
+
   @Override
   public void dispose() {
     myLowMemoryWatcher.stop();