[LL FIR] do not try to apply in-block modifications for unresolved bodies

In-block modifications are applicable only for FIR declarations with
fully resolved bodies. It can happen only during implicit type and body
phases. The implicit type phase probably is not applicable for our
particular case, but we will play safe because properties have a
complicated resolution logic, so it is hard to be sure that only body
phase will be enough.
Changes inside contracts are treated as out-of-block, so we can avoid
marker from this phase. Also, all mentioned calls inside bodies resolve
the corresponding declarations to the contract phase, so we will save a
lot with this optimization.

Q: Why in-block modification can be heavy?
A: In most cases this invalidation is pretty cheap because we have all
required information inside caches. But it is possible that we can
receive an in-block modification event for PSI without a corresponding
FIR element, or it is not yet fully resolved. In this case, for example,
we will do such redundant works as a raw FIR building or KtModule search

^KT-61842 Fixed

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/LLFirDeclarationModificationService.kt

@@ -9,6 +9,7 @@ import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationListener
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
@@ -26,13 +27,18 @@ import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.providers.analysisMessageBus
 import org.jetbrains.kotlin.analysis.providers.topics.KotlinTopics.MODULE_OUT_OF_BLOCK_MODIFICATION
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
+import org.jetbrains.kotlin.fir.declarations.FirCodeFragment
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.psi
 import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtBlockExpression
 import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtDeclarationWithBody
-import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtProperty
@@ -66,6 +72,9 @@ class LLFirDeclarationModificationService(val project: Project) : Disposable {
     private var inBlockModificationQueue: MutableSet<ChangeType.InBlock>? = null
 
     private fun addModificationToQueue(modification: ChangeType.InBlock) {
+        // There is no sense to add into the queue elements with unresolved body
+        if (!modification.blockOwner.hasFirBody) return
+
         val queue = inBlockModificationQueue ?: HashSet<ChangeType.InBlock>().also { inBlockModificationQueue = it }
         queue += modification
     }
@@ -161,7 +170,7 @@ class LLFirDeclarationModificationService(val project: Project) : Disposable {
         }
     }
 
-    private fun inBlockModification(declaration: KtElement, ktModule: KtModule) {
+    private fun inBlockModification(declaration: KtAnnotated, ktModule: KtModule) {
         val resolveSession = ktModule.getFirResolveSession(project)
         val firDeclaration = when (declaration) {
             is KtCodeFragment -> declaration.getOrBuildFirFile(resolveSession).codeFragment
@@ -173,6 +182,7 @@ class LLFirDeclarationModificationService(val project: Project) : Disposable {
         }
 
         invalidateAfterInBlockModification(firDeclaration)
+        declaration.hasFirBody = false
 
         val moduleSession = firDeclaration.llFirResolvableSession ?: errorWithFirSpecificEntries(
             "${LLFirResolvableModuleSession::class.simpleName} is not found",
@@ -210,6 +220,53 @@ class LLFirDeclarationModificationService(val project: Project) : Disposable {
     companion object {
         fun getInstance(project: Project): LLFirDeclarationModificationService =
             project.getService(LLFirDeclarationModificationService::class.java)
+
+        /**
+         * This function have to be called from Low Level FIR body transformers.
+         * It is fine to have false-positives, but false-negatives are not acceptable.
+         */
+        internal fun bodyResolved(element: FirElementWithResolveState, phase: FirResolvePhase) {
+            when (element) {
+                is FirSimpleFunction -> {
+                    // in-block modifications only applicable to functions with an explicit type,
+                    // so we mark only fully resolved functions
+                    if (phase != FirResolvePhase.BODY_RESOLVE) return
+                }
+
+                is FirProperty -> {
+                    // in-block modifications only applicable to properties with an explicit type,
+                    // but existed backing field can lead to the entire body resolution even on
+                    // implicit body phase, so we will mark this phase as fully resolved too to be safe
+                    if (phase != FirResolvePhase.BODY_RESOLVE && phase != FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) return
+                }
+
+                is FirCodeFragment -> {
+                    // in-block modifications only applicable to fully resolved code fragments
+                    if (phase != FirResolvePhase.BODY_RESOLVE) return
+                }
+
+                else -> return
+            }
+
+            val declaration = element.source?.psi as? KtAnnotated ?: return
+            when (declaration) {
+                is KtNamedFunction -> {
+                    if (declaration.isReanalyzableContainer()) {
+                        declaration.hasFirBody = true
+                    }
+                }
+
+                is KtProperty -> {
+                    if (declaration.isReanalyzableContainer() || declaration.accessors.any(KtPropertyAccessor::isReanalyzableContainer)) {
+                        declaration.hasFirBody = true
+                    }
+                }
+
+                is KtCodeFragment -> {
+                    declaration.hasFirBody = true
+                }
+            }
+        }
     }
 }
 
@@ -231,6 +288,28 @@ private sealed class ChangeType {
     }
 }
 
+/**
+ * The purpose of this property as user data is to avoid FIR building in case the [KtAnnotated]
+ * doesn't have an associated FIR with body.
+ *
+ * [KtProperty] is used as an anchor for [KtPropertyAccessor]s to avoid extra memory consumption.
+ */
+private var KtAnnotated.hasFirBody: Boolean
+    get() = when (this) {
+        is KtNamedFunction, is KtProperty, is KtCodeFragment -> getUserData(hasFirBodyKey) == true
+        is KtPropertyAccessor -> property.hasFirBody
+        else -> false
+    }
+    set(value) {
+        val declarationAnchor = if (this is KtPropertyAccessor) property else this
+        declarationAnchor.putUserData(
+            hasFirBodyKey,
+            value.takeIf { it },
+        )
+    }
+
+private val hasFirBodyKey = Key.create<Boolean?>("HAS_FIR_BODY")
+
 /**
  * Covered by org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.AbstractInBlockModificationTest
  * on the compiler side and by

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveT
 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.file.builder.LLFirLockProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableResolveSession
@@ -258,6 +259,8 @@ private class LLFirBodyTargetResolver(
             is FirScript -> target.takeUnless(FirScript::isCertainlyResolved)?.let(::resolveScript)
             else -> super.rawResolve(target)
         }
+
+        LLFirDeclarationModificationService.bodyResolved(target, resolverPhase)
     }
 }
 

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkReturnTypeRefIsResolved
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isScriptDependentDeclaration
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
@@ -89,13 +90,17 @@ internal class LLFirImplicitBodyTargetResolver(
         }
     }
 
-    override fun rawResolve(target: FirElementWithResolveState): Unit = when {
+    override fun rawResolve(target: FirElementWithResolveState) {
-        target is FirScript -> resolveScript(target)
+        when {
-        target is FirCallableDeclaration && target.attributes.callableCopySubstitutionForTypeUpdater != null -> {
+            target is FirScript -> resolveScript(target)
-            transformer.returnTypeCalculator.callableCopyTypeCalculator.computeReturnType(target)
+            target is FirCallableDeclaration && target.attributes.callableCopySubstitutionForTypeUpdater != null -> {
-            Unit
+                transformer.returnTypeCalculator.callableCopyTypeCalculator.computeReturnType(target)
+                Unit
+            }
+
+            else -> super.rawResolve(target)
         }
 
-        else -> super.rawResolve(target)
+        LLFirDeclarationModificationService.bodyResolved(target, resolverPhase)
     }
 }

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure
 import com.intellij.extapi.psi.ASTDelegatePsiElement
 import com.intellij.openapi.util.Disposer
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 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.element.builder.getNonLocalContainingOrThisDeclaration
@@ -67,6 +68,21 @@ internal fun testInBlockModification(
     testServices: TestServices,
     dumpFirFile: Boolean,
 ): String = resolveWithCaches(file) { firSession ->
+    // We are trying to invoke a test case twice inside one session to be sure that sequent modifications are work
+    val firstAttempt = doTestInBlockModification(file, elementToModify, testServices, dumpFirFile, firSession)
+    val secondAttempt = doTestInBlockModification(file, elementToModify, testServices, dumpFirFile, firSession)
+    testServices.assertions.assertEquals(firstAttempt, secondAttempt) { "Invocations must be the same" }
+
+    firstAttempt
+}
+
+private fun doTestInBlockModification(
+    file: KtFile,
+    elementToModify: PsiElement,
+    testServices: TestServices,
+    dumpFirFile: Boolean,
+    firSession: LLFirResolveSession,
+): String {
     val declaration = elementToModify.getNonLocalContainingOrThisDeclaration() ?: file
     val firDeclarationBefore = declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
     val declarationToRender = if (dumpFirFile) {
@@ -116,7 +132,7 @@ internal fun testInBlockModification(
         "The declaration must have the same in the resolved state"
     }
 
-    "BEFORE MODIFICATION:\n$textBefore\nAFTER MODIFICATION:\n$textAfterModification"
+    return "BEFORE MODIFICATION:\n$textBefore\nAFTER MODIFICATION:\n$textAfterModification"
 }
 
 /**