[LL API] Avoid unnecessary BODY context collection

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt

@@ -72,9 +72,10 @@ internal object ContextCollector {
      * Returns `null` if the context was not collected.
      */
     fun process(file: FirFile, holder: SessionHolder, targetElement: PsiElement, bodyElement: PsiElement? = targetElement): Context? {
+        val isBodyContextCollected = bodyElement != null
         val acceptedElements = targetElement.parentsWithSelf.toSet()
 
-        val contextProvider = process(computeResolveTarget(file, targetElement), holder) { candidate ->
+        val contextProvider = process(computeResolveTarget(file, targetElement), holder, isBodyContextCollected) { candidate ->
             when (candidate) {
                 targetElement -> FilterResponse.STOP
                 in acceptedElements -> FilterResponse.CONTINUE
@@ -113,11 +114,21 @@ internal object ContextCollector {
 
     /**
      * Processes the [FirFile] that owns the [target], collecting contexts for elements matching the [filter].
+     *
+     * @param target The target to process.
+     * @param holder The [SessionHolder] for the session that owns a [target].
+     * @param shouldCollectBodyContext If `true`, [ContextKind.BODY] is collected where available.
+     * @param filter The filter predicate. Context is collected only for [PsiElement]s for which the [filter] returns `true`.
      */
-    fun process(target: LLFirResolveTarget, holder: SessionHolder, filter: (PsiElement) -> FilterResponse): ContextProvider {
+    fun process(
+        target: LLFirResolveTarget,
+        holder: SessionHolder,
+        shouldCollectBodyContext: Boolean,
+        filter: (PsiElement) -> FilterResponse
+    ): ContextProvider {
         val pathIterator = target.path.iterator()
 
-        val visitor = ContextCollectorVisitor(holder, filter) {
+        val visitor = ContextCollectorVisitor(holder, shouldCollectBodyContext, filter) {
             if (pathIterator.hasNext()) pathIterator.next() else null
         }
 
@@ -134,6 +145,7 @@ internal object ContextCollector {
 
 private class ContextCollectorVisitor(
     private val holder: SessionHolder,
+    private val shouldCollectBodyContext: Boolean,
     private val filter: (PsiElement) -> FilterResponse,
     private val interceptor: () -> FirElement?
 ) : FirDefaultVisitorVoid() {
@@ -177,6 +189,10 @@ private class ContextCollectorVisitor(
             return
         }
 
+        if (kind == ContextKind.BODY && !shouldCollectBodyContext) {
+            return
+        }
+
         val key = ContextKey(psi, kind)
         if (key in result) {
             return
@@ -203,6 +219,7 @@ private class ContextCollectorVisitor(
     override fun visitAnnotationCall(annotationCall: FirAnnotationCall) {
         dumpContext(annotationCall.psi, ContextKind.SELF)
 
+        onActiveBody {
             context.forAnnotation {
                 dumpContext(annotationCall.psi, ContextKind.BODY)
 
@@ -211,15 +228,17 @@ private class ContextCollectorVisitor(
                 // Here we ignore declarations that might be inside such expressions, avoiding unnecessary tree traversal.
             }
         }
+    }
 
     override fun visitRegularClass(regularClass: FirRegularClass) = withProcessor {
         dumpContext(regularClass.psi, ContextKind.SELF)
 
         processSignatureAnnotations(regularClass)
 
-        context.withContainingClass(regularClass) {
+        onActiveBody {
             regularClass.lazyResolveToPhase(FirResolvePhase.STATUS)
 
+            context.withContainingClass(regularClass) {
                 processList(regularClass.typeParameters)
 
                 context.withRegularClass(regularClass, holder) {
@@ -232,6 +251,7 @@ private class ContextCollectorVisitor(
                     }
                 }
             }
+        }
 
         if (regularClass.isLocal) {
             context.storeClassIfNotNested(regularClass, session)
@@ -243,9 +263,10 @@ private class ContextCollectorVisitor(
 
         processSignatureAnnotations(constructor)
 
-        context.withConstructor(constructor) {
+        onActiveBody {
             constructor.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
 
+            context.withConstructor(constructor) {
                 val owningClass = context.containerIfAny as? FirRegularClass
                 context.forConstructorParameters(constructor, owningClass, holder) {
                     processList(constructor.valueParameters)
@@ -272,10 +293,12 @@ private class ContextCollectorVisitor(
                 }
             }
         }
+    }
 
     override fun visitEnumEntry(enumEntry: FirEnumEntry) {
         dumpContext(enumEntry.psi, ContextKind.SELF)
 
+        onActiveBody {
             enumEntry.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
 
             context.forEnumEntry {
@@ -286,14 +309,16 @@ private class ContextCollectorVisitor(
                 }
             }
         }
+    }
 
     override fun visitSimpleFunction(simpleFunction: FirSimpleFunction) = withProcessor {
         dumpContext(simpleFunction.psi, ContextKind.SELF)
 
-        simpleFunction.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
-
         processSignatureAnnotations(simpleFunction)
 
+        onActiveBody {
+            simpleFunction.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
+
             context.withSimpleFunction(simpleFunction, session) {
                 context.forFunctionBody(simpleFunction, holder) {
                     processList(simpleFunction.valueParameters)
@@ -310,20 +335,22 @@ private class ContextCollectorVisitor(
                 }
             }
         }
+    }
 
     override fun visitProperty(property: FirProperty) = withProcessor {
         dumpContext(property.psi, ContextKind.SELF)
 
         processSignatureAnnotations(property)
 
-        onActive {
+        onActiveBody {
             property.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
 
             context.withProperty(property) {
-                context.forPropertyInitializer {
+                dumpContext(property.psi, ContextKind.BODY)
+
                 onActive {
+                    context.forPropertyInitializer {
                         process(property.initializer)
-                    }
 
                         onActive {
                             process(property.delegate)
@@ -333,6 +360,7 @@ private class ContextCollectorVisitor(
                             process(property.backingField)
                         }
                     }
+                }
 
                 onActive {
                     processChildren(property)
@@ -350,6 +378,7 @@ private class ContextCollectorVisitor(
 
         processSignatureAnnotations(propertyAccessor)
 
+        onActiveBody {
             context.withPropertyAccessor(propertyAccessor.propertySymbol.fir, propertyAccessor, holder) {
                 dumpContext(propertyAccessor.psi, ContextKind.BODY)
 
@@ -358,12 +387,14 @@ private class ContextCollectorVisitor(
                 }
             }
         }
+    }
 
     override fun visitValueParameter(valueParameter: FirValueParameter) = withProcessor {
         dumpContext(valueParameter.psi, ContextKind.SELF)
 
         processSignatureAnnotations(valueParameter)
 
+        onActiveBody {
             context.withValueParameter(valueParameter, session) {
                 dumpContext(valueParameter.psi, ContextKind.BODY)
 
@@ -373,11 +404,14 @@ private class ContextCollectorVisitor(
             }
         }
 
+    }
+
     override fun visitAnonymousInitializer(anonymousInitializer: FirAnonymousInitializer) = withProcessor {
         dumpContext(anonymousInitializer.psi, ContextKind.SELF)
 
         processSignatureAnnotations(anonymousInitializer)
 
+        onActiveBody {
             context.withAnonymousInitializer(anonymousInitializer, session) {
                 dumpContext(anonymousInitializer.psi, ContextKind.BODY)
 
@@ -387,12 +421,14 @@ private class ContextCollectorVisitor(
                 }
             }
         }
+    }
 
     override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction) = withProcessor {
         dumpContext(anonymousFunction.psi, ContextKind.SELF)
 
         processSignatureAnnotations(anonymousFunction)
 
+        onActiveBody {
             context.withAnonymousFunction(anonymousFunction, holder, ResolutionMode.ContextIndependent) {
                 for (parameter in anonymousFunction.valueParameters) {
                     process(parameter)
@@ -411,11 +447,14 @@ private class ContextCollectorVisitor(
             }
         }
 
+    }
+
     override fun visitAnonymousObject(anonymousObject: FirAnonymousObject) = withProcessor {
         dumpContext(anonymousObject.psi, ContextKind.SELF)
 
         processSignatureAnnotations(anonymousObject)
 
+        onActiveBody {
             context.withAnonymousObject(anonymousObject, holder) {
                 dumpContext(anonymousObject.psi, ContextKind.BODY)
 
@@ -425,15 +464,19 @@ private class ContextCollectorVisitor(
             }
         }
 
+    }
+
     override fun visitBlock(block: FirBlock) = withProcessor {
         dumpContext(block.psi, ContextKind.SELF)
 
+        onActiveBody {
             context.forBlock(session) {
                 processChildren(block)
 
                 dumpContext(block.psi, ContextKind.BODY)
             }
         }
+    }
 
     override fun visitSmartCastExpression(smartCastExpression: FirSmartCastExpression) {
         if (smartCastExpression.isStable) {
@@ -514,6 +557,12 @@ private class ContextCollectorVisitor(
             block()
         }
     }
+
+    private inline fun onActiveBody(block: () -> Unit) {
+        if (isActive || shouldCollectBodyContext) {
+            block()
+        }
+    }
 }
 
 @DslMarker