[LL API] Move PSI element context collection to 'ContextCollector'

The new API is going to be reused in 'KtFirScopeProvider'.

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

@@ -5,15 +5,11 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.asResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError
 import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.codeFragmentScopeProvider
-import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
@@ -53,9 +49,8 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
-import org.jetbrains.kotlin.utils.exceptions.buildErrorWithAttachment
 import org.jetbrains.kotlin.utils.exceptions.requireWithAttachment
+import org.jetbrains.kotlin.utils.exceptions.withPsiEntry
 import org.jetbrains.kotlin.utils.findIsInstanceAnd
 
 internal object LLFirBodyLazyResolver : LLFirLazyResolver(FirResolvePhase.BODY_RESOLVE) {
@@ -180,20 +175,11 @@ private class LLFirBodyTargetResolver(
 
         firCodeFragment.codeFragmentContext = if (contextKtFile != null) {
             val contextFirFile = resolveSession.getOrBuildFirFile(contextKtFile)
-            val target = computeCodeFragmentContextResolveTarget(contextFirFile, contextPsiElement)
 
-            val contextParentKtElements = contextPsiElement.parentsWithSelf.toSet()
+            val elementContext = ContextCollector.process(contextFirFile, transformer.components, contextPsiElement, preferBody = true)
-
+                ?: errorWithAttachment("Cannot find enclosing context for ${contextPsiElement::class}") {
-            val contextProvider = ContextCollector.process(target, transformer.components) { candidate ->
+                    withPsiEntry("contextPsiElement", contextPsiElement)
-                when (candidate) {
-                    contextPsiElement -> ContextCollector.FilterResponse.STOP
-                    in contextParentKtElements -> ContextCollector.FilterResponse.CONTINUE
-                    else -> ContextCollector.FilterResponse.SKIP
                 }
-            }
-
-            val elementContext = contextProvider[contextPsiElement, ContextCollector.ContextKind.BODY]
-                ?: contextParentKtElements.firstNotNullOf { contextProvider[it, ContextCollector.ContextKind.SELF] }
 
             LLFirCodeFragmentContext(elementContext.towerDataContext.withExtraScopes(), elementContext.smartCasts)
         } else {
@@ -202,18 +188,6 @@ private class LLFirBodyTargetResolver(
         }
     }
 
-    private fun computeCodeFragmentContextResolveTarget(contextFirFile: FirFile, contextPsiElement: PsiElement): LLFirResolveTarget {
-        val contextKtDeclaration = contextPsiElement.getNonLocalContainingOrThisDeclaration()
-        if (contextKtDeclaration != null) {
-            val designationPath = FirElementFinder.collectDesignationPath(contextFirFile, contextKtDeclaration)
-            if (designationPath != null) {
-                return FirDesignationWithFile(designationPath.path, designationPath.target, contextFirFile).asResolveTarget()
-            }
-        }
-
-        return LLFirWholeFileResolveTarget(contextFirFile)
-    }
-
     override fun doLazyResolveUnderLock(target: FirElementWithResolveState) {
         when (target) {
             is FirRegularClass, is FirCodeFragment -> error("Should have been resolved in ${::doResolveWithoutLock.name}")

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

@@ -9,7 +9,11 @@ import com.intellij.openapi.progress.ProgressManager
 import com.intellij.psi.PsiElement
 import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.persistentMapOf
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.asResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.ContextKind
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.Context
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.FilterResponse
@@ -29,6 +33,7 @@ import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 
 internal object ContextCollector {
     enum class ContextKind {
@@ -55,6 +60,45 @@ internal object ContextCollector {
         SKIP
     }
 
+    /**
+     * Process the [file], collecting contexts for the [targetElement] and all its PSI tree parents.
+     */
+    fun process(file: FirFile, holder: SessionHolder, targetElement: PsiElement, preferBody: Boolean): Context? {
+        val acceptedElements = targetElement.parentsWithSelf.toSet()
+
+        val contextProvider = process(computeResolveTarget(file, targetElement), holder) { candidate ->
+            when (candidate) {
+                targetElement -> FilterResponse.STOP
+                in acceptedElements -> FilterResponse.CONTINUE
+                else -> FilterResponse.SKIP
+            }
+        }
+
+        if (preferBody) {
+            val bodyContext = contextProvider[targetElement, ContextKind.BODY]
+            if (bodyContext != null) {
+                return bodyContext
+            }
+        }
+
+        return acceptedElements.firstNotNullOfOrNull { contextProvider[it, ContextKind.SELF] }
+    }
+
+    private fun computeResolveTarget(file: FirFile, targetElement: PsiElement): LLFirResolveTarget {
+        val contextKtDeclaration = targetElement.getNonLocalContainingOrThisDeclaration()
+        if (contextKtDeclaration != null) {
+            val designationPath = FirElementFinder.collectDesignationPath(file, contextKtDeclaration)
+            if (designationPath != null) {
+                return FirDesignationWithFile(designationPath.path, designationPath.target, file).asResolveTarget()
+            }
+        }
+
+        return LLFirWholeFileResolveTarget(file)
+    }
+
+    /**
+     * Processes the [FirFile] that owns the [target], collecting contexts for elements matching the [filter].
+     */
     fun process(target: LLFirResolveTarget, holder: SessionHolder, filter: (PsiElement) -> FilterResponse): ContextProvider {
         val pathIterator = target.path.iterator()
 
@@ -67,10 +111,10 @@ internal object ContextCollector {
 
         return ContextProvider { element, kind -> visitor[element, kind] }
     }
-}
 
-internal fun interface ContextProvider {
+    fun interface ContextProvider {
-    operator fun get(element: PsiElement, kind: ContextKind): Context?
+        operator fun get(element: PsiElement, kind: ContextKind): Context?
+    }
 }
 
 private class ContextCollectorVisitor(