[LL FIR] do not treat phase from FirFile as indicator for entire file

`resolvePhase` from `FirFile` is not an indicator of resolution state for
the entire file, so we shouldn't depend on it.
The right check will be performed during `lazyResolveTargets` over `LLFirResolveTarget`

^KTIJ-26666 Fixed
^KT-61296

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

@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirSingleRe
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeClassResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LLFirLazyResolverRunner
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.withOnAirDesignation
@@ -43,7 +42,14 @@ internal class LLFirModuleLazyDeclarationResolver(val moduleComponents: LLFirMod
         toPhase: FirResolvePhase,
     ) {
         val fromPhase = target.resolvePhase
-        if (fromPhase >= toPhase) return
+
+        /**
+         * Currently [lazyResolve] on file means [LLFirWholeFileResolveTarget], but also [FirFile] itself
+         * has [resolvePhase] which does not match with the entire file resolution state.
+         * This additional [FirFile] condition can be dropped after KT-61296
+         */
+        if (target !is FirFile && fromPhase >= toPhase) return
+
         try {
             resolveContainingFileToImports(target)
             if (toPhase == FirResolvePhase.IMPORTS) return

analysis/low-level-api-fir/testdata/inBlockModification/topLevelFunctionWithTypeWithBodyDumpFile.kt

@@ -0,0 +1,7 @@
+// DUMP_FILE
+
+fun foo(): Int {
+    return <expr>doSmth</expr>("str")
+}
+
+fun doSmth(i: String) = 4

analysis/low-level-api-fir/testdata/inBlockModification/topLevelFunctionWithTypeWithBodyDumpFile.txt

@@ -0,0 +1,15 @@
+BEFORE MODIFICATION:
+FILE: [ResolvedTo(BODY_RESOLVE)] topLevelFunctionWithTypeWithBodyDumpFile.kt
+    public final [ResolvedTo(BODY_RESOLVE)] fun foo(): R|kotlin/Int| {
+        ^foo R|/doSmth|(String(str))
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun doSmth([ResolvedTo(BODY_RESOLVE)] i: R|kotlin/String|): R|kotlin/Int| {
+        ^doSmth Int(4)
+    }
+
+AFTER MODIFICATION:
+FILE: [ResolvedTo(BODY_RESOLVE)] topLevelFunctionWithTypeWithBodyDumpFile.kt
+    public final [ResolvedTo(ANNOTATIONS_ARGUMENTS_MAPPING)] fun foo(): R|kotlin/Int| { LAZY_BLOCK }
+    public final [ResolvedTo(BODY_RESOLVE)] fun doSmth([ResolvedTo(BODY_RESOLVE)] i: R|kotlin/String|): R|kotlin/Int| {
+        ^doSmth Int(4)
+    }

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/AbstractCodeFragmentInBlockModificationTest.kt

@@ -21,7 +21,7 @@ abstract class AbstractCodeFragmentInBlockModificationTest : AbstractLowLevelApi
 
         assertNull(targetElement.getNonLocalReanalyzableContainingDeclaration())
 
-        val (before, after) = testInBlockModification(ktCodeFragment, ktCodeFragment, testServices)
+        val (before, after) = testInBlockModification(ktCodeFragment, ktCodeFragment, testServices, dumpFirFile = false)
         val actualText = "BEFORE MODIFICATION:\n$before\nAFTER MODIFICATION:\n$after"
 
         testServices.assertions.assertEqualsToTestDataFileSibling(actualText)

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/AbstractInBlockModificationTest.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure
 
 import com.intellij.extapi.psi.ASTDelegatePsiElement
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazyResolveRenderer
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolveWithCaches
@@ -16,15 +17,24 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.Analys
 import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
+import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
 import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestServices
 import org.jetbrains.kotlin.test.services.assertions
 
 abstract class AbstractInBlockModificationTest : AbstractLowLevelApiSingleFileTest() {
+    override fun configureTest(builder: TestConfigurationBuilder) {
+        super.configureTest(builder)
+        builder.useDirectives(Directives)
+    }
+
     override fun doTestByFileStructure(ktFile: KtFile, moduleStructure: TestModuleStructure, testServices: TestServices) {
         val selectedElement = testServices.expressionMarkerProvider.getSelectedElementOfTypeByDirective(
             ktFile = ktFile,
@@ -33,7 +43,13 @@ abstract class AbstractInBlockModificationTest : AbstractLowLevelApiSingleFileTe
 
         val declaration = selectedElement.getNonLocalReanalyzableContainingDeclaration()
         val actual = if (declaration != null) {
-            val (before, after) = testInBlockModification(ktFile, declaration, testServices)
+            val (before, after) = testInBlockModification(
+                file = ktFile,
+                declaration = declaration,
+                testServices = testServices,
+                dumpFirFile = Directives.DUMP_FILE in moduleStructure.allDirectives,
+            )
+
             "BEFORE MODIFICATION:\n$before\nAFTER MODIFICATION:\n$after"
         } else {
             "IN-BLOCK MODIFICATION IS NOT APPLICABLE FOR THIS PLACE"
@@ -41,33 +57,55 @@ abstract class AbstractInBlockModificationTest : AbstractLowLevelApiSingleFileTe
 
         testServices.assertions.assertEqualsToTestDataFileSibling(actual)
     }
+
+    private object Directives : SimpleDirectivesContainer() {
+        val DUMP_FILE by directive("Dump the entire FirFile to the output")
+    }
 }
 
-internal fun testInBlockModification(file: KtFile, declaration: KtAnnotated, testServices: TestServices): Pair<String, String> {
-    return resolveWithCaches(file) { firSession ->
-        val firDeclarationBefore = declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
-        val declarationTextBefore = firDeclarationBefore.render()
-
-        declaration.modifyBody()
-        invalidateAfterInBlockModification(declaration)
-
-        val declarationTextAfterModification = firDeclarationBefore.render()
-        testServices.assertions.assertNotEquals(declarationTextBefore, declarationTextAfterModification) {
-            "The declaration before and after modification must be in different state"
-        }
-
-        val firDeclarationAfter = declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
-        testServices.assertions.assertEquals(firDeclarationBefore, firDeclarationAfter) {
-            "The declaration before and after must be the same"
-        }
-
-        val declarationTextAfter = firDeclarationAfter.render()
-        testServices.assertions.assertEquals(declarationTextBefore, declarationTextAfter) {
-            "The declaration must have the same in the resolved state"
-        }
-
-        Pair(declarationTextBefore, declarationTextAfterModification)
+internal fun testInBlockModification(
+    file: KtFile,
+    declaration: KtAnnotated,
+    testServices: TestServices,
+    dumpFirFile: Boolean,
+): Pair<String, String> = resolveWithCaches(file) { firSession ->
+    val firDeclarationBefore = declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
+    val declarationToRender = if (dumpFirFile) {
+        file.getOrBuildFirFile(firSession).also { it.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE) }
+    } else {
+        firDeclarationBefore
     }
+
+    val textBefore = declarationToRender.render()
+
+    declaration.modifyBody()
+    invalidateAfterInBlockModification(declaration)
+
+    val textAfterModification = declarationToRender.render()
+    testServices.assertions.assertNotEquals(textBefore, textAfterModification) {
+        "The declaration before and after modification must be in different state"
+    }
+
+    val textAfter = if (dumpFirFile) {
+        // we should resolve the entire file instead of the declaration to be sure that this declaration will be
+        // resolved by file resolution as well
+        declarationToRender.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
+        declarationToRender.render()
+    } else {
+        declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
+        declarationToRender.render()
+    }
+
+    val firDeclarationAfter = declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
+    testServices.assertions.assertEquals(firDeclarationBefore, firDeclarationAfter) {
+        "The declaration before and after must be the same"
+    }
+
+    testServices.assertions.assertEquals(textBefore, textAfter) {
+        "The declaration must have the same in the resolved state"
+    }
+
+    Pair(textBefore, textAfterModification)
 }
 
 /**

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/OutOfContentRootInBlockModificationTestGenerated.java

@@ -288,6 +288,12 @@ public class OutOfContentRootInBlockModificationTestGenerated extends AbstractOu
         runTest("analysis/low-level-api-fir/testdata/inBlockModification/topLevelFunctionWithTypeWithBody.kt");
     }
 
+    @Test
+    @TestMetadata("topLevelFunctionWithTypeWithBodyDumpFile.kt")
+    public void testTopLevelFunctionWithTypeWithBodyDumpFile() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/inBlockModification/topLevelFunctionWithTypeWithBodyDumpFile.kt");
+    }
+
     @Test
     @TestMetadata("topLevelFunctionWithTypeWithoutBody.kt")
     public void testTopLevelFunctionWithTypeWithoutBody() throws Exception {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/SourceInBlockModificationTestGenerated.java

@@ -288,6 +288,12 @@ public class SourceInBlockModificationTestGenerated extends AbstractSourceInBloc
         runTest("analysis/low-level-api-fir/testdata/inBlockModification/topLevelFunctionWithTypeWithBody.kt");
     }
 
+    @Test
+    @TestMetadata("topLevelFunctionWithTypeWithBodyDumpFile.kt")
+    public void testTopLevelFunctionWithTypeWithBodyDumpFile() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/inBlockModification/topLevelFunctionWithTypeWithBodyDumpFile.kt");
+    }
+
     @Test
     @TestMetadata("topLevelFunctionWithTypeWithoutBody.kt")
     public void testTopLevelFunctionWithTypeWithoutBody() throws Exception {