FIR IDE: implement incremental function analysis

compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt

@@ -281,14 +281,19 @@ inline fun <reified T : PsiElement> PsiElement.collectDescendantsOfType(noinline
 inline fun <reified T : PsiElement> PsiElement.collectDescendantsOfType(
     crossinline canGoInside: (PsiElement) -> Boolean,
     noinline predicate: (T) -> Boolean = { true }
-): List<T> {
-    val result = ArrayList<T>()
+): List<T> = collectDescendantsOfTypeTo(ArrayList(), canGoInside, predicate)
+
+inline fun <reified T : PsiElement, C : MutableCollection<T>> PsiElement.collectDescendantsOfTypeTo(
+    to: C,
+    crossinline canGoInside: (PsiElement) -> Boolean,
+    noinline predicate: (T) -> Boolean = { true }
+): C {
     forEachDescendantOfType<T>(canGoInside) {
         if (predicate(it)) {
-            result.add(it)
+            to.add(it)
         }
     }
-    return result
+    return to
 }
 
 // ----------- Working with offsets, ranges and texts ----------------------------------------------------------------------------------------------

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

@@ -14,6 +14,8 @@ import com.intellij.psi.util.CachedValuesManager
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.LibraryModificationTracker
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.caches.trackers.KotlinCodeBlockModificationListener
+import org.jetbrains.kotlin.idea.caches.trackers.KotlinModuleOutOfCodeBlockModificationTracker
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeDependentModulesSourcesSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeLibrariesSession
@@ -25,7 +27,7 @@ import org.jetbrains.kotlin.idea.util.psiModificationTrackerBasedCachedValue
 import java.util.concurrent.ConcurrentHashMap
 
 internal class FirIdeResolveStateService(private val project: Project) {
-    private val stateCache by psiModificationTrackerBasedCachedValue(project) {
+    private val stateCache by cachedValue(project, KotlinCodeBlockModificationListener.getInstance(project).kotlinOutOfCodeBlockTracker) {
         ConcurrentHashMap<IdeaModuleInfo, FirModuleResolveStateImpl>()
     }
 
@@ -35,7 +37,6 @@ internal class FirIdeResolveStateService(private val project: Project) {
         ConcurrentHashMap<IdeaModuleInfo, FirIdeLibrariesSession>()
     }
 
-
     private fun createResolveStateFor(moduleInfo: IdeaModuleInfo): FirModuleResolveStateImpl {
         require(moduleInfo is ModuleSourceInfo)
         val firPhaseRunner = FirPhaseRunner()

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

@@ -11,17 +11,16 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.DiagnosticsCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirElementBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.PsiToFirCache
 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.file.structure.FileStructureCache
 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.sessions.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeCurrentModuleSourcesSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeDependentModulesSourcesSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeLibrariesSession
@@ -45,16 +44,13 @@ abstract class FirModuleResolveState {
     // todo temporary, used only in completion
     abstract fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile)
 
-    // todo temporary, used only in completion
-    abstract fun getCachedMappingForCompletion(element: KtElement): FirElement?
-
     abstract fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D
 
     // todo temporary, used only in completion
-    internal abstract fun lazyResolveDeclarationForCompletion(
+    abstract fun lazyResolveDeclarationForCompletion(
         firFunction: FirDeclaration,
         containerFirFile: FirFile,
-        firIdeProvider: FirIdeProvider,
+        firIdeProvider: FirProvider,
         toPhase: FirResolvePhase,
         towerDataContextCollector: FirTowerDataContextCollector
     )
@@ -71,22 +67,21 @@ internal class FirModuleResolveStateImpl(
     val firFileBuilder: FirFileBuilder,
     val firLazyDeclarationResolver: FirLazyDeclarationResolver,
 ) : FirModuleResolveState() {
-    val psiToFirCache = PsiToFirCache(currentModuleSourcesSession.cache)
-    val elementBuilder = FirElementBuilder(firFileBuilder, firLazyDeclarationResolver)
-    private val diagnosticsCollector =
-        DiagnosticsCollector(firFileBuilder, elementBuilder, psiToFirCache, currentModuleSourcesSession.cache)
+    val fileStructureCache = FileStructureCache(firFileBuilder, firLazyDeclarationResolver)
+    val elementBuilder = FirElementBuilder()
+    private val diagnosticsCollector = DiagnosticsCollector(fileStructureCache, currentModuleSourcesSession.cache)
 
     override fun getSessionFor(moduleInfo: IdeaModuleInfo): FirSession =
         sessionProvider.getSession(moduleInfo)
 
     override fun getOrBuildFirFor(element: KtElement, toPhase: FirResolvePhase): FirElement =
-        elementBuilder.getOrBuildFirFor(element, currentModuleSourcesSession.cache, psiToFirCache, toPhase)
+        elementBuilder.getOrBuildFirFor(element, currentModuleSourcesSession.cache, fileStructureCache)
 
     override fun getDiagnostics(element: KtElement): List<Diagnostic> =
         diagnosticsCollector.getDiagnosticsFor(element)
 
     override fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile) {
-        psiToFirCache.recordElementsForCompletionFrom(fir, firFile, ktFile)
+        error("Should be called only from FirModuleResolveStateForCompletion")
     }
 
     override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
@@ -98,13 +93,10 @@ internal class FirModuleResolveStateImpl(
         return declaration
     }
 
-    override fun getCachedMappingForCompletion(element: KtElement): FirElement? =
-        psiToFirCache.getCachedMapping(element)
-
     override fun lazyResolveDeclarationForCompletion(
         firFunction: FirDeclaration,
         containerFirFile: FirFile,
-        firIdeProvider: FirIdeProvider,
+        firIdeProvider: FirProvider,
         toPhase: FirResolvePhase,
         towerDataContextCollector: FirTowerDataContextCollector
     ) {
@@ -113,6 +105,7 @@ internal class FirModuleResolveStateImpl(
             currentModuleSourcesSession.cache,
             containerFirFile,
             firIdeProvider,
+            fromPhase = firFunction.resolvePhase,
             toPhase,
             towerDataContextCollector,
             checkPCE = false

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

@@ -11,9 +11,11 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.PsiToFirCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeCurrentModuleSourcesSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeDependentModulesSourcesSession
@@ -29,29 +31,24 @@ internal class FirModuleResolveStateForCompletion(
     override val dependentModulesSourcesSession: FirIdeDependentModulesSourcesSession get() = originalState.dependentModulesSourcesSession
     override val librariesSession: FirIdeLibrariesSession get() = originalState.librariesSession
 
-    private val psiToFirCache = PsiToFirCache(originalState.currentModuleSourcesSession.cache)
+    private val fileStructureCache = originalState.fileStructureCache
+
+    private val completionMapping = mutableMapOf<KtElement, FirElement>()
 
     override fun getSessionFor(moduleInfo: IdeaModuleInfo): FirSession =
         originalState.getSessionFor(moduleInfo)
 
     override fun getOrBuildFirFor(element: KtElement, toPhase: FirResolvePhase): FirElement {
-        getCachedMappingForCompletion(element)?.let { return it }
+        completionMapping[element]?.let { return it }
         return originalState.elementBuilder.getOrBuildFirFor(
             element,
             originalState.currentModuleSourcesSession.cache,
-            psiToFirCache,
-            toPhase
+            fileStructureCache,
         )
     }
 
     override fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile) {
-        psiToFirCache.recordElementsForCompletionFrom(fir, firFile, ktFile)
-    }
-
-    override fun getCachedMappingForCompletion(element: KtElement): FirElement? {
-        psiToFirCache.getCachedMapping(element)?.let { return it }
-        originalState.psiToFirCache.getCachedMapping(element)?.let { return it }
-        return null
+        fir.accept(FirElementsRecorder(), completionMapping)
     }
 
     override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
@@ -61,7 +58,7 @@ internal class FirModuleResolveStateForCompletion(
     override fun lazyResolveDeclarationForCompletion(
         firFunction: FirDeclaration,
         containerFirFile: FirFile,
-        firIdeProvider: FirIdeProvider,
+        firIdeProvider: FirProvider,
         toPhase: FirResolvePhase,
         towerDataContextCollector: FirTowerDataContextCollector
     ) {

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

@@ -0,0 +1,35 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
+import org.jetbrains.kotlin.fir.resolve.transformers.createTransformerBasedProcessorByPhase
+import java.util.*
+
+class FirTransformerProvider(session: FirSession) : FirSessionComponent {
+    private val scopeSession = ThreadLocal.withInitial { ScopeSession() }
+    private val transformers = ThreadLocal.withInitial {
+        val scopesSession = scopeSession.get()
+        EnumMap<FirResolvePhase, FirTransformerBasedResolveProcessor>(FirResolvePhase::class.java).apply {
+            FirResolvePhase.values().forEach { resolvePhase ->
+                if (resolvePhase != FirResolvePhase.RAW_FIR) {
+                    put(resolvePhase, resolvePhase.createTransformerBasedProcessorByPhase(session, scopesSession))
+                }
+            }
+        }
+    }
+
+    fun getTransformerForPhase(phase: FirResolvePhase): FirTransformerBasedResolveProcessor =
+        transformers.get().getValue(phase)
+
+    fun getScopeSession(): ScopeSession = scopeSession.get()
+}
+
+val FirSession.firTransformerProvider: FirTransformerProvider by FirSession.sessionComponentAccessor()

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

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyAccessorCopy
 import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunctionCopy
 import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyCopy
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
@@ -43,16 +42,11 @@ object LowLevelFirApiFacade {
     class FirCompletionContext internal constructor(
         val session: FirSession,
         private val towerDataContextCollector: FirTowerDataContextCollector,
-        private val state: FirModuleResolveState,
     ) {
         fun getTowerDataContext(element: KtElement): FirTowerDataContext {
             var current: PsiElement? = element
             while (current is KtElement) {
-                val mappedFir = state.getCachedMappingForCompletion(current)
-
-                if (mappedFir is FirStatement) {
-                    towerDataContextCollector.getContext(mappedFir)?.let { return it }
-                }
+                towerDataContextCollector.getContext(current)?.let { return it }
                 current = current.parent
             }
 
@@ -78,8 +72,7 @@ object LowLevelFirApiFacade {
 
         return FirCompletionContext(
             copyFunction.session,
-            contextCollector,
-            state
+            contextCollector
         )
     }
 
@@ -101,8 +94,7 @@ object LowLevelFirApiFacade {
 
         return FirCompletionContext(
             copyProperty.session,
-            contextCollector,
-            state
+            contextCollector
         )
     }
 
@@ -145,7 +137,7 @@ object LowLevelFirApiFacade {
                 symbol = builtSetter.symbol
                 resolvePhase = minOf(builtSetter.resolvePhase, FirResolvePhase.DECLARATIONS)
                 source = builtSetter.source
-                session = state.firIdeSourcesSession
+                session = state.currentModuleSourcesSession
             }
         } else {
             builtSetter
@@ -159,7 +151,7 @@ object LowLevelFirApiFacade {
 
             resolvePhase = minOf(originalProperty.resolvePhase, FirResolvePhase.DECLARATIONS)
             source = builtProperty.source
-            session = state.firIdeSourcesSession
+            session = state.currentModuleSourcesSession
         }
     }
 

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

@@ -5,57 +5,68 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
 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.util.FirElementFinder
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.findNonLocalFirDeclaration
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureCache
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
 /**
  * Maps [KtElement] to [FirElement]
- * Stateless, caches everything into [ModuleFileCache] & [PsiToFirCache] passed into the function
+ * Stateless, caches everything into [ModuleFileCache] & [FileStructureCache] passed into the function
  */
 @ThreadSafe
-internal class FirElementBuilder(
-    private val firFileBuilder: FirFileBuilder,
-    private val firLazyDeclarationResolver: FirLazyDeclarationResolver,
-) {
+internal class FirElementBuilder {
     fun getOrBuildFirFor(
         element: KtElement,
         moduleFileCache: ModuleFileCache,
-        psiToFirCache: PsiToFirCache,
-        toPhase: FirResolvePhase,
+        fileStructureCache: FileStructureCache,
     ): FirElement {
-        val ktFile = element.containingKtFile
-        val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
+        val fileStructure = fileStructureCache.getFileStructure(element.containingKtFile, moduleFileCache)
+        val mappings = fileStructure.getStructureElementFor(element).mappings
 
-        val containerFir = when (val container = element.getNonLocalContainingDeclarationWithFqName()) {
-            is KtDeclaration -> container.findNonLocalFirDeclaration(firFileBuilder, firFile.session.firProvider, moduleFileCache)
-            null -> firFile
-            else -> error("Unsupported: ${container.text}")
+        val psi = when {
+            element is KtPropertyDelegate -> element.expression ?: element
+            element is KtQualifiedExpression && element.selectorExpression is KtCallExpression -> {
+                /*
+                 KtQualifiedExpression with KtCallExpression in selector transformed in FIR to FirFunctionCall expression
+                 Which will have a receiver as qualifier
+                 */
+                element.selectorExpression ?: error("Incomplete code:\n${element.getElementTextInContext()}")
+            }
+            else -> element
         }
-        firLazyDeclarationResolver.lazyResolveDeclaration(containerFir, moduleFileCache, toPhase, checkPCE = true)
-
-        return psiToFirCache.getFir(element, containerFir, firFile)
+        mappings[psi]?.let { return it }
+        return psi.getFirOfClosestParent(mappings)?.second
+            ?: error("FirElement is not found for:\n${element.getElementTextInContext()}")
     }
 }
 
+private fun KtElement.getFirOfClosestParent(cache: Map<KtElement, FirElement>): Pair<KtElement, FirElement>? {
+    var current: PsiElement? = this
+    while (current is KtElement) {
+        val mappedFir = cache[current]
+        if (mappedFir != null) {
+            return current to mappedFir
+        }
+        current = current.parent
+    }
 
-fun KtElement.getNonLocalContainingDeclarationWithFqName(): KtNamedDeclaration? {
-    var container = parent
+    return null
+}
+
+fun KtElement.getNonLocalContainingOrThisDeclaration(): KtNamedDeclaration? {
+    var container: PsiElement? = this
     while (container != null && container !is KtFile) {
         if (container is KtNamedDeclaration
-            && (container is KtClassOrObject || container is KtDeclarationWithBody)
+            && (container is KtClassOrObject || container is KtDeclarationWithBody || container is KtProperty || container is KtTypeAlias)
             && !KtPsiUtil.isLocal(container)
-            && container.name != null
             && container !is KtEnumEntry
             && container.containingClassOrObject !is KtEnumEntry
         ) {

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirTowerDataContextCollector.kt

@@ -5,31 +5,22 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder
 
-import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.psi.KtElement
 
-internal class FirTowerDataContextCollector {
-    private val state: MutableMap<FirElement, FirTowerDataContext> = mutableMapOf()
+class FirTowerDataContextCollector {
+    private val state: MutableMap<KtElement, FirTowerDataContext> = mutableMapOf()
 
     fun addStatementContext(statement: FirStatement, context: FirTowerDataContext) {
-        state[statement] = context
+        (statement.psi as? KtElement)?.let { state[it] = context }
     }
 
     fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext) {
-        state[declaration] = context
+        (declaration.psi as? KtElement)?.let { state[it] = context }
     }
 
-    fun getContext(statement: FirElement): FirTowerDataContext? {
-        this.state[statement]?.let { return it }
-        for ((key, value) in this.state) {
-            // TODO Rework that
-            if (key.psi == statement.psi) {
-                return value
-            }
-        }
-        return null
-    }
+    fun getContext(psi: KtElement): FirTowerDataContext? = state[psi]
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/PsiToFirCache.kt

@@ -1,160 +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.element.builder
-
-import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.fir.realPsi
-import org.jetbrains.kotlin.fir.references.*
-import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
-import org.jetbrains.kotlin.fir.types.FirTypeRef
-import org.jetbrains.kotlin.fir.types.FirUserTypeRef
-import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
-import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
-import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.isErrorElement
-import java.util.concurrent.ConcurrentHashMap
-import org.jetbrains.kotlin.idea.util.getElementTextInContext
-import org.jetbrains.kotlin.psi.*
-
-/**
- * Belongs to a [org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState]
- */
-internal class PsiToFirCache(private val moduleFileCache: ModuleFileCache) {
-    private val caches = ConcurrentHashMap<KtFile, FileCache>()
-
-    fun getCachedMapping(element: KtElement): FirElement? {
-        val ktFile = element.containingKtFile
-        val cache = caches[ktFile]
-        return cache?.getCachedMapping(element)
-    }
-
-    fun getFir(element: KtElement, containerFir: FirDeclaration, firFile: FirFile): FirElement {
-        val ktFile = element.containingKtFile
-        val cache = caches.getOrPut(ktFile) { FileCache(ktFile, firFile, moduleFileCache) }
-        return cache.getFir(element, containerFir)
-    }
-
-    //todo for completion only
-    fun recordElementsForCompletionFrom(containerFir: FirDeclaration, firFile: FirFile, ktFile: KtFile) {
-        val cache = caches.getOrPut(ktFile) { FileCache(ktFile, firFile, moduleFileCache) }
-        cache.recordElementsForCompletionFrom(containerFir)
-    }
-}
-
-internal class FileCache(val ktFile: KtFile, firFile: FirFile, moduleFileCache: ModuleFileCache) {
-    private val cache: ConcurrentHashMap<KtElement, FirElement> = ConcurrentHashMap()
-
-    fun getCachedMapping(ktElement: KtElement): FirElement? {
-        require(ktElement.containingKtFile === ktFile)
-        return cache[ktElement]
-    }
-
-    fun getFir(element: KtElement, containerFir: FirDeclaration): FirElement {
-        val psi = when {
-            element is KtPropertyDelegate -> element.expression ?: element
-            element is KtQualifiedExpression && element.selectorExpression is KtCallExpression -> {
-                /*
-                 KtQualifiedExpression with KtCallExpression in selector transformed in FIR to FirFunctionCall expression
-                 Which will have a receiver as qualifier
-                 */
-                element.selectorExpression ?: error("Incomplete code:\n${element.getElementTextInContext()}")
-            }
-            else -> element
-        }
-        cache[psi]?.let { return it }
-        recordElementsFrom(containerFir)
-
-        val (current, mappedFir) = psi.getFirOfClosestParent()
-            ?: error("FirElement is not found for:\n${element.getElementTextInContext()}")
-        if (current !== element) {
-            cache(current, mappedFir)
-        }
-
-        return mappedFir
-    }
-
-    fun cache(psi: KtElement, fir: FirElement) {
-        // todo make it thread safe
-        val existingFir = cache[psi]
-        if (existingFir != null && existingFir !== fir) {
-            when {
-                existingFir is FirTypeRef && fir is FirTypeRef && psi is KtTypeReference -> {
-                    // FirTypeRefs are often created during resolve
-                    // a lot of them with have the same source
-                    // we want to take the most "resolved one" here
-                    if (fir is FirResolvedTypeRefImpl && existingFir !is FirResolvedTypeRefImpl) {
-                        cache[psi] = fir
-                    }
-                }
-                existingFir.isErrorElement && !fir.isErrorElement -> {
-                    // TODO better handle error elements
-                    // but for now just take first non-error one if such exist
-                    cache[psi] = fir
-                }
-                existingFir.isErrorElement || fir.isErrorElement -> {
-                    // do nothing and maybe upgrade to a non-error element in the branch above in the future
-                }
-                else -> {
-                    if (DuplicatedFirSourceElementsException.IS_ENABLED) {
-                        throw DuplicatedFirSourceElementsException(existingFir, fir, psi)
-                    }
-                }
-            }
-        }
-        if (existingFir == null) {
-            cache[psi] = fir
-        }
-    }
-
-    fun recordElementsForCompletionFrom(firDeclaration: FirDeclaration) {
-        recordElementsFrom(firDeclaration)
-    }
-
-    private fun recordElementsFrom(firDeclaration: FirDeclaration) {
-        firDeclaration.accept(elementsRecorder)
-    }
-
-    private val elementsRecorder = object : FirVisitorVoid() {
-        override fun visitElement(element: FirElement) {
-            (element.realPsi as? KtElement)?.let { psi ->
-                cache(psi, element)
-            }
-            element.acceptChildren(this)
-        }
-
-        override fun visitReference(reference: FirReference) {}
-        override fun visitControlFlowGraphReference(controlFlowGraphReference: FirControlFlowGraphReference) {}
-        override fun visitNamedReference(namedReference: FirNamedReference) {}
-        override fun visitResolvedNamedReference(resolvedNamedReference: FirResolvedNamedReference) {}
-        override fun visitDelegateFieldReference(delegateFieldReference: FirDelegateFieldReference) {}
-        override fun visitBackingFieldReference(backingFieldReference: FirBackingFieldReference) {}
-        override fun visitSuperReference(superReference: FirSuperReference) {}
-        override fun visitThisReference(thisReference: FirThisReference) {}
-        override fun visitErrorTypeRef(errorTypeRef: FirErrorTypeRef) {}
-
-        override fun visitUserTypeRef(userTypeRef: FirUserTypeRef) {
-            userTypeRef.acceptChildren(this)
-        }
-    }
-
-    private fun KtElement.getFirOfClosestParent(): Pair<KtElement, FirElement>? {
-        var current: PsiElement? = this
-        while (current is KtElement) {
-            val mappedFir = cache[current]
-            if (mappedFir != null) {
-                return current to mappedFir
-            }
-            current = current.parent
-        }
-
-        return null
-    }
-}
-

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

@@ -0,0 +1,165 @@
+/*
+ * 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.diagnostics.Diagnostic
+import org.jetbrains.kotlin.fir.declarations.*
+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.findNonLocalFirDeclaration
+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
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
+import java.util.concurrent.ConcurrentHashMap
+
+
+internal class FileStructure(
+    private val ktFile: KtFile,
+    private val firFile: FirFile,
+    private val firLazyDeclarationResolver: FirLazyDeclarationResolver,
+    private val firFileBuilder: FirFileBuilder,
+    private val moduleFileCache: ModuleFileCache
+) {
+    private val firIdeProvider = firFile.session.firIdeProvider
+
+    private val structureElements = ConcurrentHashMap<KtAnnotated, FileStructureElement>()
+
+    fun getStructureElementFor(element: KtElement): FileStructureElement {
+        val container: KtAnnotated = element.getNonLocalContainingOrThisDeclaration() ?: element.containingKtFile
+        return getStructureElementForDeclaration(container)
+    }
+
+    private fun getStructureElementForDeclaration(declaration: KtAnnotated): FileStructureElement {
+        val structureElement = structureElements.compute(declaration) { _, structureElement ->
+            when {
+                structureElement == null -> createStructureElement(declaration)
+                structureElement is WithInBlockModificationFileStructureElement && !structureElement.isUpToDate() -> {
+                    createMappingsCopy(structureElement, declaration as KtNamedFunction)
+                }
+                else -> structureElement
+            }
+        }
+        return structureElement
+            ?: error("FileStructureElement for was not defined for \n${declaration.getElementTextInContext()}")
+    }
+
+    @OptIn(ExperimentalStdlibApi::class)
+    fun getAllDiagnosticsForFile(): Collection<Diagnostic> {
+        val containersForStructureElement = buildList {
+            add(ktFile)
+            ktFile.forEachDescendantOfType<KtDeclaration>(
+                canGoInside = { psi -> psi !is KtFunction && psi !is KtValVarKeywordOwner }
+            ) { declaration ->
+                if (declaration.isStructureElementContainer()) {
+                    add(declaration)
+                }
+            }
+        }
+        val structureElements = containersForStructureElement.map(::getStructureElementFor)
+        return buildList {
+            structureElements.forEach { it.diagnostics.forEach { diagnostics -> addAll(diagnostics) } }
+        }
+    }
+
+    private fun replaceFunction(from: FirSimpleFunction, to: FirSimpleFunction) {
+        val declarations = if (from.symbol.callableId.className == null) {
+            firFile.declarations as MutableList<FirDeclaration>
+        } else {
+            val classId = from.symbol.callableId.classId
+                ?: error("Class name should not be null for non-top-level & non-local declarations")
+            val containingClass = firIdeProvider.getFirClassifierByFqName(classId) 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
+        replaceFunction(originalFunction, newFunction)
+
+        try {
+            firLazyDeclarationResolver.lazyResolveDeclaration(
+                newFunction,
+                moduleFileCache,
+                FirResolvePhase.BODY_RESOLVE,
+                checkPCE = true,
+                reresolveFile = true,
+            )
+            return WithInBlockModificationFileStructureElement(
+                firFile,
+                containerKtFunction,
+                newFunction.symbol,
+                containerKtFunction.modificationStamp,
+            )
+        } catch (e: Throwable) {
+            replaceFunction(newFunction, originalFunction)
+            throw e
+        }
+    }
+
+    private fun createDeclarationStructure(declaration: KtDeclaration): FileStructureElement {
+        val firDeclaration = declaration.findNonLocalFirDeclaration(firFileBuilder, firIdeProvider, moduleFileCache)
+        firLazyDeclarationResolver.lazyResolveDeclaration(
+            firDeclaration,
+            moduleFileCache,
+            FirResolvePhase.BODY_RESOLVE,
+            checkPCE = true
+        )
+        return when {
+            declaration is KtNamedFunction && declaration.hasExplicitTypeOrUnit -> {
+                WithInBlockModificationFileStructureElement(
+                    firFile,
+                    declaration,
+                    (firDeclaration as FirSimpleFunction).symbol,
+                    declaration.modificationStamp,
+                )
+            }
+            else -> {
+                NonLocalDeclarationFileStructureElement(
+                    firFile,
+                    firDeclaration,
+                    declaration,
+                )
+            }
+        }
+    }
+
+    private fun createStructureElement(container: KtAnnotated): FileStructureElement = when (container) {
+        is KtFile -> {
+            val firFile = firFileBuilder.getFirFileResolvedToPhaseWithCaching(
+                container,
+                moduleFileCache,
+                FirResolvePhase.BODY_RESOLVE,
+                checkPCE = true
+            )
+            FileWithoutDeclarationsFileStructureElement(
+                firFile,
+                container,
+            )
+        }
+        is KtDeclaration -> createDeclarationStructure(container)
+        else -> error("Invalid container $container")
+    }
+}
+
+private fun KtDeclaration.isStructureElementContainer(): Boolean {
+    if (this !is KtClassOrObject && this !is KtDeclarationWithBody && this !is KtProperty && this !is KtTypeAlias) return false
+    if (this is KtEnumEntry) return false
+    if (containingClassOrObject is KtEnumEntry) return false
+    return !KtPsiUtil.isLocal(this)
+}
+
+

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

@@ -0,0 +1,27 @@
+/*
+ * 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.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.psi.KtFile
+import java.util.concurrent.ConcurrentHashMap
+
+/**
+ * Belongs to a [org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState]
+ */
+internal class FileStructureCache(
+    private val fileBuilder: FirFileBuilder,
+    private val firLazyDeclarationResolver: FirLazyDeclarationResolver,
+) {
+    private val cache = ConcurrentHashMap<KtFile, FileStructure>()
+
+    fun getFileStructure(ktFile: KtFile, moduleFileCache: ModuleFileCache): FileStructure = cache.computeIfAbsent(ktFile) {
+        val firFile = fileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
+        FileStructure(ktFile, firFile, firLazyDeclarationResolver, fileBuilder, moduleFileCache)
+    }
+}

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

@@ -0,0 +1,119 @@
+/*
+ * 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.diagnostics.Diagnostic
+import org.jetbrains.kotlin.fir.FirElement
+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.fir.psi
+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.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>>
+) {
+    fun diagnosticsFor(element: KtElement): List<Diagnostic> = map[element] ?: emptyList()
+
+    inline fun forEach(action: (List<Diagnostic>) -> Unit) = map.values.forEach(action)
+}
+
+internal sealed class FileStructureElement {
+    abstract val firFile: FirFile
+    abstract val psi: KtAnnotated
+    abstract val mappings: Map<KtElement, FirElement>
+    abstract val diagnostics: FileStructureElementDiagnostics
+}
+
+internal class WithInBlockModificationFileStructureElement(
+    override val firFile: FirFile,
+    override val psi: KtFunction,
+    val firSymbol: FirFunctionSymbol<*>,
+    val timestamp: Long
+) : FileStructureElement() {
+
+    override val mappings: Map<KtElement, FirElement> =
+        FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
+
+
+    fun isUpToDate(): Boolean = psi.getModificationStamp() == timestamp
+
+    override val diagnostics: FileStructureElementDiagnostics by lazy {
+        val innerDeclarations = psi
+            .collectDescendantsOfTypeTo<KtDeclaration, HashSet<KtDeclaration>>(
+                hashSetOf(psi),
+                canGoInside = { true }
+            )
+        FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(
+            innerDeclarations,
+            firFile,
+            substitution = null
+        )
+    }
+
+    companion object {
+        private val recorder = FirElementsRecorder()
+    }
+}
+
+internal class NonLocalDeclarationFileStructureElement(
+    override val firFile: FirFile,
+    fir: FirDeclaration,
+    override val psi: KtDeclaration,
+) : FileStructureElement() {
+    override val mappings: Map<KtElement, FirElement> =
+        FirElementsRecorder.recordElementsFrom(fir, recorder)
+
+    override val diagnostics: FileStructureElementDiagnostics by lazy {
+        @Suppress("RemoveExplicitTypeArguments")
+        val innerDeclarations = psi
+            .collectDescendantsOfTypeTo<KtDeclaration, HashSet<KtDeclaration>>(
+                hashSetOf(psi),
+                canGoInside = { it.safeAs<KtNamedFunction>()?.hasExplicitTypeOrUnit != true }
+            )
+
+        FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(innerDeclarations, firFile)
+    }
+
+    companion object {
+        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)) {
+                    super.visitSimpleFunction(simpleFunction, data)
+                }
+            }
+        }
+    }
+}
+
+
+internal data class FileWithoutDeclarationsFileStructureElement(
+    override val firFile: FirFile,
+    override val psi: KtFile,
+) : FileStructureElement() {
+    override val mappings: Map<KtElement, FirElement> =
+        FirElementsRecorder.recordElementsFrom(firFile, recorder)
+
+    override val diagnostics: FileStructureElementDiagnostics by lazy {
+        FirIdeStructureElementDiagnosticsCollector.collectForStructureElement(hashSetOf(psi), firFile)
+    }
+
+    companion object {
+        private val recorder = object : FirElementsRecorder() {
+            override fun visitElement(element: FirElement, data: MutableMap<KtElement, FirElement>) {
+                if (element !is FirDeclaration || element is FirFile) {
+                    super.visitElement(element, data)
+                }
+            }
+        }
+    }
+}

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

@@ -0,0 +1,83 @@
+/*
+ * 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.FirElement
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.realPsi
+import org.jetbrains.kotlin.fir.references.*
+import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.FirUserTypeRef
+import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.DuplicatedFirSourceElementsException
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.isErrorElement
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtTypeReference
+
+internal open class FirElementsRecorder : FirVisitor<Unit, MutableMap<KtElement, FirElement>>() {
+    private fun cache(psi: KtElement, fir: FirElement, cache: MutableMap<KtElement, FirElement>) {
+        val existingFir = cache[psi]
+        if (existingFir != null && existingFir !== fir) {
+            when {
+                existingFir is FirTypeRef && fir is FirTypeRef && psi is KtTypeReference -> {
+                    // FirTypeRefs are often created during resolve
+                    // a lot of them with have the same source
+                    // we want to take the most "resolved one" here
+                    if (fir is FirResolvedTypeRefImpl && existingFir !is FirResolvedTypeRefImpl) {
+                        cache[psi] = fir
+                    }
+                }
+                existingFir.isErrorElement && !fir.isErrorElement -> {
+                    // TODO better handle error elements
+                    // but for now just take first non-error one if such exist
+                    cache[psi] = fir
+                }
+                existingFir.isErrorElement || fir.isErrorElement -> {
+                    // do nothing and maybe upgrade to a non-error element in the branch above in the future
+                }
+                else -> {
+                    if (DuplicatedFirSourceElementsException.IS_ENABLED) {
+                        throw DuplicatedFirSourceElementsException(existingFir, fir, psi)
+                    }
+                }
+            }
+        }
+        if (existingFir == null) {
+            cache[psi] = fir
+        }
+    }
+
+    override fun visitElement(element: FirElement, data: MutableMap<KtElement, FirElement>) {
+        (element.realPsi as? KtElement)?.let { psi ->
+            cache(psi, element, data)
+        }
+        element.acceptChildren(this, data)
+    }
+
+    //@formatter:off
+    override fun visitReference(reference: FirReference, data: MutableMap<KtElement, FirElement>) {}
+    override fun visitControlFlowGraphReference(controlFlowGraphReference: FirControlFlowGraphReference, data: MutableMap<KtElement, FirElement>) {}
+    override fun visitNamedReference(namedReference: FirNamedReference, data: MutableMap<KtElement, FirElement>) {}
+    override fun visitResolvedNamedReference(resolvedNamedReference: FirResolvedNamedReference, data: MutableMap<KtElement, FirElement>) {}
+    override fun visitDelegateFieldReference(delegateFieldReference: FirDelegateFieldReference, data: MutableMap<KtElement, FirElement>) {}
+    override fun visitBackingFieldReference(backingFieldReference: FirBackingFieldReference, data: MutableMap<KtElement, FirElement>) {}
+    override fun visitSuperReference(superReference: FirSuperReference, data: MutableMap<KtElement, FirElement>) {}
+    override fun visitThisReference(thisReference: FirThisReference, data: MutableMap<KtElement, FirElement>) {}
+    override fun visitErrorTypeRef(errorTypeRef: FirErrorTypeRef, data: MutableMap<KtElement, FirElement>) {}
+    //@formatter:on
+
+    override fun visitUserTypeRef(userTypeRef: FirUserTypeRef, data: MutableMap<KtElement, FirElement>) {
+        userTypeRef.acceptChildren(this, data)
+    }
+
+    companion object {
+        @OptIn(ExperimentalStdlibApi::class)
+        fun recordElementsFrom(firDeclaration: FirDeclaration, recorder: FirElementsRecorder): Map<KtElement, FirElement> =
+            buildMap { firDeclaration.accept(recorder, this) }
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -13,7 +13,7 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirDesignatedBodyResolveTransformerForIDE
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingDeclarationWithFqName
+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.providers.firIdeProvider
@@ -31,12 +31,14 @@ internal class FirLazyDeclarationResolver(
         moduleFileCache: ModuleFileCache,
         toPhase: FirResolvePhase,
         towerDataContextCollector: FirTowerDataContextCollector? = null,
-        checkPCE: Boolean = false
+        checkPCE: Boolean = false,
+        reresolveFile: Boolean = false,
     ) {
         if (declaration.resolvePhase >= toPhase) return
         val firFile = declaration.getContainingFile()
             ?: error("FirFile was not found for\n${declaration.render()}")
         val provider = firFile.session.firIdeProvider
+        val fromPhase = if (reresolveFile) declaration.resolvePhase else minOf(firFile.resolvePhase, declaration.resolvePhase)
         if (checkPCE) {
             firFileBuilder.runCustomResolveWithPCECheck(firFile, moduleFileCache) {
                 runLazyResolveWithoutLock(
@@ -44,6 +46,7 @@ internal class FirLazyDeclarationResolver(
                     moduleFileCache,
                     firFile,
                     provider,
+                    fromPhase,
                     toPhase,
                     towerDataContextCollector,
                     checkPCE = true
@@ -57,6 +60,7 @@ internal class FirLazyDeclarationResolver(
                         moduleFileCache,
                         firFile,
                         provider,
+                        fromPhase,
                         toPhase,
                         towerDataContextCollector,
                         checkPCE = false
@@ -71,16 +75,17 @@ internal class FirLazyDeclarationResolver(
         moduleFileCache: ModuleFileCache,
         containerFirFile: FirFile,
         provider: FirProvider,
+        fromPhase: FirResolvePhase,
         toPhase: FirResolvePhase,
         towerDataContextCollector: FirTowerDataContextCollector? = null,
-        checkPCE: Boolean
+        checkPCE: Boolean,
     ) {
-        if (firDeclarationToResolve.resolvePhase >= toPhase) return
+        if (fromPhase >= toPhase) return
         val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
-        if (firDeclarationToResolve.resolvePhase < nonLazyPhase) {
+        if (fromPhase < nonLazyPhase) {
             firFileBuilder.runResolveWithoutLock(
                 containerFirFile,
-                fromPhase = firDeclarationToResolve.resolvePhase,
+                fromPhase = fromPhase,
                 toPhase = nonLazyPhase,
                 checkPCE = checkPCE
             )
@@ -106,7 +111,7 @@ internal class FirLazyDeclarationResolver(
                 is FirCallableDeclaration<*> -> {
                     nonLocalDeclarationToResolve.symbol.callableId.classId
                 }
-                is FirRegularClass -> {
+                is FirClassLikeDeclaration<*> -> {
                     nonLocalDeclarationToResolve.symbol.classId
                 }
                 else -> error("Unsupported: ${nonLocalDeclarationToResolve.render()}")
@@ -126,7 +131,7 @@ internal class FirLazyDeclarationResolver(
         val transformer = FirDesignatedBodyResolveTransformerForIDE(
             designation.iterator(), containerFirFile.session,
             scopeSession,
-            implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+            toPhase,
             towerDataContextCollector
         )
         executeWithoutPCE {
@@ -138,7 +143,7 @@ internal class FirLazyDeclarationResolver(
         if (this is FirFile) return this
         val ktDeclaration = psi as? KtDeclaration ?: error("FirDeclaration should have a PSI of type KtDeclaration")
         if (!KtPsiUtil.isLocal(ktDeclaration)) return this
-        val nonLocalPsi = ktDeclaration.getNonLocalContainingDeclarationWithFqName()
+        val nonLocalPsi = ktDeclaration.getNonLocalContainingOrThisDeclaration()
             ?: error("Container for local declaration cannot be null")
         return nonLocalPsi.findNonLocalFirDeclaration(firFileBuilder, provider, moduleFileCache)
     }

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

@@ -0,0 +1,19 @@
+/*
+ * 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
+
+@Suppress("NOTHING_TO_INLINE")
+internal inline fun <K, V> MutableMap<K, MutableList<V>>.addValueFor(element: K, value: V) {
+    getOrPut(element) { mutableListOf() } += value
+}
+
+internal fun <T> MutableList<T>.replaceFirst(from: T, to: T) {
+    val index = indexOf(from)
+    if (index < 0) {
+        error("$from was not found in $this")
+    }
+    set(index, to)
+}

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

@@ -7,11 +7,13 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.util
 
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 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.util.classIdIfNonLocal
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
@@ -25,27 +27,45 @@ internal fun KtDeclaration.findNonLocalFirDeclaration(
         this is KtClassOrObject -> findFir(provider)
         this is KtNamedDeclaration && (this is KtProperty || this is KtNamedFunction) -> {
             val containerClass = containingClassOrObject
-            if (containerClass != null) {
+            val declarations = if (containerClass != null) {
                 val containerClassFir = containerClass.findFir(provider)
-                containerClassFir.declarations.first { it.psi === this }
+                containerClassFir.declarations
             } else {
                 val ktFile = containingKtFile
                 val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
-                firFile.declarations.first { it.psi === this }
+                firFile.declarations
             }
+            val original = originalDeclaration
+            declarations.first { it.psi === this || it.psi == original }
         }
         this is KtConstructor<*> -> {
             val containerClassFir = containingClassOrObject?.findFir(provider)
                 ?: error("Container class should be not null for KtConstructor")
             containerClassFir.declarations.first { it.psi === this }
         }
-        else -> error("Invalid container $this")
+        this is KtTypeAlias -> findFir(provider)
+        else -> error("Invalid container $this::class")
     }
 }
 
+val ORIGINAL_DECLARATION_KEY = com.intellij.openapi.util.Key<KtDeclaration>("ORIGINAL_DECLARATION_KEY")
+
+var KtDeclaration.originalDeclaration by UserDataProperty(ORIGINAL_DECLARATION_KEY)
+
+
 private fun KtClassOrObject.findFir(provider: FirProvider): FirRegularClass {
-    val containerClassId = classIdIfNonLocal()
+    val classId = classIdIfNonLocal()
         ?: error("Container classId should not be null for non-local declaration")
-    return provider.getFirClassifierByFqName(containerClassId) as? FirRegularClass
-        ?: error("Could not find class $containerClassId")
+    return executeWithoutPCE {
+        provider.getFirClassifierByFqName(classId) as? FirRegularClass
+            ?: error("Could not find class $classId")
+    }
+}
+
+private fun KtTypeAlias.findFir(provider: FirProvider): FirTypeAlias {
+    val typeAlias = ClassId(containingKtFile.packageFqName, nameAsSafeName)
+    return executeWithoutPCE {
+        provider.getFirClassifierByFqName(typeAlias) as? FirTypeAlias
+            ?: error("Could not find type alias $typeAlias")
+    }
 }

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

@@ -0,0 +1,11 @@
+/*
+ * 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

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

@@ -52,10 +52,6 @@ internal inline fun checkCanceled() {
 internal val FirElement.isErrorElement
     get() = this is FirDiagnosticHolder
 
-@Suppress("NOTHING_TO_INLINE")
-internal inline fun <K, V> MutableMap<K, MutableList<V>>.addValueFor(element: K, value: V) {
-    getOrPut(element) { mutableListOf() } += value
-}
 
 internal fun IdeaModuleInfo.collectTransitiveDependenciesWithSelf(): List<IdeaModuleInfo> {
     val result = mutableSetOf<IdeaModuleInfo>()

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirScopeProvider.kt

@@ -128,7 +128,7 @@ internal class KtFirScopeProvider(
         originalFile: KtFile,
         positionInFakeFile: KtElement
     ): LowLevelFirApiFacade.FirCompletionContext {
-        val originalFirFile = originalFile.getOrBuildFirOfType<FirFile>(firResolveState)
+        val originalFirFile = firResolveState.getFirFile(originalFile)
         val declarationContext = EnclosingDeclarationContext.detect(originalFile, positionInFakeFile)
 
         return declarationContext.buildCompletionContext(originalFirFile, firResolveState)

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/EnclosingDeclarationContext.kt

@@ -9,6 +9,7 @@ import com.intellij.psi.util.parentsOfType
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
@@ -20,7 +21,7 @@ internal sealed class EnclosingDeclarationContext {
             if (fakeFunction != null) {
                 val originalFunction = originalFile.findDeclarationOfTypeAt<KtNamedFunction>(fakeFunction.textOffset)
                     ?: error("Cannot find original function matching to ${fakeFunction.getElementTextInContext()} in $originalFile")
-
+                fakeFunction.originalDeclaration = originalFunction
                 return FunctionContext(fakeFunction, originalFunction)
             }
 
@@ -28,6 +29,7 @@ internal sealed class EnclosingDeclarationContext {
             if (fakeProperty != null) {
                 val originalProperty = originalFile.findDeclarationOfTypeAt<KtProperty>(fakeProperty.textOffset)
                     ?: error("Cannot find original property matching to ${fakeProperty.getElementTextInContext()} in $originalFile")
+                fakeProperty.originalDeclaration = originalProperty
 
                 return PropertyContext(fakeProperty, originalProperty)
             }