[LL API] Add more flexible 'preferBodyPredicate'

The previous API required to pass the exact declaration element to
the 'ContextCollector'. Now it's possible to get a 'BODY' context for
some ancestor element.

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt

@@ -175,8 +175,9 @@ private class LLFirBodyTargetResolver(
 
         firCodeFragment.codeFragmentContext = if (contextKtFile != null) {
             val contextFirFile = resolveSession.getOrBuildFirFile(contextKtFile)
+            val sessionHolder = transformer.components
 
-            val elementContext = ContextCollector.process(contextFirFile, transformer.components, contextPsiElement, preferBody = true)
+            val elementContext = ContextCollector.process(contextFirFile, sessionHolder, contextPsiElement)
                 ?: errorWithAttachment("Cannot find enclosing context for ${contextPsiElement::class}") {
                     withPsiEntry("contextPsiElement", contextPsiElement)
                 }

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

@@ -61,9 +61,17 @@ internal object ContextCollector {
     }
 
     /**
-     * Process the [file], collecting contexts for the [targetElement] and all its PSI tree parents.
+     * Get the most precise context available for the [targetElement] in the [file].
+     *
+     * @param file The file to process.
+     * @param holder A [SessionHolder] for the session that owns a [file].
+     * @param targetElement The most precise element for which the context is required.
+     * @param bodyElement An element for which the [ContextKind.BODY] context is preferred.
+     *
+     * Returns the context of the [targetElement] if available, or of one of its tree parents.
+     * Returns `null` if the context was not collected.
      */
-    fun process(file: FirFile, holder: SessionHolder, targetElement: PsiElement, preferBody: Boolean): Context? {
+    fun process(file: FirFile, holder: SessionHolder, targetElement: PsiElement, bodyElement: PsiElement? = targetElement): Context? {
         val acceptedElements = targetElement.parentsWithSelf.toSet()
 
         val contextProvider = process(computeResolveTarget(file, targetElement), holder) { candidate ->
@@ -74,14 +82,21 @@ internal object ContextCollector {
             }
         }
 
-        if (preferBody) {
+        for (acceptedElement in acceptedElements) {
-            val bodyContext = contextProvider[targetElement, ContextKind.BODY]
+            if (acceptedElement === bodyElement) {
-            if (bodyContext != null) {
+                val bodyContext = contextProvider[acceptedElement, ContextKind.BODY]
-                return bodyContext
+                if (bodyContext != null) {
+                    return bodyContext
+                }
+            }
+
+            val elementContext = contextProvider[acceptedElement, ContextKind.SELF]
+            if (elementContext != null) {
+                return elementContext
             }
         }
 
-        return acceptedElements.firstNotNullOfOrNull { contextProvider[it, ContextKind.SELF] }
+        return null
     }
 
     private fun computeResolveTarget(file: FirFile, targetElement: PsiElement): LLFirResolveTarget {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractContextCollectorTest.kt

@@ -48,7 +48,7 @@ abstract class AbstractContextCollectorTest : AbstractAnalysisApiBasedSingleModu
         val targetElement = testServices.expressionMarkerProvider
             .getBottommostSelectedElementOfType(mainKtFile, KtElement::class.java)
 
-        val elementContext = ContextCollector.process(firFile, sessionHolder, targetElement) { it === targetElement }
+        val elementContext = ContextCollector.process(firFile, sessionHolder, targetElement)
             ?: error("Context not found for element $targetElement")
 
         val actualText = ElementContextRenderer.render(elementContext)