diff --git a/platform/core-impl/src/com/intellij/psi/stubs/DefaultStubBuilder.java b/platform/core-impl/src/com/intellij/psi/stubs/DefaultStubBuilder.java index 1a91ff737ae7..5ce880264df7 100644 --- a/platform/core-impl/src/com/intellij/psi/stubs/DefaultStubBuilder.java +++ b/platform/core-impl/src/com/intellij/psi/stubs/DefaultStubBuilder.java @@ -28,7 +28,12 @@ public class DefaultStubBuilder implements StubBuilder { } protected final @NotNull StubElement buildStubTreeFor(@NotNull ASTNode root, @NotNull StubElement parentStub) { - new StubBuildingWalkingVisitor(root, parentStub).buildStubTree(); + StubBuildCachedValuesManager.startBuildingStubs(); + try { + new StubBuildingWalkingVisitor(root, parentStub).buildStubTree(); + } finally { + StubBuildCachedValuesManager.finishBuildingStubs(); + } return parentStub; } diff --git a/platform/core-impl/src/com/intellij/psi/stubs/StubBuildCachedValuesManager.kt b/platform/core-impl/src/com/intellij/psi/stubs/StubBuildCachedValuesManager.kt new file mode 100644 index 000000000000..33fa0149a45f --- /dev/null +++ b/platform/core-impl/src/com/intellij/psi/stubs/StubBuildCachedValuesManager.kt @@ -0,0 +1,134 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.psi.stubs + +import com.intellij.openapi.util.Key +import com.intellij.psi.PsiElement +import com.intellij.psi.impl.source.tree.CompositeElement +import com.intellij.psi.util.CachedValueProvider +import com.intellij.psi.util.CachedValuesManager +import com.intellij.psi.util.ParameterizedCachedValue +import com.intellij.psi.util.ParameterizedCachedValueProvider +import org.jetbrains.annotations.ApiStatus +import java.util.concurrent.atomic.AtomicLong +import java.util.function.Function + +@ApiStatus.Internal +object StubBuildCachedValuesManager { + + private val myStubBuildId = ThreadLocal() + private val ourStubBuildIdCounter = AtomicLong() + + @JvmStatic + @ApiStatus.Internal + fun startBuildingStubs() { + myStubBuildId.set(ourStubBuildIdCounter.getAndIncrement()) + } + + @JvmStatic + @ApiStatus.Internal + fun finishBuildingStubs() { + myStubBuildId.remove() + } + + @JvmStatic + val isBuildingStubs: Boolean + get() = myStubBuildId.get() != null + + private val stubBuildId: Long? + get() = myStubBuildId.get() + + @JvmStatic + fun getCachedValueIfBuildingStubs( + dataHolder: PsiElement, + stubBuildingKey: Key>, + parameter: P, + provider: Function + ): T { + val stubBuildId = stubBuildId + if (stubBuildId != null) { + val node = dataHolder.getNode() + var current = node.getUserData(stubBuildingKey) + if (current == null || current.buildId != stubBuildId) { + current = StubBuildCachedValue(stubBuildId, provider.apply(parameter)) + node.putUserData>(stubBuildingKey, current) + } + return current.value + } + return provider.apply(parameter) + } + + @JvmStatic + fun getCachedValueStubBuildOptimized( + node: CompositeElement, + key: Key>, + stubBuildingKey: Key>, + provider: ParameterizedCachedValueProvider, + parameter: P + ): T { + val stubBuildId = stubBuildId + if (stubBuildId != null) { + var current = node.getUserData(stubBuildingKey) + if (current == null || current.buildId != stubBuildId) { + val value = provider.compute(parameter) + current = StubBuildCachedValue(stubBuildId, value.getValue()) + node.putUserData(stubBuildingKey, current) + } + return current.value + } + return CachedValuesManager.getManager(node.getManager().getProject()).getParameterizedCachedValue( + node, key, provider, false, parameter + ) + } + + // Avoid using recursion manager and other complex logic when building stubs - improves speed by 10%. + @JvmStatic + fun getCachedValueStubBuildOptimized( + dataHolder: PsiElement, + key: Key>, + stubBuildingKey: Key>, + provider: ParameterizedCachedValueProvider, + parameter: P + ): T { + val stubBuildId = stubBuildId + if (stubBuildId != null) { + val node = dataHolder.getNode() + var current = node.getUserData(stubBuildingKey) + if (current == null || current.buildId != stubBuildId) { + val value = provider.compute(parameter) + current = StubBuildCachedValue(stubBuildId, value.getValue()) + node.putUserData(stubBuildingKey, current) + } + return current.value + } + return CachedValuesManager.getManager(dataHolder.getProject()).getParameterizedCachedValue( + dataHolder, key, provider, false, parameter + ) + } + + /** + * This overload is provided for convenience. Consider + * using ParameterizedCachedValueProvider without a lambda, which + * improves performance by avoiding unnecessary lambda instantiation. + */ + @JvmStatic + fun getCachedValueStubBuildOptimized( + dataHolder: PsiElement, + stubBuildingKey: Key>, + provider: CachedValueProvider, + ): T? { + val stubBuildId = stubBuildId + if (stubBuildId != null) { + val node = dataHolder.getNode() + var current = node.getUserData(stubBuildingKey) + if (current == null || current.buildId != stubBuildId) { + val value = provider.compute() + current = StubBuildCachedValue(stubBuildId, value?.getValue()) + node.putUserData(stubBuildingKey, current) + } + return current.value + } + return CachedValuesManager.getCachedValue(dataHolder, provider) + } + + class StubBuildCachedValue(val buildId: Long, val value: T) +}