From 27a15cd76d34b199754c935fe5659eb0bbc35b29 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Wed, 2 Aug 2023 20:09:33 +0900
Subject: [PATCH] [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.
---
 .../fir/transformers/LLFirBodyLazyResolver.kt |  3 +-
 .../level/api/fir/util/ContextCollector.kt    | 29 ++++++++++++++-----
 .../api/fir/AbstractContextCollectorTest.kt   |  2 +-
 3 files changed, 25 insertions(+), 9 deletions(-)

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 03baed7e271..f0fd2333bcc 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
@@ -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)
                 }
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 312a9e1b4c8..f3fa4fe3b16 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 @@ 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) {
-            val bodyContext = contextProvider[targetElement, ContextKind.BODY]
-            if (bodyContext != null) {
-                return bodyContext
+        for (acceptedElement in acceptedElements) {
+            if (acceptedElement === bodyElement) {
+                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 {
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractContextCollectorTest.kt b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractContextCollectorTest.kt
index 5b2f2185086..c9bbe51cc3f 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractContextCollectorTest.kt
+++ b/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)