FIR IDE: introduce FileElementFactory

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/FirIdeStructureElementDiagnosticsCollector.kt

@@ -51,5 +51,28 @@ internal class FirIdeStructureElementDiagnosticsCollector private constructor(
                 collector.collectDiagnostics(firFile)
                 FileStructureElementDiagnostics(collector.result)
             }
+
+        fun collectForSingleDeclaration(firFile: FirFile, declaration: FirDeclaration): FileStructureElementDiagnostics {
+            var inCurrentDeclaration = false
+
+            return collectForStructureElement(
+                firFile,
+                onDeclarationEnter = { firDeclaration ->
+                    when {
+                        firDeclaration == declaration -> {
+                            inCurrentDeclaration = true
+                            DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                        }
+                        inCurrentDeclaration -> DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
+                        else -> DiagnosticCollectorDeclarationAction.SKIP_CURRENT_DECLARATION_AND_CHECK_NESTED
+                    }
+                },
+                onDeclarationExit = { firDeclaration ->
+                    if (declaration == firDeclaration) {
+                        inCurrentDeclaration = false
+                    }
+                }
+            )
+        }
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileElementFactory.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
+
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtNamedFunction
+
+internal object FileElementFactory {
+    /**
+     * should be consistent with [isReanalyzableContainer]
+     */
+    fun createFileStructureElement(
+        firDeclaration: FirDeclaration,
+        ktDeclaration: KtDeclaration,
+        firFile: FirFile,
+    ): FileStructureElement = when {
+        ktDeclaration is KtNamedFunction && ktDeclaration.name != null && ktDeclaration.hasExplicitTypeOrUnit ->
+            IncrementallyReanalyzableFunction(
+                firFile,
+                ktDeclaration,
+                (firDeclaration as FirSimpleFunction).symbol,
+                ktDeclaration.modificationStamp
+            )
+
+        else -> NonLocalDeclarationFileStructureElement(
+            firFile,
+            firDeclaration,
+            ktDeclaration,
+        )
+    }
+
+    /**
+     * should be consistent with [createFileStructureElement]
+     */
+    fun isReanalyzableContainer(
+        ktDeclaration: KtDeclaration,
+    ): Boolean = when {
+        ktDeclaration is KtNamedFunction && ktDeclaration.name != null && ktDeclaration.hasExplicitTypeOrUnit -> true
+        else -> false
+    }
+
+    val KtNamedFunction.hasExplicitTypeOrUnit
+        get() = hasBlockBody() || typeReference != null
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructure.kt

@@ -6,17 +6,13 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
 import org.jetbrains.kotlin.diagnostics.Diagnostic
-import org.jetbrains.kotlin.fir.containingClass
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.hasExplicitTypeOrUnit
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.replaceFirst
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
@@ -44,8 +40,8 @@ internal class FileStructure(
         val structureElement = structureElements.compute(declaration) { _, structureElement ->
             when {
                 structureElement == null -> createStructureElement(declaration)
-                structureElement is WithInBlockModificationFileStructureElement && !structureElement.isUpToDate() -> {
+                structureElement is ReanalyzableStructureElement<*> && !structureElement.isUpToDate() -> {
-                    createMappingsCopy(structureElement, declaration as KtNamedFunction)
+                    structureElement.reanalyze(declaration as KtNamedFunction, moduleFileCache, firLazyDeclarationResolver, firIdeProvider)
                 }
                 else -> structureElement
             }
@@ -72,52 +68,6 @@ internal class FileStructure(
         }
     }
 
-    private fun replaceFunction(from: FirSimpleFunction, to: FirSimpleFunction) {
-        val declarations = if (from.symbol.callableId.className == null) {
-            firFile.declarations as MutableList<FirDeclaration>
-        } else {
-            val classLikeLookupTag = from.containingClass()
-                ?: error("Class name should not be null for non-top-level & non-local declarations")
-            val containingClass = classLikeLookupTag.toSymbol(firFile.session)?.fir as FirRegularClass
-            containingClass.declarations as MutableList<FirDeclaration>
-        }
-        declarations.replaceFirst(from, to)
-    }
-
-    private fun createMappingsCopy(
-        original: WithInBlockModificationFileStructureElement,
-        containerKtFunction: KtNamedFunction
-    ): WithInBlockModificationFileStructureElement {
-        val newFunction = firIdeProvider.buildFunctionWithBody(containerKtFunction) as FirSimpleFunction
-        val originalFunction = original.firSymbol.fir as FirSimpleFunction
-
-        moduleFileCache.firFileLockProvider.withWriteLock(firFile) {
-            replaceFunction(originalFunction, newFunction)
-        }
-
-        try {
-            firLazyDeclarationResolver.lazyResolveDeclaration(
-                newFunction,
-                moduleFileCache,
-                FirResolvePhase.BODY_RESOLVE,
-                checkPCE = true,
-                reresolveFile = true,
-            )
-            return moduleFileCache.firFileLockProvider.withReadLock(firFile) {
-                WithInBlockModificationFileStructureElement(
-                    firFile,
-                    containerKtFunction,
-                    newFunction.symbol,
-                    containerKtFunction.modificationStamp,
-                )
-            }
-        } catch (e: Throwable) {
-            moduleFileCache.firFileLockProvider.withWriteLock(firFile) {
-                replaceFunction(newFunction, originalFunction)
-            }
-            throw e
-        }
-    }
 
     private fun createDeclarationStructure(declaration: KtDeclaration): FileStructureElement {
         val firDeclaration = declaration.findSourceNonLocalFirDeclaration(
@@ -133,23 +83,7 @@ internal class FileStructure(
             checkPCE = true
         )
         return moduleFileCache.firFileLockProvider.withReadLock(firFile) {
-            when {
+            FileElementFactory.createFileStructureElement(firDeclaration, declaration, firFile)
-                declaration is KtNamedFunction && declaration.hasExplicitTypeOrUnit -> {
-                    WithInBlockModificationFileStructureElement(
-                        firFile,
-                        declaration,
-                        (firDeclaration as FirSimpleFunction).symbol,
-                        declaration.modificationStamp,
-                    )
-                }
-                else -> {
-                    NonLocalDeclarationFileStructureElement(
-                        firFile,
-                        firDeclaration,
-                        declaration,
-                    )
-                }
-            }
         }
     }
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureElement.kt

@@ -8,16 +8,19 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.collectors.DiagnosticCollectorDeclarationAction
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.containingClass
-import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FirIdeStructureElementDiagnosticsCollector
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.hasExplicitTypeOrUnit
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
+import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.ktDeclaration
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.replaceFirst
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfTypeTo
-import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 internal class FileStructureElementDiagnostics(
     private val map: Map<KtElement, List<Diagnostic>>
@@ -34,43 +37,90 @@ internal sealed class FileStructureElement {
     abstract val diagnostics: FileStructureElementDiagnostics
 }
 
-internal class WithInBlockModificationFileStructureElement(
+internal sealed class ReanalyzableStructureElement<KT : KtDeclaration> : FileStructureElement() {
-    override val firFile: FirFile,
+    abstract override val psi: KT
-    override val psi: KtFunction,
+    abstract val firSymbol: AbstractFirBasedSymbol<*>
-    val firSymbol: FirFunctionSymbol<*>,
+    abstract val timestamp: Long
-    val timestamp: Long
-) : FileStructureElement() {
 
-    override val mappings: Map<KtElement, FirElement> =
+    /**
-        FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
+     * Creates new declaration by [newKtDeclaration] which will serve as replacement of [firSymbol]
+     * Also, modify [firFile] & replace old version of declaration to a new one
+     */
+    abstract fun reanalyze(
+        newKtDeclaration: KtNamedFunction,
+        cache: ModuleFileCache,
+        firLazyDeclarationResolver: FirLazyDeclarationResolver,
+        firIdeProvider: FirIdeProvider,
+    ): ReanalyzableStructureElement<KT>
 
     fun isUpToDate(): Boolean = psi.getModificationStamp() == timestamp
 
     override val diagnostics: FileStructureElementDiagnostics by lazy {
-        var inCurrentDeclaration = false
+        FirIdeStructureElementDiagnosticsCollector.collectForSingleDeclaration(firFile, firSymbol.fir as FirDeclaration)
-
-        FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(
-            firFile,
-            onDeclarationEnter = { firDeclaration ->
-                when {
-                    firDeclaration == firSymbol.fir -> {
-                        inCurrentDeclaration = true
-                        DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
-                    }
-                    inCurrentDeclaration -> DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
-                    else -> DiagnosticCollectorDeclarationAction.SKIP_CURRENT_DECLARATION_AND_CHECK_NESTED
-                }
-            },
-            onDeclarationExit = { declaration ->
-                if (declaration == firSymbol.fir) {
-                    inCurrentDeclaration = false
-                }
-            }
-        )
     }
 
     companion object {
-        private val recorder = FirElementsRecorder()
+        val recorder = FirElementsRecorder()
+    }
+}
+
+internal class IncrementallyReanalyzableFunction(
+    override val firFile: FirFile,
+    override val psi: KtNamedFunction,
+    override val firSymbol: FirFunctionSymbol<*>,
+    override val timestamp: Long
+) : ReanalyzableStructureElement<KtNamedFunction>() {
+    override val mappings: Map<KtElement, FirElement> =
+        FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
+
+    private fun replaceFunction(from: FirSimpleFunction, to: FirSimpleFunction) {
+        val declarations = if (from.symbol.callableId.className == null) {
+            firFile.declarations as MutableList<FirDeclaration>
+        } else {
+            val classLikeLookupTag = from.containingClass()
+                ?: error("Class name should not be null for non-top-level & non-local declarations")
+            val containingClass = classLikeLookupTag.toSymbol(firFile.session)?.fir as FirRegularClass
+            containingClass.declarations as MutableList<FirDeclaration>
+        }
+        declarations.replaceFirst(from, to)
+    }
+
+    override fun reanalyze(
+        newKtDeclaration: KtNamedFunction,
+        cache: ModuleFileCache,
+        firLazyDeclarationResolver: FirLazyDeclarationResolver,
+        firIdeProvider: FirIdeProvider,
+    ): IncrementallyReanalyzableFunction {
+        val newFunction = firIdeProvider.buildFunctionWithBody(newKtDeclaration) as FirSimpleFunction
+        val originalFunction = firSymbol.fir as FirSimpleFunction
+
+        cache.firFileLockProvider.withWriteLock(firFile) {
+            replaceFunction(originalFunction, newFunction)
+        }
+
+        //todo remap symbol under firFile write lock
+        try {
+            firLazyDeclarationResolver.lazyResolveDeclaration(
+                newFunction,
+                cache,
+                FirResolvePhase.BODY_RESOLVE,
+                checkPCE = true,
+                reresolveFile = true,
+            )
+            return cache.firFileLockProvider.withReadLock(firFile) {
+                IncrementallyReanalyzableFunction(
+                    firFile,
+                    newKtDeclaration,
+                    newFunction.symbol,
+                    newKtDeclaration.modificationStamp,
+                )
+            }
+        } catch (e: Throwable) {
+            cache.firFileLockProvider.withWriteLock(firFile) {
+                replaceFunction(newFunction, originalFunction)
+            }
+            throw e
+        }
     }
 }
 
@@ -92,7 +142,7 @@ internal class NonLocalDeclarationFileStructureElement(
                         inCurrentDeclaration = true
                         DiagnosticCollectorDeclarationAction.CHECK_CURRENT_DECLARATION_AND_CHECK_NESTED
                     }
-                    (firDeclaration.psi as? KtNamedFunction)?.hasExplicitTypeOrUnit == true -> {
+                    FileElementFactory.isReanalyzableContainer(firDeclaration.ktDeclaration) -> {
                         DiagnosticCollectorDeclarationAction.SKIP
                     }
                     inCurrentDeclaration -> {
@@ -113,7 +163,7 @@ internal class NonLocalDeclarationFileStructureElement(
         private val recorder = object : FirElementsRecorder() {
             override fun visitSimpleFunction(simpleFunction: FirSimpleFunction, data: MutableMap<KtElement, FirElement>) {
                 val psi = simpleFunction.psi as? KtNamedFunction ?: return super.visitSimpleFunction(simpleFunction, data)
-                if (!psi.hasExplicitTypeOrUnit || KtPsiUtil.isLocal(psi)) {
+                if (!FileElementFactory.isReanalyzableContainer(psi) || KtPsiUtil.isLocal(psi)) {
                     super.visitSimpleFunction(simpleFunction, data)
                 }
             }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/psiUtils.kt

@@ -1,11 +0,0 @@
-/*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.idea.fir.low.level.api.util
-
-import org.jetbrains.kotlin.psi.KtNamedFunction
-
-internal val KtNamedFunction.hasExplicitTypeOrUnit
-    get() = hasBlockBody() || typeReference != null