IJPL-222467 Python: optimize access to caches during stub building with StubBuildCachedValuesManager

GitOrigin-RevId: 37c5ae56d98b1ca234dd5a1707f6780fccceb50b

platform/core-api/src/com/intellij/psi/stubs/StubBuildCachedValuesManager.kt

@@ -133,7 +133,7 @@ object StubBuildCachedValuesManager {
   ): T {
     val stubBuildId = stubBuildId
     if (stubBuildId != null) {
-      val node = dataHolder.getNode()
+      val node = dataHolder.getNode() ?: dataHolder
       var current = node.getUserData(stubBuildingKey)
       if (current == null || current.buildId != stubBuildId) {
         myComputingCachedValue.set(true)
@@ -194,7 +194,7 @@ object StubBuildCachedValuesManager {
   ): T? {
     val stubBuildId = stubBuildId
     if (stubBuildId != null) {
-      val node = dataHolder.getNode()
+      val node = dataHolder.getNode() ?: dataHolder
       var current = node.getUserData(stubBuildingKey)
       if (current == null || current.buildId != stubBuildId) {
         myComputingCachedValue.set(true)

python/python-ast/src/com/jetbrains/python/ast/PyAstTargetExpression.java

@@ -6,7 +6,11 @@ import com.intellij.psi.PsiComment;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiNameIdentifierOwner;
 import com.intellij.psi.PsiNamedElement;
-import com.intellij.psi.util.*;
+import com.intellij.psi.stubs.StubBuildCachedValuesManager.StubBuildCachedValueProvider;
+import com.intellij.psi.util.CachedValueProvider;
+import com.intellij.psi.util.PsiModificationTracker;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.PyTokenTypes;
@@ -18,6 +22,8 @@ import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import static com.intellij.psi.stubs.StubBuildCachedValuesManager.getCachedValueStubBuildOptimized;
+
 
 @ApiStatus.Experimental
 public interface PyAstTargetExpression extends PyAstQualifiedExpression, PsiNamedElement, PsiNameIdentifierOwner, PyAstDocStringOwner,
@@ -128,11 +134,16 @@ public interface PyAstTargetExpression extends PyAstQualifiedExpression, PsiName
    */
   default @Nullable PyAstExpression findAssignedValue() {
     PyPsiUtilsCore.assertValid(this);
-    return CachedValuesManager.getCachedValue(this,
-                                              () -> CachedValueProvider.Result
-                                                .create(findAssignedValueInternal(), PsiModificationTracker.MODIFICATION_COUNT));
+    return getCachedValueStubBuildOptimized(this, FIND_ASSIGNED_VALUE_PROVIDER);
   }
 
+  StubBuildCachedValueProvider<PyAstExpression, PyAstTargetExpression>
+    FIND_ASSIGNED_VALUE_PROVIDER = new StubBuildCachedValueProvider<>(
+    "python.findAssignedValue",
+    expression -> CachedValueProvider.Result
+      .create(expression.findAssignedValueInternal(), PsiModificationTracker.MODIFICATION_COUNT)
+  );
+
   private @Nullable PyAstExpression findAssignedValueInternal() {
     final PyAstAssignmentStatement assignment = PsiTreeUtil.getParentOfType(this, PyAstAssignmentStatement.class);
     if (assignment != null) {

python/python-ast/src/com/jetbrains/python/codeInsight/dataflow/scope/ScopeUtilCore.java

@@ -4,15 +4,16 @@ import com.intellij.extapi.psi.StubBasedPsiElementBase;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.StubBasedPsiElement;
+import com.intellij.psi.stubs.StubBuildCachedValuesManager.StubBuildCachedValueProvider;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.CachedValueProvider;
-import com.intellij.psi.util.CachedValuesManager;
 import com.intellij.psi.util.PsiModificationTracker;
 import com.jetbrains.python.ast.*;
 import com.jetbrains.python.ast.controlFlow.AstScopeOwner;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Nullable;
 
+import static com.intellij.psi.stubs.StubBuildCachedValuesManager.getCachedValueStubBuildOptimized;
 import static com.intellij.psi.util.PsiTreeUtil.getParentOfType;
 import static com.intellij.psi.util.PsiTreeUtil.isAncestor;
 
@@ -33,10 +34,16 @@ public final class ScopeUtilCore {
       final PsiElement context = element.getContext();
       return context instanceof AstScopeOwner ? (AstScopeOwner)context : getScopeOwner(context);
     }
-    return CachedValuesManager.getCachedValue(element, () -> CachedValueProvider.Result
-      .create(calculateScopeOwner(element), PsiModificationTracker.MODIFICATION_COUNT));
+    return getCachedValueStubBuildOptimized(element, GET_SCOPE_OWNER_PROVIDER);
   }
 
+  private static final StubBuildCachedValueProvider<AstScopeOwner, PsiElement>
+    GET_SCOPE_OWNER_PROVIDER = new StubBuildCachedValueProvider<>(
+    "python.scopeOwner",
+    element -> CachedValueProvider.Result
+      .create(calculateScopeOwner(element), PsiModificationTracker.MODIFICATION_COUNT)
+  );
+
   private static @Nullable AstScopeOwner calculateScopeOwner(@Nullable PsiElement element) {
     if (element instanceof StubBasedPsiElement<?> stubBasedElement) {
       final StubElement<?> stub = stubBasedElement.getStub();

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyVersionSpecificStubBase.kt

@@ -5,12 +5,13 @@ import com.intellij.openapi.util.Version
 import com.intellij.psi.PsiElement
 import com.intellij.psi.stubs.*
 import com.intellij.psi.util.CachedValueProvider
-import com.intellij.psi.util.CachedValuesManager
+import com.intellij.psi.stubs.StubBuildCachedValuesManager.StubBuildCachedValueProvider
 import com.jetbrains.python.PyTokenTypes
 import com.jetbrains.python.psi.*
 import com.jetbrains.python.psi.impl.PyEvaluator
 import com.jetbrains.python.psi.impl.PyPsiUtils
 import com.jetbrains.python.psi.stubs.PyVersionSpecificStub
+import com.intellij.psi.stubs.StubBuildCachedValuesManager.getCachedValueStubBuildOptimized
 
 internal abstract class PyVersionSpecificStubBase<T : PsiElement>(
   parent: StubElement<*>?,
@@ -26,26 +27,30 @@ internal fun getChildrenStubs(stub: StubElement<*>, languageLevel: LanguageLevel
 }
 
 internal fun evaluateVersionsForElement(element: PsiElement): ImmutableRangeSet<Version> {
-  return CachedValuesManager.getCachedValue(element) {
-    val parent = element.parent
-    var result: ImmutableRangeSet<Version>
-    if (parent == null) {
-      result = ImmutableRangeSet.of(Range.all())
-    }
-    else {
-      result = evaluateVersionsForElement(parent)
-      if (parent is PyIfPart || parent is PyElsePart) {
-        val grandParent = parent.parent
-        if (grandParent is PyIfStatement && element === (parent as PyStatementPart).statementList) {
-          val versions = evaluateVersionRangeForIfStatementPart(grandParent, parent)
-          if (versions != null) {
-            result = result.intersection(versions)
-          }
+  return getCachedValueStubBuildOptimized(element, EVAL_VERSIONS_PROVIDER)
+}
+
+private val EVAL_VERSIONS_PROVIDER = StubBuildCachedValueProvider<ImmutableRangeSet<Version>, PsiElement>(
+  "python.versionsForElement"
+) { element ->
+  val parent = element.parent
+  var result: ImmutableRangeSet<Version>
+  if (parent == null) {
+    result = ImmutableRangeSet.of(Range.all())
+  }
+  else {
+    result = evaluateVersionsForElement(parent)
+    if (parent is PyIfPart || parent is PyElsePart) {
+      val grandParent = parent.parent
+      if (grandParent is PyIfStatement && element === (parent as PyStatementPart).statementList) {
+        val versions = evaluateVersionRangeForIfStatementPart(grandParent, parent)
+        if (versions != null) {
+          result = result.intersection(versions)
         }
       }
     }
-    CachedValueProvider.Result.create(result, element)
   }
+  CachedValueProvider.Result.create(result, element)
 }
 
 private fun evaluateVersionRangeForIfStatementPart(ifStatement: PyIfStatement, ifStatementPart: PsiElement): RangeSet<Version>? {