[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
2023-09-15 21:16:42 +02:00
parent f7fe33bd98
commit 1d0785e246
4 changed files with 112 additions and 9 deletions
@@ -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
@@ -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)
    }
 }

@@ -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 {
-        target is FirScript -> resolveScript(target)
-        target is FirCallableDeclaration && target.attributes.callableCopySubstitutionForTypeUpdater != null -> {
-            transformer.returnTypeCalculator.callableCopyTypeCalculator.computeReturnType(target)
-            Unit
+    override fun rawResolve(target: FirElementWithResolveState) {
+        when {
+            target is FirScript -> resolveScript(target)
+            target is FirCallableDeclaration && target.attributes.callableCopySubstitutionForTypeUpdater != null -> {
+                transformer.returnTypeCalculator.callableCopyTypeCalculator.computeReturnType(target)
+                Unit
+            }
+
+            else -> super.rawResolve(target)
        }

-        else -> super.rawResolve(target)
+        LLFirDeclarationModificationService.bodyResolved(target, resolverPhase)
    }
 }
@@ -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"
 }

 /**