[LL API] Support on-air resolution in 'FirCodeFragment'

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFil
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.resolve
 import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.CodeFragmentCapturedValueAnalyzer
 import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.CompilationPeerCollector
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
 import org.jetbrains.kotlin.backend.jvm.*
@@ -39,7 +40,6 @@ import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.diagnostics.Severity
 import org.jetbrains.kotlin.fir.backend.*
 import org.jetbrains.kotlin.fir.backend.jvm.*
-import org.jetbrains.kotlin.fir.declarations.FirCodeFragment
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.pipeline.applyIrGenerationExtensions
@@ -235,7 +235,7 @@ internal class KtFirCompilerFacility(
         resolveSession: LLFirResolveSession,
         configuration: CompilerConfiguration,
     ): CodeFragmentMappings {
-        val codeFragment = mainFirFile.declarations.single() as FirCodeFragment
+        val codeFragment = mainFirFile.codeFragment
 
         val capturedData = CodeFragmentCapturedValueAnalyzer.analyze(resolveSession, codeFragment)
 

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

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirResolvableS
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirElementFinder
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
@@ -51,8 +52,12 @@ object LowLevelFirApiFacadeForResolveOnAir {
         else -> true
     }
 
-    private fun PsiElement.onAirGetNonLocalContainingOrThisDeclaration(): KtDeclaration? {
-        return getNonLocalContainingOrThisDeclaration { declaration ->
+    private fun PsiElement.getNonLocalContainingOrThisDeclarationCodeFragmentAware(predicate: (KtDeclaration) -> Boolean): KtElement? {
+        return getNonLocalContainingOrThisDeclaration(predicate) ?: containingFile as? KtCodeFragment
+    }
+
+    private fun PsiElement.onAirGetNonLocalContainingOrThisDeclaration(): KtElement? {
+        return getNonLocalContainingOrThisDeclarationCodeFragmentAware { declaration ->
             declaration.isApplicableForOnAirResolve() && !declaration.canBePartOfParentDeclaration
         }
     }
@@ -176,7 +181,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         require(originalFirResolveSession is LLFirResolvableResolveSession)
         require(elementToAnalyze !is KtFile) { "KtFile for dependency element not supported" }
 
-        val minimalCopiedDeclaration = elementToAnalyze.getNonLocalContainingOrThisDeclaration {
+        val minimalCopiedDeclaration = elementToAnalyze.getNonLocalContainingOrThisDeclarationCodeFragmentAware {
             it.isApplicableForOnAirResolve()
         }
 
@@ -193,7 +198,12 @@ object LowLevelFirApiFacadeForResolveOnAir {
                 withPsiEntry("originalFile", originalKtFile)
             }
 
-        recordOriginalDeclaration(targetDeclaration = copiedNonLocalDeclaration, originalDeclaration = originalNonLocalDeclaration)
+        if (copiedNonLocalDeclaration is KtNamedDeclaration && originalNonLocalDeclaration is KtNamedDeclaration) {
+            recordOriginalDeclaration(
+                targetDeclaration = copiedNonLocalDeclaration,
+                originalDeclaration = originalNonLocalDeclaration
+            )
+        }
 
         val collector = FirTowerDataContextAllElementsCollector()
         val copiedFirDeclaration = runBodyResolveOnAir(
@@ -260,7 +270,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         collector: FirResolveContextCollector?,
         forcedResolvePhase: FirResolvePhase? = null,
     ): FirElement {
-        val minimalOriginalDeclarationToReplace = originalPlace.getNonLocalContainingOrThisDeclaration {
+        val minimalOriginalDeclarationToReplace = originalPlace.getNonLocalContainingOrThisDeclarationCodeFragmentAware {
             it.isApplicableForOnAirResolve()
         }
 
@@ -286,7 +296,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
             }
         }
 
-        val originalDesignationPath = FirElementFinder.collectDesignationPath(originalFirFile, originalDeclaration)
+        val originalDesignationPath = computeDesignation(originalDeclaration, originalFirFile)
             ?: errorWithFirSpecificEntries(message = "Impossible to collect designation", fir = originalFirFile, psi = originalDeclaration)
 
         val originalFirDeclaration = originalDesignationPath.target
@@ -321,17 +331,31 @@ object LowLevelFirApiFacadeForResolveOnAir {
         return newFirDeclaration
     }
 
-    private fun requiredResolvePhase(
-        container: KtDeclaration,
-        elementToReplace: PsiElement,
-    ): FirResolvePhase = when {
-        bodyResolveRequired(container, elementToReplace) -> FirResolvePhase.BODY_RESOLVE
-        annotationMappingRequired(container, elementToReplace) -> FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING
+    private fun computeDesignation(originalDeclaration: KtElement, originalFirFile: FirFile): FirElementFinder.FirDeclarationDesignation? {
+        if (originalDeclaration is KtCodeFragment) {
+            val firCodeFragment = originalFirFile.codeFragment
+            return FirElementFinder.FirDeclarationDesignation(emptyList(), firCodeFragment)
+        }
 
-        /**
-         * Currently it is a minimal phase there we can collect [FirResolveContextCollector]
-         */
-        else -> FirResolvePhase.ARGUMENTS_OF_ANNOTATIONS
+        require(originalDeclaration is KtDeclaration)
+        return FirElementFinder.collectDesignationPath(originalFirFile, originalDeclaration)
+    }
+
+    private fun requiredResolvePhase(
+        container: KtElement,
+        elementToReplace: PsiElement,
+    ): FirResolvePhase {
+        assert(container is KtDeclaration || container is KtCodeFragment)
+
+        return when {
+            container !is KtDeclaration -> FirResolvePhase.BODY_RESOLVE
+            bodyResolveRequired(container, elementToReplace) -> FirResolvePhase.BODY_RESOLVE
+            annotationMappingRequired(container, elementToReplace) -> FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING
+            else -> {
+                /** Currently it is a minimal phase there we can collect [FirResolveContextCollector] */
+                FirResolvePhase.ARGUMENTS_OF_ANNOTATIONS
+            }
+        }
     }
 
     private fun annotationMappingRequired(
@@ -374,6 +398,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
             }
             is KtClassOrObject -> check(ktDeclaration.body)
             is KtScript -> check(ktDeclaration.blockExpression)
+            is KtCodeFragment -> true
             is KtTypeAlias -> false
             else -> errorWithAttachment("Not supported type ${ktDeclaration::class}") {
                 withPsiEntry("declaration", ktDeclaration)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt

@@ -75,9 +75,12 @@ internal class FirElementBuilder(
     fun getOrBuildFirFor(
         element: KtElement,
         firResolveSession: LLFirResolveSession,
-    ): FirElement? = when (element) {
-        is KtFile -> getOrBuildFirForKtFile(element)
-        else -> getFirForNonKtFileElement(element, firResolveSession)
+    ): FirElement? {
+        return if (element is KtFile && element !is KtCodeFragment) {
+            getOrBuildFirForKtFile(element)
+        } else {
+            getFirForNonKtFileElement(element, firResolveSession)
+        }
     }
 
     private fun getOrBuildFirForKtFile(ktFile: KtFile): FirFile {
@@ -90,7 +93,7 @@ internal class FirElementBuilder(
         element: KtElement,
         firResolveSession: LLFirResolveSession,
     ): FirElement? {
-        require(element !is KtFile)
+        require(element !is KtFile || element is KtCodeFragment)
 
         if (!doKtElementHasCorrespondingFirElement(element)) {
             return null
@@ -247,6 +250,13 @@ internal val KtTypeParameter.containingDeclaration: KtDeclaration?
 internal val KtDeclaration.canBePartOfParentDeclaration: Boolean get() = this is KtPropertyAccessor || this is KtParameter || this is KtTypeParameter
 
 internal fun PsiElement.getNonLocalContainingOrThisDeclaration(predicate: (KtDeclaration) -> Boolean = { true }): KtDeclaration? {
+    return getNonLocalContainingDeclaration(parentsWithSelf, predicate)
+}
+
+internal fun getNonLocalContainingDeclaration(
+    elementsToCheck: Sequence<PsiElement>,
+    predicate: (KtDeclaration) -> Boolean = { true }
+): KtDeclaration? {
     var candidate: KtDeclaration? = null
 
     fun propose(declaration: KtDeclaration) {
@@ -255,7 +265,7 @@ internal fun PsiElement.getNonLocalContainingOrThisDeclaration(predicate: (KtDec
         }
     }
 
-    for (parent in parentsWithSelf) {
+    for (parent in elementsToCheck) {
         candidate?.let { notNullCandidate ->
             if (parent is KtEnumEntry ||
                 parent is KtCallableDeclaration &&

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt

@@ -9,6 +9,7 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
@@ -54,7 +55,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
             session: FirSession,
             scopeProvider: FirScopeProvider,
             designation: FirDesignation,
-            rootNonLocalDeclaration: KtDeclaration,
+            rootNonLocalDeclaration: KtElement,
             replacement: RawFirReplacement?,
         ): FirDeclaration {
             val replacementApplier = replacement?.Applier()
@@ -316,6 +317,10 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
                     }
                 }
                 is KtDestructuringDeclaration -> visitor.convertDestructuringDeclaration(declarationToBuild)
+                is KtCodeFragment -> {
+                    val firFile = visitor.convertElement(declarationToBuild, originalDeclaration) as FirFile
+                    firFile.codeFragment
+                }
                 else -> visitor.convertElement(declarationToBuild, originalDeclaration)
             } as FirDeclaration
         }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/state/LLFirResolvableResolveSession.kt

@@ -13,14 +13,13 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveCompone
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getModule
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
-import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
+import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.retryOnInvalidSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirDeclarationForCompiledElementSearcher
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
@@ -133,7 +132,7 @@ internal abstract class LLFirResolvableResolveSession(
             "Declaration should be resolvable module, instead it had ${module::class}"
         }
 
-        val nonLocalDeclaration = ktDeclaration.getNonLocalContainingOrThisDeclaration()
+        val nonLocalDeclaration = getNonLocalContainingDeclaration(ktDeclaration.parentsWithSelfCodeFragmentAware)
             ?: errorWithAttachment("Declaration should have non-local container") {
                 withPsiEntry("ktDeclaration", ktDeclaration, ::getModule)
                 withEntry("module", module) { it.moduleDescription }

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

@@ -8,9 +8,9 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.CodeFragmentCapturedValueAnalyzer
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.test.framework.utils.indented
-import org.jetbrains.kotlin.fir.declarations.FirCodeFragment
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.renderer.FirRenderer
@@ -30,7 +30,7 @@ abstract class AbstractCodeFragmentCapturingTest : AbstractLowLevelApiCodeFragme
         val resolveSession = module.getFirResolveSession(project)
         val firFile = ktCodeFragment.getOrBuildFirFile(resolveSession)
 
-        val firCodeFragment = firFile.declarations.single() as FirCodeFragment
+        val firCodeFragment = firFile.codeFragment
         firCodeFragment.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
 
         val capturedSymbols = CodeFragmentCapturedValueAnalyzer.analyze(resolveSession, firCodeFragment).symbols