[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
the target ancestor.

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

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt
index 9f76ddd..a337b55 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt
@@ -172,8 +172,9 @@
 
         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) { it === 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

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt
index af23b7f..6249622 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt
@@ -61,9 +61,17 @@
     }
 
     /**
-     * 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 preferBodyPredicate A predicate filtering elements for those the [ContextKind.BODY] 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, preferBodyPredicate: (PsiElement) -> Boolean): Context? {
         val acceptedElements = targetElement.parentsWithSelf.toSet()
 
         val contextProvider = process(computeResolveTarget(file, targetElement), holder) { candidate ->
@@ -74,14 +82,21 @@
             }
         }
 
-        if (preferBody) {
-            val bodyContext = contextProvider[targetElement, ContextKind.BODY]
-            if (bodyContext != null) {
-                return bodyContext
+        for (acceptedElement in acceptedElements) {
+            if (preferBodyPredicate(acceptedElement)) {
+                val bodyContext = contextProvider[acceptedElement, ContextKind.BODY]
+                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 {