FIR IDE: rewrite low level API

- Cache ModuleResolveState for module till the world changes
- Resolve every file under a lock
- All creation of raw fir files and resolve of them happens in FirFileBuilder
- Lazy resolve of fir elements happens in FirElementBuilder

Caching works like the following:
- FirModuleResolveState holds PsiToFirCache & DiagnosticsCollector & FileCacheForModuleProvider
- FileCacheForModuleProvider holds a mapping from ModuleInfo to ModuleFileCache
- ModuleFileCache caches
    - KtFile -> FirFile mapping
    - ClassId -> FirClassLikeDeclaration, CallableId -> FirCallableSymbol
        which used in corresponding FirProvider
    - mapping from declaration to it's file
        which used in corresponding FirProvider
    - locks for fir file resolving
- PsiToFirCache provides mapping from KtElement to  FirElement
- DiagnosticsCollector collects diagnostics for file and caches them

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.
  */
 
@@ -65,7 +65,7 @@ class FirJavaModuleBasedSession private constructor(
 
 
     init {
-        sessionProvider.sessionCache[moduleInfo] = this
+        sessionProvider.registerSession(moduleInfo, this)
     }
 }
 
@@ -118,14 +118,18 @@ class FirLibrarySession private constructor(
     }
 
     init {
-        sessionProvider.sessionCache[moduleInfo] = this
+        sessionProvider.registerSession(moduleInfo, this)
     }
 }
 
-class FirProjectSessionProvider(override val project: Project) : FirSessionProvider {
+open class FirProjectSessionProvider(override val project: Project) : FirSessionProvider {
     override fun getSession(moduleInfo: ModuleInfo): FirSession? {
         return sessionCache[moduleInfo]
     }
 
-    val sessionCache = mutableMapOf<ModuleInfo, FirSession>()
+    fun registerSession(moduleInfo: ModuleInfo, session: FirSession) {
+        sessionCache[moduleInfo] = session
+    }
+
+    protected open val sessionCache: MutableMap<ModuleInfo, FirSession> = mutableMapOf()
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/FirSymbolProvider.kt

@@ -55,8 +55,6 @@ abstract class FirSymbolProvider : FirSessionComponent {
     open fun getNestedClassesNamesInClass(classId: ClassId): Set<Name> = emptySet()
 
     abstract fun getPackage(fqName: FqName): FqName? // TODO: Replace to symbol sometime
-
-    open fun getAllCallableNamesInPackage(): Set<Name> = emptySet()
 }
 
 fun FirSession.getNestedClassifierScope(lookupTag: ConeClassLikeLookupTag): FirScope? =

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirCompositeSymbolProvider.kt

@@ -46,8 +46,4 @@ class FirCompositeSymbolProvider(val providers: List<FirSymbolProvider>) : FirSy
     override fun getNestedClassesNamesInClass(classId: ClassId): Set<Name> {
         return providers.flatMapTo(mutableSetOf()) { it.getNestedClassesNamesInClass(classId) }
     }
-
-    override fun getAllCallableNamesInPackage(): Set<Name> {
-        return providers.flatMapTo(mutableSetOf()) { it.getAllCallableNamesInPackage() }
-    }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirProviderImpl.kt

@@ -76,12 +76,6 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: KotlinSc
         klass.accept(FirRecorder, state to owner.file)
     }
 
-    override fun getAllCallableNamesInPackage(): Set<Name> {
-        return state.callableMap.keys.asSequence()
-            .filter { it.className == null }
-            .mapTo(mutableSetOf()) { it.callableName }
-    }
-
     private val FirAnnotatedDeclaration.file: FirFile
         get() = when (this) {
             is FirFile -> this
@@ -92,7 +86,6 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: KotlinSc
     private fun recordFile(file: FirFile, state: State) {
         val packageName = file.packageFqName
         state.fileMap.merge(packageName, listOf(file)) { a, b -> a + b }
-
         file.acceptChildren(FirRecorder, state to file)
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSupertypesResolution.kt

@@ -237,7 +237,7 @@ private class FirSupertypeResolverVisitor(
                 val outerClassFir = classId.outerClassId?.let(session.firProvider::getFirClassifierByFqName) as? FirRegularClass
                 prepareScopeForNestedClasses(outerClassFir ?: return ImmutableList.empty())
             }
-            else -> prepareFileScopes(session.firProvider.getFirClassifierContainerFile(classId))
+            else -> prepareFileScopes(session.firProvider.getFirClassifierContainerFile(classLikeDeclaration.symbol))
         }
 
         return result.pushIfNotNull(classLikeDeclaration.typeParametersScope())

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/PhaseUtils.kt

@@ -1,10 +1,11 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.fir.resolve.transformers
 
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSymbolOwner
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
@@ -33,7 +34,11 @@ fun <D> AbstractFirBasedSymbol<D>.phasedFir(
             else -> null
         }
             ?: throw AssertionError("Cannot get container file by symbol: $this (${result.render()})")
-        containingFile.runResolve(toPhase = requiredPhase, fromPhase = availablePhase)
+
+        val resolver = fir.session.phasedFirFileResolver
+            ?: error("phasedFirFileResolver should be defined when working with FIR in phased mode")
+        resolver.resolveFile(containingFile, fromPhase = availablePhase, toPhase = requiredPhase)
     }
     return result
 }
+

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/PhasedFirFileResolver.kt

@@ -0,0 +1,17 @@
+/*
+ * 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.fir.resolve.transformers
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+
+abstract class PhasedFirFileResolver : FirSessionComponent {
+    abstract fun resolveFile(firFile: FirFile, fromPhase: FirResolvePhase, toPhase: FirResolvePhase)
+}
+
+internal val FirSession.phasedFirFileResolver: PhasedFirFileResolver? by FirSession.nullableSessionComponentAccessor()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirPackageMemberScope.kt

@@ -58,6 +58,6 @@ class FirPackageMemberScope(val fqName: FqName, val session: FirSession) : FirSc
     }
 
     override fun getCallableNames(): Set<Name> {
-        return symbolProvider.getAllCallableNamesInPackage()
+        return symbolProvider.getAllCallableNamesInPackage(fqName)
     }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.
  */
 
@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.fir.types.impl.*
 import org.jetbrains.kotlin.fir.utils.ArrayMapAccessor
 import org.jetbrains.kotlin.fir.utils.ComponentArrayOwner
+import org.jetbrains.kotlin.fir.utils.NullableArrayMapAccessor
 import org.jetbrains.kotlin.fir.utils.TypeRegistry
 import org.jetbrains.kotlin.utils.Jsr305State
 import kotlin.reflect.KClass
@@ -21,6 +22,10 @@ abstract class FirSession(val sessionProvider: FirSessionProvider?) : ComponentA
         inline fun <reified T : FirSessionComponent> sessionComponentAccessor(): ArrayMapAccessor<FirSessionComponent, FirSessionComponent, T> {
             return generateAccessor(T::class)
         }
+
+        inline fun <reified T : FirSessionComponent> nullableSessionComponentAccessor(): NullableArrayMapAccessor<FirSessionComponent, FirSessionComponent, T> {
+            return generateNullableAccessor(T::class)
+        }
     }
 
     open val moduleInfo: ModuleInfo? get() = null

idea/idea-analysis/src/org/jetbrains/kotlin/idea/stubindex/KotlinFullClassNameIndex.java

@@ -1,17 +1,6 @@
 /*
- * Copyright 2010-2015 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
+ * 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.stubindex;
@@ -28,7 +17,7 @@ import org.jetbrains.kotlin.psi.KtClassOrObject;
 import java.util.Collection;
 
 public class KotlinFullClassNameIndex extends StringStubIndexExtension<KtClassOrObject> {
-    private static final StubIndexKey<String, KtClassOrObject> KEY = KotlinIndexUtil.createIndexKey(KotlinFullClassNameIndex.class);
+    public static final StubIndexKey<String, KtClassOrObject> KEY = KotlinIndexUtil.createIndexKey(KotlinFullClassNameIndex.class);
 
     private static final KotlinFullClassNameIndex ourInstance = new KotlinFullClassNameIndex();
 

idea/idea-fir/tests/org/jetbrains/kotlin/idea/ideaFirTestUtils.kt

@@ -5,7 +5,7 @@
 
 package org.jetbrains.kotlin.idea
 
-import org.jetbrains.kotlin.idea.fir.low.level.api.DuplicatedFirSourceElementsException
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.DuplicatedFirSourceElementsException
 
 fun Throwable.shouldBeRethrown(): Boolean = when (this) {
     is DuplicatedFirSourceElementsException -> true

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

@@ -1,71 +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
-
-import com.intellij.openapi.module.impl.scopes.ModuleScopeProviderImpl
-import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
-import org.jetbrains.kotlin.fir.scopes.FirScope
-import org.jetbrains.kotlin.fir.symbols.CallableId
-import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
-import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
-import org.jetbrains.kotlin.idea.stubindex.KotlinFullClassNameIndex
-import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtElement
-
-internal class FirIdeAllSourceDependenciesSymbolProvider(
-    project: Project,
-    moduleInfo: ModuleSourceInfo,
-    private val sessionProvider: FirIdeSessionProvider,
-    private val firBuilder: FirBuilder,
-) : FirSymbolProvider() {
-    private val scope = ModuleScopeProviderImpl(moduleInfo.module).moduleWithDependenciesScope
-    private val indexHelper = IndexHelper(project)
-    private val packageExistenceChecker = PackageExistenceCheckerForMultipleModules(project, moduleInfo.dependencies())
-
-    private val cache = PsiToFirCacheImpl()
-
-    override fun getClassLikeSymbolByFqName(classId: ClassId): FirClassLikeSymbol<*>? {
-        return executeOrReturnDefaultValueOnPCE(null) {
-            val ktClass = indexHelper.firstMatchingOrNull(
-                KotlinFullClassNameIndex.KEY,
-                classId.asSingleFqName().asString(),
-                scope
-            ) { candidate -> candidate.containingKtFile.packageFqName == classId.packageFqName } ?: return null
-            (ktClass.toFir() as FirClassLikeDeclaration<*>).symbol
-        }
-    }
-
-
-    @OptIn(ExperimentalStdlibApi::class)
-    override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
-        val callableId = CallableId(packageFqName, name)
-        return executeOrReturnDefaultValueOnPCE(emptyList()) {
-            buildList {
-                indexHelper.getTopLevelFunctions(callableId, scope).mapTo(this) { (it.toFir() as FirCallableDeclaration<*>).symbol }
-                indexHelper.getTopLevelProperties(callableId, scope).mapTo(this) { (it.toFir() as FirCallableDeclaration<*>).symbol }
-            }
-        }
-    }
-
-    private fun KtElement.toFir(): FirElement =
-        getOrBuildFir(cache, firBuilder, sessionProvider, FirResolvePhase.RAW_FIR)
-
-    override fun getNestedClassifierScope(classId: ClassId): FirScope? {
-        val classifier = getClassLikeSymbolByFqName(classId) ?: return null
-        return classifier.fir.session.firIdeProvider.getNestedClassifierScope(classId)
-    }
-
-    override fun getPackage(fqName: FqName): FqName? =
-        fqName.takeIf(packageExistenceChecker::isPackageExists)
-}

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

@@ -5,56 +5,89 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
+import com.intellij.openapi.module.impl.scopes.ModuleWithDependentsScope
 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analyzer.ModuleInfo
-import org.jetbrains.kotlin.fir.FirModuleBasedSession
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.registerCheckersComponent
-import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
+import org.jetbrains.kotlin.fir.extensions.BunchOfRegisteredExtensions
+import org.jetbrains.kotlin.fir.extensions.extensionService
+import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
-import org.jetbrains.kotlin.fir.registerCommonComponents
-import org.jetbrains.kotlin.fir.registerResolveComponents
 import org.jetbrains.kotlin.fir.resolve.calls.jvm.registerJvmCallConflictResolverFactory
-import org.jetbrains.kotlin.fir.resolve.firProvider
+import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCompositeSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
+import org.jetbrains.kotlin.fir.resolve.transformers.PhasedFirFileResolver
 import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
+import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCacheImpl
+import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeModuleLibraryDependenciesSession
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.collectTransitiveDependenciesWithSelf
 
 
-internal class FirIdeJavaModuleBasedSession(
+internal class FirIdeJavaModuleBasedSession private constructor(
     moduleInfo: ModuleInfo,
-    sessionProvider: FirProjectSessionProvider
+    sessionProvider: FirSessionProvider,
+    val firFileBuilder: FirFileBuilder,
 ) : FirModuleBasedSession(moduleInfo, sessionProvider) {
+    val cache get() = firIdeProvider.cache
+
     companion object {
+        /**
+         * Should be invoked only under a [moduleInfo]-based lock
+         */
         fun create(
             project: Project,
-            moduleInfo: ModuleInfo,
-            sessionProvider: FirProjectSessionProvider,
-            scope: GlobalSearchScope
+            moduleInfo: ModuleSourceInfo,
+            sessionProvider: FirSessionProvider,
         ): FirIdeJavaModuleBasedSession {
-            return FirIdeJavaModuleBasedSession(moduleInfo, sessionProvider).apply {
+            val scopeProvider = KotlinScopeProvider(::wrapScopeWithJvmMapped)
+            val firBuilder = FirFileBuilder(sessionProvider as FirIdeSessionProvider, scopeProvider)
+            return FirIdeJavaModuleBasedSession(moduleInfo, sessionProvider, firBuilder).apply {
+                val cache = ModuleFileCacheImpl(this)
+                val phasedFirFileResolver = IdePhasedFirFileResolver(firBuilder, cache)
+                val dependentModules = moduleInfo.collectTransitiveDependenciesWithSelf().filterIsInstance<ModuleSourceInfo>()
+                val searchScope = ModuleWithDependentsScope(project, dependentModules.map { it.module })
+
                 registerCommonComponents()
                 registerResolveComponents()
                 registerCheckersComponent()
-                registerJvmCallConflictResolverFactory()
 
-                val firIdeProvider = FirIdeProvider(project, scope, this, KotlinScopeProvider(::wrapScopeWithJvmMapped))
+                val provider = FirIdeProvider(
+                    project,
+                    this,
+                    moduleInfo,
+                    scopeProvider,
+                    firFileBuilder,
+                    cache,
+                    searchScope
+                )
 
-                register(FirProvider::class, firIdeProvider)
-                register(FirIdeProvider::class, firIdeProvider)
+                register(FirProvider::class, provider)
+                register(FirIdeProvider::class, provider)
+
+                register(PhasedFirFileResolver::class, phasedFirFileResolver)
 
                 register(
                     FirSymbolProvider::class,
                     FirCompositeSymbolProvider(
-                        listOf(
-                            firProvider,
-                            JavaSymbolProvider(this, sessionProvider.project, scope),
-                            FirIdeModuleDependenciesSymbolProvider(this)
-                        )
+                        @OptIn(ExperimentalStdlibApi::class)
+                        buildList {
+                            add(provider)
+                            add(JavaSymbolProvider(this@apply, sessionProvider.project, searchScope))
+                            add(FirIdeModuleLibraryDependenciesSession.create(moduleInfo, sessionProvider, project).firSymbolProvider)
+                        }
                     ) as FirSymbolProvider
                 )
+                registerJvmCallConflictResolverFactory()
+                extensionService.registerExtensions(BunchOfRegisteredExtensions.empty())
             }
         }
     }

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

@@ -14,18 +14,20 @@ import org.jetbrains.kotlin.fir.resolve.providers.impl.FirDependenciesSymbolProv
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.isLibraryClasses
 import org.jetbrains.kotlin.idea.caches.resolve.IDEPackagePartProvider
+import javax.annotation.concurrent.NotThreadSafe
 
 
+@NotThreadSafe
+//todo make thread safe
 internal class FirIdeModuleLibraryDependenciesSymbolProvider(
     session: FirIdeJavaModuleBasedSession
 ) : FirDependenciesSymbolProviderImpl(session) {
-    //todo invalidate cache on libraries changed
     override val dependencyProviders: List<FirSymbolProvider> by lazy {
         val moduleInfo = session.moduleInfo
         val sessionProvider = session.sessionProvider ?: error("No session provider found")
         val project = sessionProvider.project
 
-        moduleInfo.dependenciesWithoutSelf()
+        val dependencies = moduleInfo.dependenciesWithoutSelf()
             .filterIsInstance<IdeaModuleInfo>()
             .mapNotNull { dependencyInfo ->
                 val dependencySession = if (dependencyInfo.isLibraryClasses()) {
@@ -37,6 +39,7 @@ internal class FirIdeModuleLibraryDependenciesSymbolProvider(
                 } else null
                 dependencySession?.firSymbolProvider
             }.toList()
+        dependencies
     }
 }
 

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

@@ -1,176 +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
-
-import com.intellij.openapi.progress.ProcessCanceledException
-import com.intellij.openapi.project.Project
-import com.intellij.psi.search.GlobalSearchScope
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.builder.RawFirBuilder
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
-import org.jetbrains.kotlin.fir.resolve.providers.FirProviderInternals
-import org.jetbrains.kotlin.fir.resolve.providers.impl.FirProviderImpl
-import org.jetbrains.kotlin.fir.scopes.FirScope
-import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
-import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
-import org.jetbrains.kotlin.idea.stubindex.*
-import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtNamedFunction
-
-internal class FirIdeProvider(
-    val project: Project,
-    val scope: GlobalSearchScope,
-    val session: FirSession,
-    kotlinScopeProvider: KotlinScopeProvider
-) : FirProvider() {
-    private val cacheProvider = FirProviderImpl(session, kotlinScopeProvider)
-
-    data class FirFileWithStamp(val file: FirFile, val stamp: Long)
-
-    private val files = mutableMapOf<KtFile, FirFileWithStamp>()
-
-    override val isPhasedFirAllowed: Boolean
-        get() = true
-
-    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>? {
-        return cacheProvider.getFirClassifierByFqName(classId) ?: run {
-            try {
-                val classes = KotlinFullClassNameIndex.getInstance().get(classId.asSingleFqName().asString(), project, scope)
-                val ktClass = classes.firstOrNull {
-                    classId.packageFqName == it.containingKtFile.packageFqName
-                } ?: return null // TODO: what if two of them?
-
-                val ktFile = ktClass.containingKtFile
-                getOrBuildFile(ktFile)
-
-                cacheProvider.getFirClassifierByFqName(classId)
-            } catch (e: ProcessCanceledException) {
-                return null
-            }
-        }
-    }
-
-    fun getOrBuildFile(ktFile: KtFile): FirFile {
-        val modificationStamp = ktFile.modificationStamp
-        files[ktFile]?.let { (firFile, stamp) ->
-            if (stamp == modificationStamp) {
-                return firFile
-            }
-        }
-        return synchronized(ktFile) {
-            var fileWithStamp = files[ktFile]
-            if (fileWithStamp != null && fileWithStamp.stamp == modificationStamp) {
-                fileWithStamp.file
-            } else {
-                val file = RawFirBuilder(session, cacheProvider.kotlinScopeProvider, stubMode = false).buildFirFile(ktFile)
-                cacheProvider.recordFile(file)
-                fileWithStamp = FirFileWithStamp(file, modificationStamp)
-                files[ktFile] = fileWithStamp
-                file
-            }
-        }
-    }
-
-    fun buildFunctionWithBody(ktNamedFunction: KtNamedFunction): FirFunction<*> {
-        return RawFirBuilder(session, cacheProvider.kotlinScopeProvider, stubMode = false).buildFunctionWithBody(ktNamedFunction)
-    }
-
-    fun getFile(ktFile: KtFile): FirFile? {
-        val (firFile, stamp) = files[ktFile] ?: return null
-        if (stamp == ktFile.modificationStamp) return firFile
-        return null
-    }
-
-    override fun getClassLikeSymbolByFqName(classId: ClassId): FirClassLikeSymbol<*>? {
-        return getFirClassifierByFqName(classId)?.symbol
-    }
-
-    override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
-        val packagePrefix = if (packageFqName.isRoot) "" else "$packageFqName."
-        return try {
-            val topLevelFunctions = KotlinTopLevelFunctionFqnNameIndex.getInstance()["$packagePrefix$name", project, scope]
-            val topLevelProperties = KotlinTopLevelPropertyFqnNameIndex.getInstance()["$packagePrefix$name", project, scope]
-            topLevelFunctions.forEach { getOrBuildFile(it.containingKtFile) }
-            topLevelProperties.forEach { getOrBuildFile(it.containingKtFile) }
-            cacheProvider.getTopLevelCallableSymbols(packageFqName, name)
-        } catch (e: ProcessCanceledException) {
-            emptyList()
-        }
-    }
-
-    override fun getFirClassifierContainerFile(fqName: ClassId): FirFile {
-        getFirClassifierByFqName(fqName) // Necessary to ensure cacheProvider contains this classifier
-        return cacheProvider.getFirClassifierContainerFile(fqName)
-    }
-
-    override fun getFirClassifierContainerFileIfAny(fqName: ClassId): FirFile? {
-        getFirClassifierByFqName(fqName) // Necessary to ensure cacheProvider contains this classifier
-        return cacheProvider.getFirClassifierContainerFileIfAny(fqName)
-    }
-
-    override fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile {
-        return getFirClassifierContainerFileIfAny(symbol)
-            ?: error("Couldn't find container for ${symbol.classId}")
-    }
-
-    override fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile? {
-        val psi = symbol.fir.source?.psi
-        if (psi is KtElement) {
-            return try {
-                val ktFile = psi.containingKtFile
-                getOrBuildFile(ktFile)
-            } catch (e: ProcessCanceledException) {
-                null
-            }
-        }
-        return getFirClassifierContainerFileIfAny(symbol.classId)
-    }
-
-    override fun getFirCallableContainerFile(symbol: FirCallableSymbol<*>): FirFile? {
-        return cacheProvider.getFirCallableContainerFile(symbol)
-    }
-
-    override fun getFirFilesByPackage(fqName: FqName): List<FirFile> {
-        return try {
-            val files = KotlinExactPackagesIndex.getInstance()[fqName.asString(), project, scope]
-            files.forEach { getOrBuildFile(it) }
-            cacheProvider.getFirFilesByPackage(fqName)
-        } catch (e: ProcessCanceledException) {
-            emptyList()
-        }
-    }
-
-    override fun getNestedClassifierScope(classId: ClassId): FirScope? {
-        getFirClassifierByFqName(classId)
-        return cacheProvider.getNestedClassifierScope(classId)
-    }
-
-    @FirProviderInternals
-    override fun recordGeneratedClass(owner: FirAnnotatedDeclaration, klass: FirRegularClass) {
-        // TODO: check that this implementation is correct
-        cacheProvider.recordGeneratedClass(owner, klass)
-    }
-
-    @FirProviderInternals
-    override fun recordGeneratedMember(owner: FirAnnotatedDeclaration, klass: FirDeclaration) {
-        // TODO: check that this implementation is correct
-        cacheProvider.recordGeneratedMember(owner, klass)
-    }
-
-    // TODO this should be reworked because [FirIdeProvider] should not have such method
-    override fun getAllCallableNamesInPackage(): Set<Name> {
-        return cacheProvider.getAllCallableNamesInPackage()
-    }
-}
-
-internal val FirSession.firIdeProvider: FirIdeProvider by FirSession.sessionComponentAccessor()

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

@@ -5,59 +5,32 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
-import com.intellij.openapi.components.ServiceManager
+import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.ModificationTracker
-import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
-import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+import org.jetbrains.kotlin.idea.util.psiModificationTrackerBasedCachedValue
+import java.util.concurrent.ConcurrentHashMap
+import org.jetbrains.kotlin.idea.caches.project.*
+
+internal class FirIdeResolveStateService(private val project: Project) {
+    private val stateCache by psiModificationTrackerBasedCachedValue(project) {
+        ConcurrentHashMap<IdeaModuleInfo, FirModuleResolveStateImpl>()
+    }
+
+    private fun createResolveStateFor(moduleInfo: IdeaModuleInfo): FirModuleResolveStateImpl {
+        val sessionProvider = FirIdeSessionProvider(project)
+        val session = FirIdeJavaModuleBasedSession.create(
+            project,
+            moduleInfo as ModuleSourceInfo,
+            sessionProvider,
+        ).apply { sessionProvider.registerSession(moduleInfo, this) }
+
+        return FirModuleResolveStateImpl(moduleInfo, session, sessionProvider, session.firFileBuilder, session.cache)
+    }
+
+    fun getResolveState(moduleInfo: IdeaModuleInfo): FirModuleResolveStateImpl =
+        stateCache.getOrPut(moduleInfo) { createResolveStateFor(moduleInfo) }
 
-internal interface FirIdeResolveStateService {
     companion object {
-        fun getInstance(project: Project): FirIdeResolveStateService =
-            ServiceManager.getService(project, FirIdeResolveStateService::class.java)!!
+        fun getInstance(project: Project): FirIdeResolveStateService = project.service()
     }
-
-    val fallbackModificationTracker: ModificationTracker?
-
-    fun getResolveState(moduleInfo: IdeaModuleInfo): FirModuleResolveStateImpl
-}
-
-private class FirModuleData(val state: FirModuleResolveStateImpl, val modificationTracker: ModificationTracker?) {
-    val modificationCount: Long = modificationTracker?.modificationCount ?: Long.MIN_VALUE
-
-    fun isOutOfDate(): Boolean {
-        val currentModCount = modificationTracker?.modificationCount
-        return currentModCount != null && currentModCount > modificationCount
-    }
-}
-
-internal class FirIdeResolveStateServiceImpl(val project: Project) : FirIdeResolveStateService {
-    private val stateCache = mutableMapOf<IdeaModuleInfo, FirModuleData>()
-
-    private fun createResolveState(): FirModuleResolveStateImpl {
-        val provider = FirProjectSessionProvider(project)
-        return FirModuleResolveStateImpl(provider)
-    }
-
-    private fun createModuleData(): FirModuleData {
-        val state = createResolveState()
-        // We want to invalidate cache on every PSI change for now
-        // This is needed for working with high level API until the proper caching is implemented
-        return FirModuleData(state, fallbackModificationTracker)
-    }
-
-    // TODO: multi thread protection
-    override fun getResolveState(moduleInfo: IdeaModuleInfo): FirModuleResolveStateImpl {
-        var moduleData = stateCache.getOrPut(moduleInfo) {
-            createModuleData()
-        }
-        if (moduleData.isOutOfDate()) {
-            moduleData = createModuleData()
-            stateCache[moduleInfo] = moduleData
-        }
-        return moduleData.state
-    }
-
-    override val fallbackModificationTracker: ModificationTracker? =
-        org.jetbrains.kotlin.analyzer.KotlinModificationTrackerService.getInstance(project).modificationTracker
 }

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

@@ -0,0 +1,30 @@
+/*
+ * 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 com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
+import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
+import org.jetbrains.kotlin.psi.KtElement
+import java.util.concurrent.ConcurrentHashMap
+import javax.annotation.concurrent.ThreadSafe
+
+@ThreadSafe
+internal class FirIdeSessionProvider(project: Project) : FirProjectSessionProvider(project) {
+    override val sessionCache = ConcurrentHashMap<ModuleInfo, FirSession>()
+
+    fun getSession(psi: KtElement): FirSession {
+        val moduleInfo = psi.getModuleInfo()
+        return getSession(moduleInfo)
+    }
+
+    override fun getSession(moduleInfo: ModuleInfo): FirSession {
+        return sessionCache[moduleInfo]
+            ?: error("$moduleInfo")
+    }
+}

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

@@ -5,196 +5,89 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
-import com.google.common.collect.MapMaker
-import com.intellij.openapi.project.Project
-import com.intellij.psi.PsiElement
-import com.intellij.psi.PsiElementVisitor
-import com.intellij.psi.impl.source.tree.LeafPsiElement
-import org.jetbrains.kotlin.cfg.pseudocode.containingDeclarationForPseudocode
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.extensions.BunchOfRegisteredExtensions
-import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.registerExtensions
-import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
-import org.jetbrains.kotlin.fir.render
-import org.jetbrains.kotlin.fir.types.FirTypeRef
-import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
-import org.jetbrains.kotlin.idea.caches.project.LibrarySourceInfo
-import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
-import org.jetbrains.kotlin.idea.util.getElementTextInContext
+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.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.providers.FirIdeProvider
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtTypeReference
 
 abstract class FirModuleResolveState {
-    internal abstract val sessionProvider: FirProjectSessionProvider
+    abstract val moduleInfo: IdeaModuleInfo
+    abstract val firSession: FirSession
 
-    internal fun getSession(psi: KtElement): FirSession {
-        val moduleInfo = psi.getModuleInfo()
-        return getSession(psi.project, moduleInfo)
-    }
+    abstract fun getSessionFor(moduleInfo: IdeaModuleInfo): FirSession
 
-    internal fun getSession(project: Project, moduleInfo: IdeaModuleInfo): FirSession {
-        sessionProvider.getSession(moduleInfo)?.let { return it }
-        val lock = when (moduleInfo) {
-            is ModuleSourceInfo -> moduleInfo.module
-            is LibrarySourceInfo -> moduleInfo.library
-            else -> TODO(moduleInfo.toString())
-        }
-        return synchronized(lock) {
-            val session = sessionProvider.getSession(moduleInfo) ?: FirIdeJavaModuleBasedSession.create(
-                project, moduleInfo, sessionProvider, moduleInfo.contentScope()
-            ).also { moduleBasedSession ->
-                sessionProvider.sessionCache[moduleInfo] = moduleBasedSession
-            }
-            session.also {
-                it.extensionService.registerExtensions(BunchOfRegisteredExtensions.empty())
-            }
-        }
-    }
+    abstract fun getOrBuildFirFor(element: KtElement, toPhase: FirResolvePhase): FirElement
 
-    internal abstract fun getCachedMapping(psi: KtElement): FirElement?
+    abstract fun getDiagnostics(element: KtElement): List<Diagnostic>
 
-    internal abstract fun getDiagnostics(psi: KtElement): List<Diagnostic>
+    // todo temporary, used only in completion
+    abstract fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile)
 
-    internal abstract fun hasDiagnosticsForFile(file: KtFile): Boolean
+    // todo temporary, used only in completion
+    abstract fun getCachedMappingForCompletion(element: KtElement): FirElement?
 
-    internal abstract fun record(psi: KtElement, fir: FirElement)
+    // todo temporary, used only in completion
+    internal abstract fun lazyResolveFunctionForCompletion(
+        firFunction: FirFunction<*>,
+        containerFirFile: FirFile,
+        firIdeProvider: FirIdeProvider,
+        toPhase: FirResolvePhase,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>
+    )
 
-    internal abstract fun record(psi: KtElement, diagnostic: Diagnostic)
-
-    internal abstract fun setDiagnosticsForFile(file: KtFile, fir: FirFile, diagnostics: Iterable<FirDiagnostic<*>> = emptyList())
 }
 
-internal abstract class PsiToFirCache {
-    abstract operator fun get(psi: KtElement): FirElement?
-    abstract operator fun set(psi: KtElement, fir: FirElement)
-    abstract fun remove(psi: PsiElement)
-}
 
-internal class PsiToFirCacheImpl : PsiToFirCache() {
-    private val cache = mutableMapOf<KtElement, FirElement>()
+internal open class FirModuleResolveStateImpl(
+    override val moduleInfo: IdeaModuleInfo,
+    override val firSession: FirSession,
+    private val sessionProvider: FirIdeSessionProvider,
+    val firFileBuilder: FirFileBuilder,
+    val fileCache: ModuleFileCache,
+) : FirModuleResolveState() {
+    val psiToFirCache = PsiToFirCache(fileCache)
+    val elementBuilder = FirElementBuilder(firFileBuilder)
+    private val diagnosticsCollector = DiagnosticsCollector(firFileBuilder, fileCache)
 
-    override fun get(psi: KtElement): FirElement? = cache[psi]
+    override fun getSessionFor(moduleInfo: IdeaModuleInfo): FirSession =
+        sessionProvider.getSession(moduleInfo)
 
-    override fun set(psi: KtElement, fir: FirElement) {
-        cache[psi] = fir
+    override fun getOrBuildFirFor(element: KtElement, toPhase: FirResolvePhase): FirElement =
+        elementBuilder.getOrBuildFirFor(element, fileCache, psiToFirCache, toPhase)
+
+    override fun getDiagnostics(element: KtElement): List<Diagnostic> =
+        diagnosticsCollector.getDiagnosticsFor(element)
+
+    override fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile) {
+        psiToFirCache.recordElementsForCompletionFrom(fir, firFile, ktFile)
     }
 
-    override fun remove(psi: PsiElement) {
-        cache.remove(psi)
-    }
-}
+    override fun getCachedMappingForCompletion(element: KtElement): FirElement? =
+        psiToFirCache.getCachedMapping(element)
 
-open class FirModuleResolveStateImpl(override val sessionProvider: FirProjectSessionProvider) : FirModuleResolveState() {
-    internal open val cache: PsiToFirCache = PsiToFirCacheImpl()
-
-    private val diagnosticCache = mutableMapOf<KtElement, MutableList<Diagnostic>>()
-
-    private val diagnosedFiles = mutableMapOf<KtFile, Long>()
-
-    override fun getCachedMapping(psi: KtElement): FirElement? = cache[psi]
-
-    override fun getDiagnostics(psi: KtElement): List<Diagnostic> {
-        return diagnosticCache[psi] ?: emptyList()
-    }
-
-    override fun hasDiagnosticsForFile(file: KtFile): Boolean {
-        val previousStamp = diagnosedFiles[file] ?: return false
-        if (file.modificationStamp == previousStamp) {
-            return true
-        }
-        diagnosedFiles.remove(file)
-        file.accept(object : PsiElementVisitor() {
-            override fun visitElement(element: PsiElement) {
-                cache.remove(element)
-                diagnosticCache.remove(element)
-                element.acceptChildren(this)
-                super.visitElement(element)
-            }
-        })
-        return false
-    }
-
-    override fun record(psi: KtElement, fir: 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 record(psi: KtElement, diagnostic: Diagnostic) {
-        val list = diagnosticCache.getOrPut(psi) { mutableListOf() }
-        list += diagnostic
-    }
-
-    override fun setDiagnosticsForFile(file: KtFile, fir: FirFile, diagnostics: Iterable<FirDiagnostic<*>>) {
-        for (diagnostic in diagnostics) {
-            require(diagnostic is FirPsiDiagnostic<*>)
-            val psi = diagnostic.element.psi as? KtElement ?: continue
-            record(psi, diagnostic.asPsiBasedDiagnostic())
-        }
-
-        diagnosedFiles[file] = file.modificationStamp
-    }
-}
-
-class DuplicatedFirSourceElementsException(
-    existingFir: FirElement,
-    newFir: FirElement,
-    psi: KtElement
-) : IllegalStateException() {
-    override val message: String? = """|The PSI element should be used only once as a real PSI source of FirElement,
-       |the elements ${if (existingFir.source === newFir.source) "HAVE" else "DON'T HAVE"} the same instances of source elements 
-       |
-       |existing FIR element is $existingFir with text:
-       |${existingFir.render().trim()}
-       |
-       |new FIR element is $newFir with text:
-       | ${newFir.render().trim()}
-       |
-       |PSI element is $psi with text in context:
-       |${psi.getElementTextInContext()}""".trimMargin()
-
-
-    companion object {
-        // The are some cases which are still generates FIR elements with duplicated source elements
-        // Then such case is met, it's better to be fixed
-        // but exception reporting can be easily disabled by setting this to false
-        var IS_ENABLED = true
+    override fun lazyResolveFunctionForCompletion(
+        firFunction: FirFunction<*>,
+        containerFirFile: FirFile,
+        firIdeProvider: FirIdeProvider,
+        toPhase: FirResolvePhase,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>
+    ) {
+        elementBuilder.runLazyResolveForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextForStatement)
     }
 }
 

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

@@ -5,27 +5,63 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
-import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
+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.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.PsiToFirCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
 
-class FirModuleResolveStateForCompletion(
-    mainState: FirModuleResolveStateImpl
-) : FirModuleResolveStateImpl(mainState.sessionProvider) {
-    override val cache: PsiToFirCache = PsiToFirCacheForCompletion(mainState.cache)
-}
+internal class FirModuleResolveStateForCompletion(
+    private val originalState: FirModuleResolveStateImpl
+) : FirModuleResolveState() {
+    override val moduleInfo: IdeaModuleInfo get() = originalState.moduleInfo
+    override val firSession: FirSession get() = originalState.firSession
 
-private class PsiToFirCacheForCompletion(private val delegate: PsiToFirCache) : PsiToFirCache() {
-    private val cache = mutableMapOf<KtElement, FirElement>()
+    private val psiToFirCache = PsiToFirCache(originalState.fileCache)
 
-    override fun get(psi: KtElement): FirElement? =
-        cache[psi] ?: delegate[psi]
+    override fun getSessionFor(moduleInfo: IdeaModuleInfo): FirSession =
+        originalState.getSessionFor(moduleInfo)
 
-    override fun set(psi: KtElement, fir: FirElement) {
-        cache[psi] = fir
+    override fun getOrBuildFirFor(element: KtElement, toPhase: FirResolvePhase): FirElement {
+        getCachedMappingForCompletion(element)?.let { return it }
+        return originalState.elementBuilder.getOrBuildFirFor(
+            element,
+            originalState.fileCache,
+            psiToFirCache,
+            toPhase
+        )
     }
 
-    override fun remove(psi: PsiElement) {
-        cache.remove(psi)
+    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
+    }
+
+    override fun lazyResolveFunctionForCompletion(
+        firFunction: FirFunction<*>,
+        containerFirFile: FirFile,
+        firIdeProvider: FirIdeProvider,
+        toPhase: FirResolvePhase,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>
+    ) {
+        originalState.lazyResolveFunctionForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextForStatement)
+    }
+
+    override fun getDiagnostics(element: KtElement): List<Diagnostic> {
+        error("Diagnostics should not be retrieved in completion")
     }
 }

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

@@ -1,326 +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
-
-import com.intellij.openapi.progress.ProgressIndicatorProvider
-import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.fir.references.*
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
-import org.jetbrains.kotlin.fir.resolve.ResolutionMode
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
-import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredCallableSymbols
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
-import org.jetbrains.kotlin.fir.resolve.transformers.runResolve
-import org.jetbrains.kotlin.fir.scopes.impl.FirPackageMemberScope
-import org.jetbrains.kotlin.fir.symbols.CallableId
-import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
-import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
-import org.jetbrains.kotlin.fir.types.FirUserTypeRef
-import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
-import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
-import org.jetbrains.kotlin.fir.visitors.compose
-import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
-
-private val FirResolvePhase.stubMode: Boolean
-    get() = this <= FirResolvePhase.DECLARATIONS
-
-private fun KtClassOrObject.relativeFqName(): FqName {
-    val className = this.nameAsSafeName
-    val parentFqName = this.containingClassOrObject?.relativeFqName()
-    return parentFqName?.child(className) ?: FqName.topLevel(className)
-}
-
-private fun FirFile.findCallableMember(
-    provider: FirProvider, callableMember: KtCallableDeclaration,
-    packageFqName: FqName, klassFqName: FqName?, declName: Name
-): FirCallableDeclaration<*> {
-    if (klassFqName != null) {
-        return provider.getClassDeclaredCallableSymbols(ClassId(packageFqName, klassFqName, false), declName)
-            .find { symbol: FirCallableSymbol<*> ->
-                symbol.fir.psi == callableMember
-            }?.fir
-            ?: error("Cannot find FIR callable declaration ${CallableId(packageFqName, klassFqName, declName)}")
-    }
-    // NB: not sure it's correct to use member scope provider from here (because of possible changes)
-    val memberScope = FirPackageMemberScope(this.packageFqName, session)
-    var result: FirCallableDeclaration<*>? = null
-    val processor = { symbol: FirCallableSymbol<*> ->
-        val fir = symbol.fir
-        if (result == null && fir.psi == callableMember) {
-            result = fir
-        }
-    }
-    if (callableMember is KtNamedFunction || callableMember is KtConstructor<*>) {
-        memberScope.processFunctionsByName(declName, processor)
-    } else {
-        memberScope.processPropertiesByName(declName, processor)
-    }
-
-    return result
-        ?: error("Cannot find FIR callable declaration ${CallableId(packageFqName, klassFqName, declName)}")
-}
-
-internal fun KtCallableDeclaration.getOrBuildFir(
-    state: FirModuleResolveState,
-    phase: FirResolvePhase = FirResolvePhase.DECLARATIONS
-): FirCallableDeclaration<*> {
-    val session = state.getSession(this)
-
-    val file = this.containingKtFile
-    val packageFqName = file.packageFqName
-    val klassFqName = this.containingClassOrObject?.relativeFqName()
-    val declName = this.nameAsSafeName
-
-    val firProvider = session.firIdeProvider
-    val firFile = firProvider.getOrBuildFile(file)
-    val firMemberSymbol = firFile.findCallableMember(firProvider, this, packageFqName, klassFqName, declName).symbol
-    val firMemberDeclaration = firMemberSymbol.fir
-    if (firMemberDeclaration.resolvePhase >= phase) {
-        return firMemberDeclaration
-    }
-    synchronized(firFile) {
-        firMemberDeclaration.runResolve(firFile, firProvider, phase, state)
-    }
-    return firMemberDeclaration
-}
-
-internal fun KtClassOrObject.getOrBuildFir(
-    state: FirModuleResolveState,
-    phase: FirResolvePhase = FirResolvePhase.DECLARATIONS
-): FirMemberDeclaration {
-    val session = state.getSession(this)
-
-    val file = this.containingKtFile
-    val packageFqName = file.packageFqName
-    val klassFqName = this.relativeFqName()
-
-    val firProvider = session.firIdeProvider
-    val firFile = firProvider.getOrBuildFile(file)
-
-    val firClassOrEnumEntry = if (this is KtEnumEntry) {
-        val firEnumClass = firProvider.getFirClassifierByFqName(ClassId(packageFqName, klassFqName.parent(), false)) as FirRegularClass
-        firEnumClass.declarations.first { it is FirEnumEntry && it.name == this.nameAsSafeName } as FirMemberDeclaration
-    } else {
-        firProvider.getFirClassifierByFqName(ClassId(packageFqName, klassFqName, false)) as FirRegularClass
-    }
-    if (firClassOrEnumEntry.resolvePhase >= phase) {
-        return firClassOrEnumEntry
-    }
-    synchronized(firFile) {
-        firClassOrEnumEntry.runResolve(firFile, firProvider, phase, state)
-    }
-    return firClassOrEnumEntry
-}
-
-private fun KtFile.getOrBuildRawFirFile(state: FirModuleResolveState): Pair<FirIdeProvider, FirFile> {
-    val session = state.getSession(this)
-    val firProvider = session.firIdeProvider
-    return firProvider to firProvider.getOrBuildFile(this)
-}
-
-internal fun KtFile.getOrBuildFir(
-    state: FirModuleResolveState,
-    phase: FirResolvePhase = FirResolvePhase.DECLARATIONS
-): FirFile {
-    val (firProvider, firFile) = getOrBuildRawFirFile(state)
-    if (phase <= FirResolvePhase.DECLARATIONS && firFile.resolvePhase >= phase) {
-        return firFile
-    }
-    synchronized(firFile) {
-        firFile.runResolve(firFile, firProvider, phase, state)
-    }
-    return firFile
-}
-
-internal fun KtFile.getOrBuildFirWithDiagnostics(state: FirModuleResolveState): FirFile {
-    val (_, firFile) = getOrBuildRawFirFile(state)
-    val currentResolvePhase = firFile.resolvePhase
-    if (currentResolvePhase < FirResolvePhase.BODY_RESOLVE) {
-        synchronized(firFile) {
-            firFile.runResolve(toPhase = FirResolvePhase.BODY_RESOLVE, fromPhase = currentResolvePhase)
-        }
-    }
-
-    ProgressIndicatorProvider.checkCanceled() // ???
-    if (state.hasDiagnosticsForFile(this)) return firFile
-
-    FirIdeDiagnosticsCollector(firFile.session, state).collectDiagnostics(firFile)
-    state.setDiagnosticsForFile(this, firFile)
-    return firFile
-}
-
-internal fun FirDeclaration.runResolve(
-    file: FirFile,
-    firProvider: FirIdeProvider,
-    toPhase: FirResolvePhase,
-    state: FirModuleResolveState,
-    towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
-) {
-    val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
-    file.runResolve(toPhase = nonLazyPhase, fromPhase = this.resolvePhase)
-    if (toPhase <= nonLazyPhase) return
-    val designation = mutableListOf<FirDeclaration>(file)
-    if (this !is FirFile) {
-        val id = when (this) {
-            is FirCallableDeclaration<*> -> {
-                this.symbol.callableId.classId
-            }
-            is FirRegularClass -> {
-                this.symbol.classId
-            }
-            else -> error("Unsupported: ${render()}")
-        }
-        val outerClasses = generateSequence(id) { classId ->
-            classId.outerClassId
-        }.mapTo(mutableListOf()) { firProvider.getFirClassifierByFqName(it)!! }
-        designation += outerClasses.asReversed()
-        if (this is FirCallableDeclaration<*>) {
-            designation += this
-        }
-    }
-    if (designation.all { it.resolvePhase >= toPhase }) {
-        return
-    }
-    val scopeSession = ScopeSession()
-    val transformer = FirDesignatedBodyResolveTransformerForIDE(
-        designation.iterator(), state.getSession(psi as KtElement),
-        scopeSession,
-        implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
-        towerDataContextForStatement
-    )
-    file.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
-}
-
-private class FirDesignatedBodyResolveTransformerForIDE(
-    private val designation: Iterator<FirElement>,
-    session: FirSession,
-    scopeSession: ScopeSession,
-    implicitTypeOnly: Boolean,
-    private val towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null
-) : FirBodyResolveTransformer(
-    session,
-    phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
-    implicitTypeOnly = implicitTypeOnly,
-    scopeSession = scopeSession,
-    returnTypeCalculator = createReturnTypeCalculatorForIDE(session, scopeSession)
-) {
-
-    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
-        if (designation.hasNext()) {
-            designation.next().visitNoTransform(this, data)
-            return declaration.compose()
-        }
-
-        return super.transformDeclarationContent(declaration, data)
-    }
-
-    override fun onBeforeStatementResolution(statement: FirStatement) {
-        if (towerDataContextForStatement == null) return
-        towerDataContextForStatement[statement] = context.towerDataContext
-    }
-}
-
-
-private fun KtElement.getNonLocalContainingDeclarationWithFqName(): KtDeclaration? {
-    var container = parent
-    while (container != null && container !is KtFile) {
-        if (container is KtDeclaration
-            && (container is KtClassOrObject || container is KtDeclarationWithBody)
-            && !KtPsiUtil.isLocal(container)
-            && container.name != null
-            && container !is KtEnumEntry
-            && container.containingClassOrObject !is KtEnumEntry
-        ) {
-            return container
-        }
-        container = container.parent
-    }
-    return null
-}
-
-internal fun KtElement.getOrBuildFir(
-    state: FirModuleResolveState,
-    phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
-): FirElement {
-    val containerFir: FirDeclaration =
-        when (val container = getNonLocalContainingDeclarationWithFqName()) {
-            is KtCallableDeclaration -> container.getOrBuildFir(state, phase)
-            is KtClassOrObject -> container.getOrBuildFir(state, phase)
-            null -> containingKtFile.getOrBuildFir(state, phase)
-            else -> error("Unsupported: ${container.text}")
-        }
-
-    val psi = when (this) {
-        is KtPropertyDelegate -> this.expression ?: this
-        else -> this
-    }
-    return state.getCachedMapping(this) ?: run {
-        state.recordElementsFrom(containerFir)
-
-        val (current, mappedFir) = psi.getFirOfClosestParent(state) ?: error("FirElement is not found for: $text")
-        if (current != this) {
-            state.record(current, mappedFir)
-        }
-
-        mappedFir
-    }
-}
-
-internal fun KtElement.getFirOfClosestParent(state: FirModuleResolveState): Pair<KtElement, FirElement>? {
-    var current: PsiElement? = this
-    while (current is KtElement) {
-        val mappedFir = state.getCachedMapping(current)
-        if (mappedFir != null) {
-            return current to mappedFir
-        }
-        current = current.parent
-    }
-
-    return null
-}
-
-internal fun FirModuleResolveState.recordElementsFrom(containerFir: FirDeclaration) {
-    containerFir.accept(object : FirVisitorVoid() {
-        override fun visitElement(element: FirElement) {
-            (element.realPsi as? KtElement)?.let {
-                record(it, 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)
-        }
-    })
-}

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

@@ -0,0 +1,21 @@
+/*
+ * 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.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.transformers.PhasedFirFileResolver
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+
+internal class IdePhasedFirFileResolver(
+    private val firFileBuilder: FirFileBuilder,
+    private val cache: ModuleFileCache
+) : PhasedFirFileResolver() {
+    override fun resolveFile(firFile: FirFile, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
+        firFileBuilder.runResolve(firFile, cache, fromPhase, toPhase)
+    }
+}

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

@@ -0,0 +1,77 @@
+/*
+ * 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 com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.stubs.StubIndex
+import com.intellij.psi.stubs.StubIndexKey
+import org.jetbrains.kotlin.fir.symbols.CallableId
+import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper.Companion.asStringForIndexes
+import org.jetbrains.kotlin.idea.stubindex.*
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.KtProperty
+
+internal class IndexHelper(val project: Project, private val scope: GlobalSearchScope) {
+    private val stubIndex: StubIndex = StubIndex.getInstance()
+
+    private inline fun <INDEX_KEY : Any, reified PSI : PsiElement> firstMatchingOrNull(
+        stubKey: StubIndexKey<INDEX_KEY, PSI>,
+        key: INDEX_KEY,
+        crossinline filter: (PSI) -> Boolean
+    ): PSI? {
+        var result: PSI? = null
+        stubIndex.processElements(
+            stubKey, key, project, scope, PSI::class.java,
+        ) { candidate ->
+            if (filter(candidate)) {
+                result = candidate
+                return@processElements false // do not continue searching over PSI
+            }
+            true
+        }
+        return result
+    }
+
+    fun classFromIndexByClassId(classId: ClassId) = firstMatchingOrNull(
+        KotlinFullClassNameIndex.KEY,
+        classId.asStringForIndexes(),
+    ) { candidate -> candidate.containingKtFile.packageFqName == classId.packageFqName }
+
+    fun typeAliasFromIndexByClassId(classId: ClassId) = firstMatchingOrNull(
+        KotlinTopLevelTypeAliasFqNameIndex.KEY,
+        classId.asStringForIndexes(),
+    ) { candidate -> candidate.containingKtFile.packageFqName == classId.packageFqName }
+
+
+    fun getTopLevelProperties(callableId: CallableId): Collection<KtProperty> =
+        KotlinTopLevelPropertyFqnNameIndex.getInstance()[callableId.asStringForIndexes(), project, scope]
+
+    fun getTopLevelFunctions(callableId: CallableId): Collection<KtNamedFunction> =
+        KotlinTopLevelFunctionFqnNameIndex.getInstance()[callableId.asStringForIndexes(), project, scope]
+
+    fun getTopLevelPropertiesInPackage(packageFqName: FqName): Collection<KtProperty> =
+        KotlinTopLevelPropertyByPackageIndex.getInstance().get(packageFqName.asStringForIndexes(), project, scope)
+
+    fun getTopLevelFunctionsInPackage(packageFqName: FqName): Collection<KtFunction> =
+        KotlinTopLevelFunctionByPackageIndex.getInstance().get(packageFqName.asStringForIndexes(), project, scope)
+
+
+    companion object {
+        private fun CallableId.asStringForIndexes(): String =
+            (if (packageName.isRoot) callableName.asString() else toString()).replace('/', '.')
+
+        private fun FqName.asStringForIndexes(): String =
+            asString().replace('/', '.')
+
+        private fun ClassId.asStringForIndexes(): String =
+            asSingleFqName().asStringForIndexes()
+    }
+}

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

@@ -12,27 +12,28 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.psi
 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
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
 
 object LowLevelFirApiFacade {
     fun getResolveStateFor(element: KtElement): FirModuleResolveState =
         element.firResolveState()
 
-    fun getResolveStateForCompletion(element: KtElement, mainState: FirModuleResolveStateImpl): FirModuleResolveStateForCompletion {
-        return FirModuleResolveStateForCompletion(mainState)
+    fun getResolveStateForCompletion(element: KtElement, originalState: FirModuleResolveState): FirModuleResolveState {
+        check(originalState is FirModuleResolveStateImpl)
+        return FirModuleResolveStateForCompletion(originalState)
     }
 
-    fun getSessionFor(element: KtElement, resolveState: FirModuleResolveState): FirSession =
-        resolveState.getSession(element)
+    fun getSessionFor(element: KtElement): FirSession =
+        getResolveStateFor(element).getSessionFor(element.getModuleInfo())
 
     fun getOrBuildFirFor(element: KtElement, resolveState: FirModuleResolveState, phase: FirResolvePhase): FirElement =
-        element.getOrBuildFir(resolveState, phase)
-
-
-    fun getFirOfClosestParent(element: KtElement): FirElement? = element.getFirOfClosestParent(element.firResolveState())?.second
+        resolveState.getOrBuildFirFor(element, phase)
 
     class FirCompletionContext internal constructor(
         val session: FirSession,
@@ -42,7 +43,7 @@ object LowLevelFirApiFacade {
         fun getTowerDataContext(element: KtElement): FirTowerDataContext {
             var current: PsiElement? = element
             while (current is KtElement) {
-                val mappedFir = state.getCachedMapping(current)
+                val mappedFir = state.getCachedMappingForCompletion(current)
 
                 if (mappedFir is FirStatement) {
                     towerDataContextForStatement[mappedFir]?.let { return it }
@@ -56,7 +57,7 @@ object LowLevelFirApiFacade {
                 current = current.parent
             }
 
-            error("No context for $element")
+            error("No context for ${element.getElementTextInContext()}")
         }
     }
 
@@ -71,10 +72,8 @@ object LowLevelFirApiFacade {
         val towerDataContextForStatement = mutableMapOf<FirStatement, FirTowerDataContext>()
 
         val function = builtFunction.apply {
-            runResolve(firFile, firIdeProvider, phase, state, towerDataContextForStatement)
-
-            // TODO this PSI caching should be somewhere else
-            state.recordElementsFrom(this)
+            state.lazyResolveFunctionForCompletion(this, firFile, firIdeProvider, phase, towerDataContextForStatement)
+            state.recordPsiToFirMappingsForCompletionFrom(this, firFile, element.containingKtFile)
         }
 
         return FirCompletionContext(
@@ -85,8 +84,6 @@ object LowLevelFirApiFacade {
     }
 
     fun getDiagnosticsFor(element: KtElement, resolveState: FirModuleResolveState): Collection<Diagnostic> {
-        val file = element.containingKtFile
-        file.getOrBuildFirWithDiagnostics(resolveState)
         return resolveState.getDiagnostics(element)
     }
 }

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

@@ -0,0 +1,39 @@
+/*
+ * 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 com.intellij.openapi.components.service
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.analyzer.PackageOracleFactory
+import org.jetbrains.kotlin.idea.caches.resolve.IdePackageOracleFactory
+import org.jetbrains.kotlin.name.FqName
+
+internal sealed class PackageExistenceChecker {
+    abstract fun isPackageExists(packageFqName: FqName): Boolean
+}
+
+internal class PackageExistenceCheckerForSingleModule(
+    project: Project,
+    module: ModuleInfo
+) : PackageExistenceChecker() {
+    private val oracle =
+        project.service<IdePackageOracleFactory>().createOracle(module)
+    override fun isPackageExists(packageFqName: FqName): Boolean = oracle.packageExists(packageFqName)
+}
+
+internal class PackageExistenceCheckerForMultipleModules(
+    project: Project,
+    modules: List<ModuleInfo>
+) : PackageExistenceChecker() {
+    private val oracles = run {
+        val factory = project.service<IdePackageOracleFactory>()
+        modules.map { factory.createOracle(it) }
+    }
+
+    override fun isPackageExists(packageFqName: FqName): Boolean =
+        oracles.any { oracle -> oracle.packageExists(packageFqName) }
+}

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

@@ -0,0 +1,48 @@
+/*
+ * 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.diagnostics
+
+import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+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.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import java.util.concurrent.ConcurrentHashMap
+
+internal class DiagnosticsCollector(
+    private val firFileBuilder: FirFileBuilder,
+    private val cache: ModuleFileCache,
+) {
+    private val diagnosticsForFile = ConcurrentHashMap<KtFile, DiagnosticsForFile>()
+
+    fun getDiagnosticsFor(element: KtElement): List<Diagnostic> {
+        val ktFile = element.containingKtFile
+        val diagnostics = diagnosticsForFile.computeIfAbsent(ktFile) {
+            val firFile = firFileBuilder.getFirFileResolvedToPhaseWithCaching(ktFile, cache, toPhase = FirResolvePhase.BODY_RESOLVE)
+            DiagnosticsForFile.collectDiagnosticsForFile(firFile)
+        }
+        return diagnostics.getDiagnosticsFor(element)
+    }
+}
+
+private class DiagnosticsForFile private constructor(private val diagnostics: Map<KtElement, List<Diagnostic>>) {
+    fun getDiagnosticsFor(element: KtElement): List<Diagnostic> = diagnostics[element].orEmpty()
+
+    companion object {
+        /**
+         * Collects diagnostics for given [firFile]
+         * Should be called under [firFile]-based lock
+         */
+        fun collectDiagnosticsForFile(firFile: FirFile): DiagnosticsForFile {
+            require(firFile.resolvePhase >= FirResolvePhase.BODY_RESOLVE) {
+                "To collect diagnostics at least FirResolvePhase.BODY_RESOLVE is needed, but file ${firFile.name} was resolved to ${firFile.resolvePhase}"
+            }
+            return DiagnosticsForFile(FirIdeDiagnosticsCollector.collect(firFile))
+        }
+    }
+}

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

@@ -3,17 +3,23 @@
  * 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
+package org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics
 
+import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
 import org.jetbrains.kotlin.fir.analysis.collectors.registerAllComponents
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.addValueFor
 import org.jetbrains.kotlin.psi.KtElement
 
-internal class FirIdeDiagnosticsCollector(session: FirSession, private val resolveState: FirModuleResolveState) : AbstractDiagnosticCollector(session) {
+internal class FirIdeDiagnosticsCollector private constructor(
+    session: FirSession,
+) : AbstractDiagnosticCollector(session) {
+    private val result = mutableMapOf<KtElement, MutableList<Diagnostic>>()
 
     init {
         registerAllComponents()
@@ -23,7 +29,7 @@ internal class FirIdeDiagnosticsCollector(session: FirSession, private val resol
         override fun report(diagnostic: FirDiagnostic<*>?) {
             if (diagnostic !is FirPsiDiagnostic<*>) return
             val psi = diagnostic.element.psi as? KtElement ?: return
-            resolveState.record(psi, diagnostic.asPsiBasedDiagnostic())
+            result.addValueFor(psi, diagnostic.asPsiBasedDiagnostic())
         }
     }
 
@@ -41,4 +47,16 @@ internal class FirIdeDiagnosticsCollector(session: FirSession, private val resol
     override fun runCheck(block: (DiagnosticReporter) -> Unit) {
         block(reporter)
     }
+
+    companion object {
+        /**
+         * Collects diagnostics for given [firFile]
+         * Should be called under [firFile]-based lock
+         */
+        fun collect(firFile: FirFile): Map<KtElement, List<Diagnostic>> =
+            FirIdeDiagnosticsCollector(firFile.session).let { collector ->
+                collector.collectDiagnostics(firFile)
+                collector.result
+            }
+    }
 }

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

@@ -0,0 +1,37 @@
+/*
+ * 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 org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
+import org.jetbrains.kotlin.psi.KtElement
+
+class DuplicatedFirSourceElementsException(
+    existingFir: FirElement,
+    newFir: FirElement,
+    psi: KtElement
+) : IllegalStateException() {
+    override val message: String? = """|The PSI element should be used only once as a real PSI source of FirElement,
+       |the elements ${if (existingFir.source === newFir.source) "HAVE" else "DON'T HAVE"} the same instances of source elements 
+       |
+       |existing FIR element is $existingFir with text:
+       |${existingFir.render().trim()}
+       |
+       |new FIR element is $newFir with text:
+       | ${newFir.render().trim()}
+       |
+       |PSI element is $psi with text in context:
+       |${psi.getElementTextInContext()}""".trimMargin()
+
+
+    companion object {
+        // The are some cases which are still generates FIR elements with duplicated source elements
+        // Then such case is met, it's better to be fixed
+        // but exception reporting can be easily disabled by setting this to false
+        var IS_ENABLED = false
+    }
+}

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

@@ -0,0 +1,48 @@
+/*
+ * 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 org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.compose
+
+internal class FirDesignatedBodyResolveTransformerForIDE(
+    private val designation: Iterator<FirElement>,
+    session: FirSession,
+    scopeSession: ScopeSession,
+    implicitTypeOnly: Boolean,
+    private val towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null
+) : FirBodyResolveTransformer(
+    session,
+    phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+    implicitTypeOnly = implicitTypeOnly,
+    scopeSession = scopeSession,
+    returnTypeCalculator = createReturnTypeCalculatorForIDE(session, scopeSession)
+) {
+
+    override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
+        if (designation.hasNext()) {
+            designation.next().visitNoTransform(this, data)
+            return declaration.compose()
+        }
+
+        return super.transformDeclarationContent(declaration, data)
+    }
+
+    override fun onBeforeStatementResolution(statement: FirStatement) {
+        if (towerDataContextForStatement == null) return
+        towerDataContextForStatement[statement] = context.towerDataContext
+    }
+}

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

@@ -0,0 +1,134 @@
+/*
+ * 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 org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+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
+import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.FirElementFinder
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+import javax.annotation.concurrent.ThreadSafe
+
+/**
+ * Maps [KtElement] to [FirElement]
+ * Stateless, caches everything into [ModuleFileCache] & [PsiToFirCache] passed into the function
+ */
+@ThreadSafe
+internal class FirElementBuilder(
+    private val firFileBuilder: FirFileBuilder,
+) {
+    fun getOrBuildFirFor(
+        element: KtElement,
+        moduleFileCache: ModuleFileCache,
+        psiToFirCache: PsiToFirCache,
+        toPhase: FirResolvePhase,
+    ): FirElement {
+        val ktFile = element.containingKtFile
+        val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
+
+        val containerFir = when (val container = element.getNonLocalContainingDeclarationWithFqName()) {
+            is KtDeclaration -> {
+                FirElementFinder.findElementByPsiIn<FirDeclaration>(firFile, container)
+                    ?: error("Declaration was not found in FIR file which was build by declaration KtFile, declaration is\n${container.getElementTextInContext()}")
+            }
+            null -> firFile
+            else -> error("Unsupported: ${container.text}")
+        }
+
+        firFileBuilder.runCustomResolve(firFile, moduleFileCache) {
+            runLazyResolveWithoutLock(containerFir, firFile, firFile.session.firIdeProvider, toPhase)
+        }
+
+        return psiToFirCache.getFir(element, containerFir, firFile)
+    }
+
+    /**
+     * Should be invoked under lock
+     */
+    private fun runLazyResolveWithoutLock(
+        firDeclarationToResolve: FirDeclaration,
+        containerFirFile: FirFile,
+        firIdeProvider: FirIdeProvider,
+        toPhase: FirResolvePhase,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
+    ) {
+        val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
+        if (firDeclarationToResolve.resolvePhase < nonLazyPhase) {
+            firFileBuilder.runResolveWithoutLock(containerFirFile, fromPhase = firDeclarationToResolve.resolvePhase, toPhase = nonLazyPhase)
+        }
+        if (toPhase <= nonLazyPhase) return
+        val designation = mutableListOf<FirDeclaration>(containerFirFile)
+        if (firDeclarationToResolve !is FirFile) {
+            val id = when (firDeclarationToResolve) {
+                is FirCallableDeclaration<*> -> {
+                    firDeclarationToResolve.symbol.callableId.classId
+                }
+                is FirRegularClass -> {
+                    firDeclarationToResolve.symbol.classId
+                }
+                else -> error("Unsupported: ${firDeclarationToResolve.render()}")
+            }
+            val outerClasses = generateSequence(id) { classId ->
+                classId.outerClassId
+            }.mapTo(mutableListOf()) { firIdeProvider.getFirClassifierByFqName(it)!! }
+            designation += outerClasses.asReversed()
+            if (firDeclarationToResolve is FirCallableDeclaration<*>) {
+                designation += firDeclarationToResolve
+            }
+        }
+        if (designation.all { it.resolvePhase >= toPhase }) {
+            return
+        }
+        val scopeSession = ScopeSession()
+        val transformer = FirDesignatedBodyResolveTransformerForIDE(
+            designation.iterator(), containerFirFile.session,
+            scopeSession,
+            implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+            towerDataContextForStatement
+        )
+        containerFirFile.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
+    }
+
+    //TODO for completion only
+    fun runLazyResolveForCompletion(
+        firFunction: FirFunction<*>,
+        containerFirFile: FirFile,
+        firIdeProvider: FirIdeProvider,
+        toPhase: FirResolvePhase,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>
+    ) {
+        runLazyResolveWithoutLock(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextForStatement)
+    }
+}
+
+
+
+private fun KtElement.getNonLocalContainingDeclarationWithFqName(): KtDeclaration? {
+    var container = parent
+    while (container != null && container !is KtFile) {
+        if (container is KtDeclaration
+            && (container is KtClassOrObject || container is KtDeclarationWithBody)
+            && !KtPsiUtil.isLocal(container)
+            && container.name != null
+            && container !is KtEnumEntry
+            && container.containingClassOrObject !is KtEnumEntry
+        ) {
+            return container
+        }
+        container = container.parent
+    }
+    return null
+}

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

@@ -0,0 +1,156 @@
+/*
+ * 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 org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtTypeReference
+import java.util.concurrent.ConcurrentHashMap
+import org.jetbrains.kotlin.psi.KtPropertyDelegate
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
+
+/**
+ * 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
+            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/builder/FirFileBuilder.kt

@@ -0,0 +1,79 @@
+/*
+ * 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.builder
+
+import org.jetbrains.kotlin.fir.builder.RawFirBuilder
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.createTransformerBasedProcessorByPhase
+import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeSessionProvider
+import org.jetbrains.kotlin.psi.KtFile
+import javax.annotation.concurrent.ThreadSafe
+
+/**
+ * Responsible for building [FirFile] by [KtFile]
+ * Stateless, all caches are stored in [ModuleFileCache] passed into corresponding functions
+ */
+@ThreadSafe
+internal class FirFileBuilder(
+    private val sessionProvider: FirIdeSessionProvider,
+    private val scopeProvider: FirScopeProvider
+) {
+    /**
+     * Builds a [FirFile] by given [ktFile] and records it's parenting info if it not present in [cache]
+     * [FirFile] building a happens at most once per each [KtFile]
+     */
+    fun buildRawFirFileWithCaching(
+        ktFile: KtFile,
+        cache: ModuleFileCache
+    ): FirFile = cache.fileCached(ktFile) {
+        RawFirBuilder(cache.session, scopeProvider, stubMode = false).buildFirFile(ktFile)
+    }
+
+    fun getFirFileResolvedToPhaseWithCaching(
+        ktFile: KtFile,
+        cache: ModuleFileCache,
+        @Suppress("SameParameterValue") toPhase: FirResolvePhase
+    ): FirFile {
+        val firFile = buildRawFirFileWithCaching(ktFile, cache)
+        if (toPhase > FirResolvePhase.RAW_FIR) {
+            cache.firFileLockProvider.withLock(firFile) {
+                //add lock for implit type resolve phase & super type
+                runResolveWithoutLock(firFile, fromPhase = firFile.resolvePhase, toPhase = toPhase)
+            }
+        }
+        return firFile
+    }
+
+    /**
+     * Runs [resolve] function (which is considered to do some resolve on [firFile]) under a lock for [firFile]
+     */
+    inline fun <R> runCustomResolve(firFile: FirFile, cache: ModuleFileCache, resolve: () -> R): R =
+        cache.firFileLockProvider.withLock(firFile) { resolve() }
+
+    fun runResolve(firFile: FirFile, cache: ModuleFileCache, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
+        cache.firFileLockProvider.withLock(firFile) {
+            runResolveWithoutLock(firFile, fromPhase, toPhase)
+        }
+    }
+
+    fun runResolveWithoutLock(firFile: FirFile, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
+        assert(fromPhase <= toPhase) {
+            "Trying to resolve file ${firFile.name} from $fromPhase to $toPhase"
+        }
+        val scopeSession = ScopeSession()
+        var currentPhase = fromPhase
+        while (currentPhase < toPhase) {
+            currentPhase = currentPhase.next
+            val phaseProcessor = currentPhase.createTransformerBasedProcessorByPhase(firFile.session, scopeSession)
+            phaseProcessor.processFile(firFile)
+        }
+    }
+}
+
+

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/LockProvider.kt

@@ -0,0 +1,33 @@
+/*
+ * 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.builder
+
+import com.google.common.collect.MapMaker
+import java.util.concurrent.ConcurrentMap
+import java.util.concurrent.locks.ReadWriteLock
+import java.util.concurrent.locks.ReentrantLock
+import kotlin.concurrent.withLock
+
+internal class LockProvider<KEY, out LOCK>(private val createLock: () -> LOCK) {
+    private val locks: ConcurrentMap<KEY, LOCK> = MapMaker().weakKeys().makeMap()
+    fun getLockFor(key: KEY) = locks.getOrPut(key) { createLock() }
+}
+
+internal inline fun <KEY, R> LockProvider<KEY, ReadWriteLock>.withReadLock(key: KEY, action: () -> R): R {
+    val readLock = getLockFor(key).readLock()
+    return readLock.withLock { action() }
+}
+
+internal inline fun <KEY, R> LockProvider<KEY, ReadWriteLock>.withWriteLock(key: KEY, action: () -> R): R {
+    val writeLock = getLockFor(key).writeLock()
+    return writeLock.withLock { action() }
+}
+
+
+internal inline fun <KEY, R> LockProvider<KEY, ReentrantLock>.withLock(key: KEY, action: () -> R): R {
+    val lock = getLockFor(key)
+    return lock.withLock { action() }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/ModuleFileCache.kt

@@ -0,0 +1,70 @@
+/*
+ * 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.builder
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.symbols.CallableId
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.psi.KtFile
+import java.util.*
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.locks.ReentrantLock
+import java.util.concurrent.locks.ReentrantReadWriteLock
+import javax.annotation.concurrent.ThreadSafe
+import kotlin.concurrent.withLock
+
+/**
+ * Caches mapping [KtFile] -> [FirFile] of module [moduleInfo]
+ */
+@ThreadSafe
+internal abstract class ModuleFileCache {
+    abstract val session: FirSession
+
+    /**
+     * Maps [ClassId] to corresponding classifiers
+     * If classifier with required [ClassId] is not found in given module then map contains [Optional.EMPTY]
+     */
+    abstract val classifierByClassId: ConcurrentHashMap<ClassId, Optional<FirClassLikeDeclaration<*>>>
+
+    /**
+     * Maps [CallableId] to corresponding callable
+     * If callable with required [CallableId]] is not found in given module then map contains emptyList
+     */
+    abstract val callableByCallableId: ConcurrentHashMap<CallableId, List<FirCallableSymbol<*>>>
+
+    /**
+     * @return [FirFile] by [file] if it was previously built or runs [createValue] otherwise
+     * The [createValue] is run under the lock so [createValue] is executed at most once for each [KtFile]
+     */
+    abstract fun fileCached(file: KtFile, createValue: () -> FirFile): FirFile
+
+    abstract fun getCachedFirFile(ktFile: KtFile): FirFile?
+
+    // todo make it ReadWriteLock and allow access fir elements only under read lock
+    // for now locks only held for resolve
+    // but there can be a situation when we are accessing some fir element in one thread without lock
+    // in the same time other thread performs resolve of it
+    // which can cause weird errors on user side
+    abstract val firFileLockProvider: LockProvider<FirFile, ReentrantLock>
+}
+
+internal class ModuleFileCacheImpl(override val session: FirSession) : ModuleFileCache() {
+    private val ktFileToFirFile = ConcurrentHashMap<KtFile, FirFile>()
+
+    override val classifierByClassId: ConcurrentHashMap<ClassId, Optional<FirClassLikeDeclaration<*>>> = ConcurrentHashMap()
+    override val callableByCallableId: ConcurrentHashMap<CallableId, List<FirCallableSymbol<*>>> = ConcurrentHashMap()
+
+    override fun fileCached(file: KtFile, createValue: () -> FirFile): FirFile =
+        ktFileToFirFile.computeIfAbsent(file) { createValue() }
+
+    override fun getCachedFirFile(ktFile: KtFile): FirFile? = ktFileToFirFile[ktFile]
+
+
+    override val firFileLockProvider: LockProvider<FirFile, ReentrantLock> = LockProvider { ReentrantLock() }
+}

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

@@ -0,0 +1,134 @@
+/*
+ * 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.providers
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.builder.RawFirBuilder
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
+import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirProviderInternals
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirAccessorSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper
+import org.jetbrains.kotlin.idea.fir.low.level.api.PackageExistenceCheckerForMultipleModules
+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.collectTransitiveDependenciesWithSelf
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtNamedFunction
+
+internal class FirIdeProvider(
+    project: Project,
+    val session: FirSession,
+    moduleInfo: ModuleSourceInfo,
+    private val kotlinScopeProvider: KotlinScopeProvider,
+    firFileBuilder: FirFileBuilder,
+    val cache: ModuleFileCache,
+    searchScope: GlobalSearchScope,
+) : FirProvider() {
+    private val indexHelper = IndexHelper(project, searchScope)
+    private val packageExistenceChecker = PackageExistenceCheckerForMultipleModules(
+        project,
+        moduleInfo.collectTransitiveDependenciesWithSelf().filterIsInstance<ModuleSourceInfo>()
+    )
+
+    private val providerHelper = FirProviderHelper(
+        cache,
+        firFileBuilder,
+        indexHelper,
+        packageExistenceChecker,
+    )
+
+    override val isPhasedFirAllowed: Boolean get() = true
+
+    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>? =
+        providerHelper.getFirClassifierByFqName(classId)
+
+    override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> =
+        providerHelper.getTopLevelCallableSymbols(packageFqName, name)
+
+    override fun getPackage(fqName: FqName): FqName? =
+        providerHelper.getPackage(fqName)
+
+    override fun getNestedClassifierScope(classId: ClassId): FirScope? =
+        providerHelper.getNestedClassifierScope(classId)
+
+    override fun getClassLikeSymbolByFqName(classId: ClassId): FirClassLikeSymbol<*>? {
+        return getFirClassifierByFqName(classId)?.symbol
+    }
+
+    override fun getFirClassifierContainerFile(fqName: ClassId): FirFile {
+        return getFirClassifierContainerFileIfAny(fqName)
+            ?: error("Couldn't find container for ${fqName}")
+    }
+
+    override fun getFirClassifierContainerFileIfAny(fqName: ClassId): FirFile? {
+        val fir = getFirClassifierByFqName(fqName) ?: return null // Necessary to ensure cacheProvider contains this classifier
+        return providerHelper.getContainingFirFile(fir.symbol)
+    }
+
+    override fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile {
+        return getFirClassifierContainerFileIfAny(symbol)
+            ?: error("Couldn't find container for ${symbol.classId}")
+    }
+
+    override fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile? =
+        providerHelper.getContainingFirFile(symbol)
+
+
+    override fun getFirCallableContainerFile(symbol: FirCallableSymbol<*>): FirFile? {
+        symbol.overriddenSymbol?.let {
+            return getFirCallableContainerFile(it)
+        }
+        if (symbol is FirAccessorSymbol) {
+            val fir = symbol.fir
+            if (fir is FirSyntheticProperty) {
+                return getFirCallableContainerFile(fir.getter.delegate.symbol)
+            }
+        }
+        return providerHelper.getContainingFirFile(symbol)
+    }
+
+    override fun getFirFilesByPackage(fqName: FqName): List<FirFile> = error("Should not be called in FIR IDE")
+
+
+    // TODO move out of here
+    // used only for completion
+    fun buildFunctionWithBody(ktNamedFunction: KtNamedFunction): FirFunction<*> {
+        return RawFirBuilder(session, kotlinScopeProvider, stubMode = false).buildFunctionWithBody(ktNamedFunction)
+    }
+
+
+    @FirProviderInternals
+    override fun recordGeneratedClass(owner: FirAnnotatedDeclaration, klass: FirRegularClass) {
+        TODO()
+    }
+
+    @FirProviderInternals
+    override fun recordGeneratedMember(owner: FirAnnotatedDeclaration, klass: FirDeclaration) {
+        TODO()
+    }
+
+    // TODO this should be reworked because [FirIdeProvider] should not have such method
+    // used only in completion
+    override fun getAllCallableNamesInPackage(fqName: FqName): Set<Name> {
+        return hashSetOf<Name>().apply {
+            indexHelper.getTopLevelPropertiesInPackage(fqName).mapNotNullTo(this) { it.nameAsName }
+            indexHelper.getTopLevelFunctionsInPackage(fqName).mapNotNullTo(this) { it.nameAsName }
+        }
+    }
+}
+
+internal val FirSession.firIdeProvider: FirIdeProvider by FirSession.sessionComponentAccessor()

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

@@ -0,0 +1,95 @@
+/*
+ * 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.providers
+
+import com.google.common.collect.Sets
+import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.impl.nestedClassifierScope
+import org.jetbrains.kotlin.fir.symbols.CallableId
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.FirElementFinder
+import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper
+import org.jetbrains.kotlin.idea.fir.low.level.api.PackageExistenceChecker
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeOrReturnDefaultValueOnPCE
+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.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtEnumEntry
+import org.jetbrains.kotlin.psi.KtFile
+import java.util.*
+
+internal class FirProviderHelper(
+    private val cache: ModuleFileCache,
+    private val firFileBuilder: FirFileBuilder,
+    private val indexHelper: IndexHelper,
+    private val packageExistenceChecker: PackageExistenceChecker,
+) {
+    fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>? {
+        return executeOrReturnDefaultValueOnPCE(null) {
+            cache.classifierByClassId.computeIfAbsent(classId) {
+                val ktClass = indexHelper.classFromIndexByClassId(classId)
+                    ?: indexHelper.typeAliasFromIndexByClassId(classId)
+                    ?: return@computeIfAbsent Optional.empty()
+                if (ktClass is KtEnumEntry) return@computeIfAbsent Optional.empty()
+                val firFile = firFileBuilder.buildRawFirFileWithCaching(ktClass.containingKtFile, cache)
+                val classifier = FirElementFinder.findElementIn<FirClassLikeDeclaration<*>>(firFile) { classifier ->
+                    classifier.symbol.classId == classId
+                }
+                    ?: error("Classifier $classId was found in file ${ktClass.containingKtFile.virtualFilePath} but was not found in FirFile")
+                Optional.of(classifier)
+            }.getOrNull()
+        }
+    }
+
+
+    fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
+        val callableId = CallableId(packageFqName, name)
+        return executeOrReturnDefaultValueOnPCE(emptyList()) {
+            cache.callableByCallableId.computeIfAbsent(callableId) {
+                val files = Sets.newIdentityHashSet<KtFile>().apply {
+                    indexHelper.getTopLevelFunctions(callableId).mapTo(this) { it.containingKtFile }
+                    indexHelper.getTopLevelProperties(callableId).mapTo(this) { it.containingKtFile }
+                }
+                @OptIn(ExperimentalStdlibApi::class)
+                buildList {
+                    files.forEach { ktFile ->
+                        val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, cache)
+                        firFile.collectCallableDeclarationsTo(this, name)
+                    }
+                }
+            }
+        }
+    }
+
+    fun getContainingFirFile(symbol: FirBasedSymbol<*>): FirFile? {
+        val ktFile = symbol.fir.psi?.containingFile as? KtFile ?: return null
+        return cache.getCachedFirFile(ktFile)
+    }
+
+    private fun FirFile.collectCallableDeclarationsTo(list: MutableList<FirCallableSymbol<*>>, name: Name) {
+        declarations.mapNotNullTo(list) { declaration ->
+            if (declaration is FirCallableDeclaration<*> && declaration.symbol.callableId.callableName == name) {
+                declaration.symbol
+            } else null
+        }
+    }
+
+    fun getNestedClassifierScope(classId: ClassId): FirScope? =
+        (getFirClassifierByFqName(classId) as? FirRegularClass)?.let {
+            nestedClassifierScope(it)
+        }
+
+    fun getPackage(fqName: FqName): FqName? =
+        fqName.takeIf(packageExistenceChecker::isPackageExists)
+}

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

@@ -0,0 +1,10 @@
+/*
+ * 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.providers
+
+import java.util.*
+
+internal fun <T: Any> Optional<T>.getOrNull(): T? = orElse(null)

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

@@ -0,0 +1,73 @@
+/*
+ * 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.sessions
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
+import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
+import org.jetbrains.kotlin.fir.java.deserialization.KotlinDeserializedJvmSymbolsProvider
+import org.jetbrains.kotlin.fir.registerCommonComponents
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.impl.FirBuiltinSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCloneableSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCompositeSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
+import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
+import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.caches.resolve.IDEPackagePartProvider
+import org.jetbrains.kotlin.idea.search.minus
+import org.jetbrains.kotlin.load.java.JavaClassFinderImpl
+import org.jetbrains.kotlin.load.kotlin.VirtualFileFinderFactory
+
+class FirIdeModuleLibraryDependenciesSession private constructor(
+    sessionProvider: FirProjectSessionProvider,
+) : FirSession(sessionProvider) {
+    companion object {
+        fun create(
+            moduleInfo: ModuleSourceInfo,
+            sessionProvider: FirProjectSessionProvider,
+            project: Project,
+        ): FirIdeModuleLibraryDependenciesSession {
+            val searchScope = moduleInfo.module.moduleWithLibrariesScope
+            val javaClassFinder = JavaClassFinderImpl().apply {
+                setProjectInstance(project)
+                setScope(searchScope)
+            }
+            val packagePartProvider = IDEPackagePartProvider(searchScope)
+
+            val kotlinClassFinder = VirtualFileFinderFactory.getInstance(project).create(searchScope)
+            return FirIdeModuleLibraryDependenciesSession(sessionProvider).apply {
+                registerCommonComponents()
+
+                val javaSymbolProvider = JavaSymbolProvider(this, sessionProvider.project, searchScope)
+
+                val kotlinScopeProvider = KotlinScopeProvider(::wrapScopeWithJvmMapped)
+
+                register(
+                    FirSymbolProvider::class,
+                    FirCompositeSymbolProvider(
+                        listOf(
+                            KotlinDeserializedJvmSymbolsProvider(
+                                this,
+                                project,
+                                packagePartProvider,
+                                javaSymbolProvider,
+                                kotlinClassFinder,
+                                javaClassFinder,
+                                kotlinScopeProvider
+                            ),
+                            FirBuiltinSymbolProvider(this, kotlinScopeProvider),
+                            FirCloneableSymbolProvider(this, kotlinScopeProvider),
+                            javaSymbolProvider,
+                        )
+                    )
+                )
+            }
+        }
+    }
+}

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

@@ -0,0 +1,31 @@
+/*
+ * 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.fir.FirElement
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.psi.KtElement
+
+object FirElementFinder {
+    inline fun <reified E : FirElement> findElementIn(container: FirElement, crossinline predicate: (E) -> Boolean): E? {
+        var result: E? = null
+        container.accept(object : FirVisitorVoid() {
+            override fun visitElement(element: FirElement) {
+                if (result != null) return
+                if (element is E && predicate(element)) {
+                    result = element
+                } else {
+                    element.acceptChildren(this)
+                }
+            }
+        })
+        return result
+    }
+
+    inline fun <reified E : FirElement> findElementByPsiIn(container: FirElement, ktElement: KtElement): E? =
+        findElementIn(container) { it.psi === ktElement }
+}

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

@@ -0,0 +1,38 @@
+/*
+ * 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 com.intellij.openapi.progress.ProcessCanceledException
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.diagnostics.FirDiagnosticHolder
+import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+
+
+internal inline fun <T> executeOrReturnDefaultValueOnPCE(defaultValue: T, action: () -> T): T =
+    try {
+        action()
+    } catch (e: ProcessCanceledException) {
+        defaultValue
+    }
+
+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>()
+    fun collect(module: IdeaModuleInfo) {
+        if (module in result) return
+        result += module
+        module.dependencies().forEach(::collect)
+    }
+    collect(this)
+    return result.toList()
+}

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

@@ -1,12 +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
-
-import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.diagnostics.FirDiagnosticHolder
-
-internal val FirElement.isErrorElement
-    get() = this is FirDiagnosticHolder

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/AbstractFirLazyResolveTest.kt

@@ -13,11 +13,13 @@ import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.search.FileTypeIndex
 import com.intellij.testFramework.LightProjectDescriptor
 import junit.framework.TestCase
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.productionSourceInfo
+import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.jsonUtils.getString
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
@@ -64,7 +66,7 @@ abstract class AbstractFirLazyResolveTest : KotlinLightCodeInsightFixtureTestCas
         } as KtExpression
 
         val resolveState = expressionToResolve.firResolveState()
-        val resultsDump = when (val firElement = expressionToResolve.getOrBuildFir(resolveState)) {
+        val resultsDump = when (val firElement = resolveState.getOrBuildFirFor(expressionToResolve, FirResolvePhase.BODY_RESOLVE)) {
             is FirResolvedImport -> buildString {
                 append("import ")
                 append(firElement.packageFqName)
@@ -100,9 +102,9 @@ abstract class AbstractFirLazyResolveTest : KotlinLightCodeInsightFixtureTestCas
         val files = FileTypeIndex.getFiles(KotlinFileType.INSTANCE, contentScope)
         for (file in files) {
             val psiFile = psiManager.findFile(file) as KtFile
-            val session = resolveState.getSession(psiFile)
+            val session = resolveState.firSession
             val firProvider = session.firIdeProvider
-            val firFile = firProvider.getFile(psiFile) ?: continue
+            val firFile = firProvider.cache.getCachedFirFile(psiFile) ?: continue
             KotlinTestUtils.assertEqualsToFile(File(expectedTxtPath(file)), firFile.render())
         }
     }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirUtils.kt

@@ -29,7 +29,6 @@ fun KtElement.getOrBuildFir(
     phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
 ) = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState, phase)
 
-fun KtElement.getFirOfClosestParent() = LowLevelFirApiFacade.getFirOfClosestParent(this)
 
 inline fun <reified E : FirElement> KtElement.getOrBuildFirSafe(
     resolveState: FirModuleResolveState,

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

@@ -22,7 +22,6 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.fir.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
-import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.frontend.api.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.scopes.KtFirScopeProvider
@@ -126,7 +125,7 @@ private constructor(
 
     override fun createContextDependentCopy(): KtAnalysisSession {
         check(!isContextSession) { "Cannot create context-dependent copy of KtAnalysis session from a context dependent one" }
-        val contextResolveState = LowLevelFirApiFacade.getResolveStateForCompletion(element, firResolveState as FirModuleResolveStateImpl)
+        val contextResolveState = LowLevelFirApiFacade.getResolveStateForCompletion(element, firResolveState)
         return KtFirAnalysisSession(
             element,
             firSession,
@@ -148,7 +147,7 @@ private constructor(
     }
 
     private fun resolveCall(firCall: FirFunctionCall, callExpression: KtExpression): CallInfo? {
-        val session = LowLevelFirApiFacade.getSessionFor(callExpression, firResolveState)
+        val session = firResolveState.firSession
         val resolvedFunctionSymbol = firCall.calleeReference.toTargetSymbol(session, firSymbolBuilder)
         val resolvedCalleeSymbol = (firCall.calleeReference as? FirResolvedNamedReference)?.resolvedSymbol
         return when {
@@ -203,7 +202,7 @@ private constructor(
         @Deprecated("Please use org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProviderKt.analyze")
         fun createForElement(element: KtElement): KtFirAnalysisSession {
             val firResolveState = LowLevelFirApiFacade.getResolveStateFor(element)
-            val firSession = LowLevelFirApiFacade.getSessionFor(element, firResolveState)
+            val firSession = firResolveState.firSession
             val project = element.project
             val typeContext = ConeTypeCheckerContext(isErrorTypeEqualsToAnything = true, isStubTypeEqualsToAnything = true, firSession)
             val token = ReadActionConfinementValidityToken(project)

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/testUtils.kt

@@ -11,7 +11,7 @@ import com.intellij.openapi.util.io.FileUtil
 import com.intellij.testFramework.LightPlatformTestCase
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
-import org.jetbrains.kotlin.idea.fir.low.level.api.DuplicatedFirSourceElementsException
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.DuplicatedFirSourceElementsException
 import java.io.File
 
 fun <R> executeOnPooledThreadInReadAction(action: () -> R): R =

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/util/ApplicationUtils.kt

@@ -1,17 +1,6 @@
 /*
- * Copyright 2010-2015 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
+ * 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.util.application
@@ -30,6 +19,15 @@ fun <T> runWriteAction(action: () -> T): T {
     return ApplicationManager.getApplication().runWriteAction<T>(action)
 }
 
+fun <T> runWriteActionInEdt(action: () -> T): T {
+    var result: T? = null
+    ApplicationManager.getApplication().invokeLater {
+        result = ApplicationManager.getApplication().runWriteAction<T>(action)
+    }
+    return result!!
+}
+
+
 fun Project.executeWriteCommand(name: String, command: () -> Unit) {
     CommandProcessor.getInstance().executeCommand(this, { runWriteAction(command) }, name, null)
 }

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/util/PsiModificationTrackerBasedCachedValue.kt

@@ -0,0 +1,41 @@
+/*
+ * 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.util
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.util.CachedValueProvider
+import com.intellij.psi.util.CachedValuesManager
+import com.intellij.psi.util.PsiModificationTracker
+import kotlin.properties.ReadOnlyProperty
+import kotlin.reflect.KProperty
+
+/**
+ * Creates a value which will be cached until until any physical PSI change happens
+ * To create one please use [psiModificationTrackerBasedCachedValue]
+ *
+ * @see com.intellij.psi.util.CachedValue
+ * @see com.intellij.psi.util.PsiModificationTracker.MODIFICATION_COUNT
+ */
+class PsiModificationTrackerBasedCachedValue<T>(project: Project, createValue: () -> T) : ReadOnlyProperty<Any?, T> {
+    private val cachedValue = CachedValuesManager.getManager(project).createCachedValue {
+        CachedValueProvider.Result(
+            createValue(),
+            PsiModificationTracker.MODIFICATION_COUNT
+        )
+    }
+
+    override fun getValue(thisRef: Any?, property: KProperty<*>): T = cachedValue.value
+}
+
+/**
+ * Creates a value which will be cached until until any physical PSI change happens
+ *
+ * @see com.intellij.psi.util.CachedValue
+ * @see com.intellij.psi.util.PsiModificationTracker.MODIFICATION_COUNT
+ * @see PsiModificationTrackerBasedCachedValue
+ */
+fun <T> psiModificationTrackerBasedCachedValue(project: Project, createValue: () -> T): PsiModificationTrackerBasedCachedValue<T> =
+    PsiModificationTrackerBasedCachedValue(project, createValue)

idea/resources-fir/META-INF/plugin.xml

@@ -88,8 +88,7 @@ The Kotlin FIR plugin provides language support in IntelliJ IDEA and Android Stu
     <applicationService serviceImplementation="org.jetbrains.kotlin.idea.PluginStartupService"/>
     <highlightRangeExtension implementation="org.jetbrains.kotlin.idea.fir.highlighter.KotlinFirPsiChecker"/>
     <annotator language="kotlin" implementationClass="org.jetbrains.kotlin.idea.fir.highlighter.KotlinFirPsiChecker"/>
-    <projectService serviceInterface="org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateService"
-                    serviceImplementation="org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateServiceImpl"/>
+    <projectService serviceImplementation="org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateService"/>
     <projectService serviceImplementation="org.jetbrains.kotlin.resolve.jvm.KotlinJavaPsiFacade"/>
 
     <completion.contributor language="kotlin"
@@ -157,10 +156,16 @@ The Kotlin FIR plugin provides language support in IntelliJ IDEA and Android Stu
     <projectService serviceInterface="org.jetbrains.kotlin.idea.framework.LibraryEffectiveKindProvider"
                     serviceImplementation="org.jetbrains.kotlin.idea.framework.LibraryEffectiveKindProviderImpl"/>
 
+    <applicationService
+        serviceInterface="org.jetbrains.kotlin.platform.DefaultIdeTargetPlatformKindProvider"
+        serviceImplementation="org.jetbrains.kotlin.platform.impl.IdeaDefaultIdeTargetPlatformKindProvider"/>
+
   </extensions>
 
 
   <extensions defaultExtensionNs="com.intellij">
+
+    <projectService serviceImplementation="org.jetbrains.kotlin.idea.caches.resolve.IdePackageOracleFactory"/>
     <projectService serviceImplementation="org.jetbrains.kotlin.idea.caches.trackers.KotlinCodeBlockModificationListener"/>
 
     <projectService serviceInterface="org.jetbrains.kotlin.load.kotlin.MetadataFinderFactory"