[low level fir] lazyly get KtSymbol for builtins

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/utils/ktSymbolUtils.kt

@@ -8,12 +8,12 @@ package org.jetbrains.kotlin.analysis.api.fir.utils
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirModuleResolveState
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.KtModuleBasedModuleData
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirBuiltinsModuleData
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirModuleSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirKtModuleBasedModuleData
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.firModuleData
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.fir.dispatchReceiverClassOrNull
-import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
 import org.jetbrains.kotlin.fir.resolve.toFirRegularClassSymbol
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
@@ -27,15 +27,17 @@ val KtSymbol.firSymbol: FirBasedSymbol<*>
 
 fun FirBasedSymbol<*>.getContainingKtModule(resolveState: LLFirModuleResolveState): KtModule {
     val target = when (this) {
-        // callable fake overrides have use-site FirModuleData
+        is FirCallableSymbol -> {
-        is FirCallableSymbol -> dispatchReceiverClassOrNull()?.toFirRegularClassSymbol(resolveState.rootModuleSession)
+            // callable fake overrides have use-site FirModuleData
-            ?: this
+            dispatchReceiverClassOrNull()?.toFirRegularClassSymbol(resolveState.rootModuleSession) ?: this
-
+        }
         else -> this
     }
-    val moduleData = target.moduleData
+    return when (val moduleData = target.firModuleData) {
-    requireIsInstance<KtModuleBasedModuleData>(moduleData)
+        is LLFirKtModuleBasedModuleData -> moduleData.ktModule
-    return moduleData.module
+        is LLFirBuiltinsModuleData -> resolveState.project.getService(ProjectStructureProvider::class.java).getStdlibWithBuiltinsModule(moduleData.useSiteKtModule)
+            ?: error("Builtins not found for the ${moduleData.useSiteKtModule.moduleDescription}")
+    }
 }
 
 fun KtSymbol.getContainingKtModule(resolveState: LLFirModuleResolveState): KtModule =

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostics/AbstractFirIdeDiagnosticsCollector.kt

@@ -5,7 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics
 
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.module
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.firKtModuleBasedModuleData
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.CheckersComponentInternal
@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
 import org.jetbrains.kotlin.fir.analysis.collectors.components.*
 import org.jetbrains.kotlin.fir.analysis.jvm.checkers.JvmDeclarationCheckers
 import org.jetbrains.kotlin.fir.analysis.jvm.checkers.JvmExpressionCheckers
-import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.platform.SimplePlatform
 import org.jetbrains.kotlin.platform.jvm.JvmPlatform
 
@@ -40,7 +39,7 @@ private object CheckersFactory {
         reporter: DiagnosticReporter,
         useExtendedCheckers: Boolean
     ): List<AbstractDiagnosticCollectorComponent> {
-        val module = session.moduleData.module
+        val module = session.firKtModuleBasedModuleData.ktModule
         val platform = module.platform.componentPlatforms.first()
         val declarationCheckers = createDeclarationCheckers(useExtendedCheckers, platform)
         val expressionCheckers = createExpressionCheckers(useExtendedCheckers, platform)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/LLFirModuleData.kt

@@ -0,0 +1,101 @@
+/*
+ * Copyright 2010-2022 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.analysis.low.level.api.fir.project.structure
+
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.fir.FirModuleData
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.moduleData
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
+
+sealed class LLFirModuleData : FirModuleData()
+
+val FirDeclaration.firModuleData: LLFirModuleData
+    get() {
+        return moduleData as LLFirModuleData
+    }
+
+val FirSession.firModuleData: LLFirModuleData
+    get() {
+        return moduleData as LLFirModuleData
+    }
+
+val FirSession.firKtModuleBasedModuleData: LLFirKtModuleBasedModuleData
+    get() {
+        return moduleData as LLFirKtModuleBasedModuleData
+    }
+
+
+val FirBasedSymbol<*>.firModuleData: LLFirModuleData
+    get() = fir.firModuleData
+
+class LLFirBuiltinsModuleData(val useSiteKtModule: KtModule) : LLFirModuleData() {
+    override val name: Name
+        get() = Name.special("<builtins for ${useSiteKtModule.moduleDescription}>")
+
+    override val dependencies: List<FirModuleData> get() = emptyList()
+    override val dependsOnDependencies: List<FirModuleData> get() = emptyList()
+    override val friendDependencies: List<FirModuleData> get() = emptyList()
+
+    override val platform: TargetPlatform get() = useSiteKtModule.platform
+    override val analyzerServices: PlatformDependentAnalyzerServices get() = useSiteKtModule.analyzerServices
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as LLFirBuiltinsModuleData
+
+        if (useSiteKtModule != other.useSiteKtModule) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        return useSiteKtModule.hashCode()
+    }
+}
+
+class LLFirKtModuleBasedModuleData(
+    val ktModule: KtModule,
+) : LLFirModuleData() {
+    override val name: Name get() = Name.special("<${ktModule.moduleDescription}>")
+
+    override val dependencies: List<FirModuleData> by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        ktModule.directRegularDependencies.map(::LLFirKtModuleBasedModuleData)
+    }
+
+    override val dependsOnDependencies: List<FirModuleData> by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        ktModule.directRefinementDependencies.map(::LLFirKtModuleBasedModuleData)
+    }
+
+    override val friendDependencies: List<FirModuleData> by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        ktModule.directRefinementDependencies.map(::LLFirKtModuleBasedModuleData)
+    }
+
+    override val platform: TargetPlatform get() = ktModule.platform
+
+    override val analyzerServices: PlatformDependentAnalyzerServices get() = ktModule.analyzerServices
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as LLFirKtModuleBasedModuleData
+
+        if (ktModule != other.ktModule) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        return ktModule.hashCode()
+    }
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/KtModuleBasedModuleData.kt

@@ -11,39 +11,3 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
 
-class KtModuleBasedModuleData(
-    val module: KtModule,
-) : FirModuleData() {
-    override val name: Name get() = Name.special("<${module.moduleDescription}>")
-
-    override val dependencies: List<FirModuleData> by lazy(LazyThreadSafetyMode.PUBLICATION) {
-        module.directRegularDependencies.map(::KtModuleBasedModuleData)
-    }
-
-    override val dependsOnDependencies: List<FirModuleData> by lazy(LazyThreadSafetyMode.PUBLICATION) {
-        module.directRefinementDependencies.map(::KtModuleBasedModuleData)
-    }
-
-    override val friendDependencies: List<FirModuleData> by lazy(LazyThreadSafetyMode.PUBLICATION) {
-        module.directRefinementDependencies.map(::KtModuleBasedModuleData)
-    }
-
-    override val platform: TargetPlatform get() = module.platform
-
-    override val analyzerServices: PlatformDependentAnalyzerServices get() = module.analyzerServices
-
-    override fun equals(other: Any?): Boolean {
-        if (this === other) return true
-        if (javaClass != other?.javaClass) return false
-
-        other as KtModuleBasedModuleData
-
-        if (module != other.module) return false
-
-        return true
-    }
-
-    override fun hashCode(): Int {
-        return module.hashCode()
-    }
-}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionFactory.kt

@@ -17,12 +17,15 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.FirFileBuild
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.ModuleFileCacheImpl
 import org.jetbrains.kotlin.analysis.low.level.api.fir.fir.caches.FirThreadSafeCachesFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyDeclarationResolver
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirKtModuleBasedModuleData
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirBuiltinsModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.analysis.providers.createAnnotationResolver
 import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.createPackageProvider
+import org.jetbrains.kotlin.analysis.utils.errors.checkIsInstance
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.config.LanguageVersionSettingsImpl
 import org.jetbrains.kotlin.fir.*
@@ -88,7 +91,7 @@ internal object LLFirSessionFactory {
         sessionsCache[module] = session
 
         return session.apply session@{
-            val moduleData = KtModuleBasedModuleData(module).apply { bindSession(this@session) }
+            val moduleData = LLFirKtModuleBasedModuleData(module).apply { bindSession(this@session) }
             registerModuleData(moduleData)
             register(FirKotlinScopeProvider::class, scopeProvider)
 
@@ -216,7 +219,7 @@ internal object LLFirSessionFactory {
         checkCanceled()
         val searchScope = project.moduleScopeProvider.getModuleLibrariesScope(sourceModule)
         LLFirLibrariesSession(project, builtinTypes).apply session@{
-            registerModuleData(KtModuleBasedModuleData(sourceModule).apply { bindSession(this@session) })
+            registerModuleData(LLFirKtModuleBasedModuleData(sourceModule).apply { bindSession(this@session) })
             registerIdeComponents(project)
             register(FirPhaseManager::class, FirPhaseCheckingPhaseManager)
             registerCommonComponents(languageVersionSettings)
@@ -251,11 +254,11 @@ internal object LLFirSessionFactory {
 
     private fun createModuleDataProviderWithLibraryDependencies(sourceModule: KtModule, session: LLFirSession): ModuleDataProvider {
         val regularDependenciesOnLibs =
-            sourceModule.directRegularDependenciesOfType<KtBinaryModule>().map { KtModuleBasedModuleData(it) }
+            sourceModule.directRegularDependenciesOfType<KtBinaryModule>().map { LLFirKtModuleBasedModuleData(it) }
         val friendDependenciesOnLibs =
-            sourceModule.directFriendDependenciesOfType<KtBinaryModule>().map { KtModuleBasedModuleData(it) }
+            sourceModule.directFriendDependenciesOfType<KtBinaryModule>().map { LLFirKtModuleBasedModuleData(it) }
         val dependsOnDependenciesOnLibs =
-            sourceModule.directRefinementDependenciesOfType<KtBinaryModule>().map { KtModuleBasedModuleData(it) }
+            sourceModule.directRefinementDependenciesOfType<KtBinaryModule>().map { LLFirKtModuleBasedModuleData(it) }
 
         val allDependencies = buildList {
             addAll(regularDependenciesOnLibs)
@@ -271,7 +274,8 @@ internal object LLFirSessionFactory {
 
         val moduleDataWithFilters: Map<FirModuleData, LibraryPathFilter.LibraryList> =
             allDependencies.associateWith { moduleData ->
-                val ktBinaryModule = moduleData.module as KtBinaryModule
+                checkIsInstance<LLFirKtModuleBasedModuleData>(moduleData)
+                val ktBinaryModule = moduleData.ktModule as KtBinaryModule
                 val moduleBinaryRoots = ktBinaryModule.getBinaryRoots().mapTo(mutableSetOf()) { it.toAbsolutePath() }
                 LibraryPathFilter.LibraryList(moduleBinaryRoots)
             }
@@ -282,12 +286,12 @@ internal object LLFirSessionFactory {
     fun createBuiltinsAndCloneableSession(
         project: Project,
         builtinTypes: BuiltinTypes,
-        stdlibModule: KtModule,
+        useSiteModule: KtModule,
         languageVersionSettings: LanguageVersionSettings = LanguageVersionSettingsImpl.DEFAULT,
         configureSession: (LLFirSession.() -> Unit)? = null,
     ): LLFirBuiltinsAndCloneableSession {
         return LLFirBuiltinsAndCloneableSession(project, builtinTypes).apply session@{
-            val moduleData = KtModuleBasedModuleData(stdlibModule).apply {
+            val moduleData = LLFirBuiltinsModuleData(useSiteModule).apply {
                 bindSession(this@session)
             }
             registerIdeComponents(project)
@@ -339,7 +343,7 @@ internal object LLFirSessionFactory {
         sessionsCache[module] = session
 
         return session.apply session@{
-            val moduleData = KtModuleBasedModuleData(module).apply { bindSession(this@session) }
+            val moduleData = LLFirKtModuleBasedModuleData(module).apply { bindSession(this@session) }
             registerModuleData(moduleData)
             register(FirKotlinScopeProvider::class, scopeProvider)
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProvider.kt

@@ -8,8 +8,10 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analysis.api.impl.barebone.annotations.Immutable
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirKtModuleBasedModuleData
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.NoCacheForModuleException
+import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
 import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionProvider
@@ -20,8 +22,11 @@ class LLFirSessionProvider internal constructor(
     internal val rootModuleSession: LLFirResolvableModuleSession,
     private val moduleToSession: Map<KtModule, LLFirResolvableModuleSession>
 ) : FirSessionProvider() {
-    override fun getSession(moduleData: FirModuleData): FirSession? =
+
-        moduleToSession[moduleData.module]
+    override fun getSession(moduleData: FirModuleData): FirSession? {
+        requireIsInstance<LLFirKtModuleBasedModuleData>(moduleData)
+        return moduleToSession[moduleData.ktModule]
+    }
 
     fun getSession(module: KtModule): FirSession? =
         moduleToSession[module]

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProviderStorage.kt

@@ -11,6 +11,7 @@ import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.persistentMapOf
 import kotlinx.collections.immutable.toPersistentMap
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirPhaseRunner
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.firKtModuleBasedModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.addValueFor
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.executeWithoutPCE
 import org.jetbrains.kotlin.analysis.project.structure.*
@@ -19,7 +20,6 @@ import org.jetbrains.kotlin.analysis.providers.createModuleWithoutDependenciesOu
 import org.jetbrains.kotlin.analysis.utils.caches.getValue
 import org.jetbrains.kotlin.analysis.utils.caches.softCachedValue
 import org.jetbrains.kotlin.fir.BuiltinTypes
-import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.moduleData
 import java.util.concurrent.ConcurrentHashMap
 
@@ -35,9 +35,8 @@ class LLFirSessionProviderStorage(val project: Project) {
         val firPhaseRunner = LLFirPhaseRunner()
 
         val builtinTypes = BuiltinTypes()
-        val stdlibModule = project.getService(ProjectStructureProvider::class.java).getStdlibWithBuiltinsModule(rootModule)
+
-            ?: error("Stdlib was not found for ${rootModule.moduleDescription}")
+        val builtinsAndCloneableSession = LLFirSessionFactory.createBuiltinsAndCloneableSession(project, builtinTypes, rootModule)
-        val builtinsAndCloneableSession = LLFirSessionFactory.createBuiltinsAndCloneableSession(project, builtinTypes, stdlibModule)
         val cache = sessionsCache.getOrPut(rootModule) { FromModuleViewSessionCache() }
         val (sessions, session) = cache.withMappings(project) { mappings ->
             val sessions = mutableMapOf<KtModule, LLFirResolvableModuleSession>().apply { putAll(mappings) }
@@ -85,7 +84,7 @@ private class FromModuleViewSessionCache {
     private var mappings: PersistentMap<KtModule, FirSessionWithModificationTracker> = persistentMapOf()
 
     val sessionInvalidator: LLFirSessionInvalidator = LLFirSessionInvalidator { session ->
-        mappings[session.moduleData.module]?.invalidate()
+        mappings[session.firKtModuleBasedModuleData.ktModule]?.invalidate()
     }
 
 
@@ -125,7 +124,7 @@ private class FromModuleViewSessionCache {
         }
         return wasSessionInvalidated.entries
             .mapNotNull { (session, wasInvalidated) -> session.takeUnless { wasInvalidated } }
-            .associate { session -> session.firSession.moduleData.module to session.firSession }
+            .associate { session -> session.firSession.firKtModuleBasedModuleData.ktModule to session.firSession }
     }
 
     private fun <T> Collection<T>.reversedDependencies(getDependencies: (T) -> List<T>): Map<T, List<T>> {
@@ -144,11 +143,12 @@ private class FirSessionWithModificationTracker(
     val firSession: LLFirResolvableModuleSession,
 ) {
     private val modificationTracker =
-        when (val moduleInfo = firSession.moduleData.module) {
+        when (val ktModule = firSession.firKtModuleBasedModuleData.ktModule) {
-            is KtSourceModule -> moduleInfo.createModuleWithoutDependenciesOutOfBlockModificationTracker(project)
+            is KtSourceModule -> ktModule.createModuleWithoutDependenciesOutOfBlockModificationTracker(project)
             else -> ModificationTracker.NEVER_CHANGED
         }
 
+
     private val timeStamp = modificationTracker.modificationCount
 
     @Volatile
@@ -159,9 +159,4 @@ private class FirSessionWithModificationTracker(
     }
 
     val isValid: Boolean get() = !isInvalidated && modificationTracker.modificationCount == timeStamp
-}
+}
-
-internal val FirModuleData.module: KtModule get() = moduleUnsafe()
-
-internal inline fun <reified T : KtModule> FirModuleData.moduleUnsafe(): T = (this as KtModuleBasedModuleData).module as T
-internal inline fun <reified T : KtModule> FirModuleData.moduleInfoSafe(): T? = (this as KtModuleBasedModuleData).module as? T