[low level fir] lazyly get KtSymbol for builtins

2022-04-11 14:33:59 +02:00
parent 66df69ea2c
commit 0fa791e669
7 changed files with 145 additions and 75 deletions
@@ -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.sessions.LLFirModuleSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirBuiltinsModuleData
+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 -> dispatchReceiverClassOrNull()?.toFirRegularClassSymbol(resolveState.rootModuleSession)
-            ?: this
-
+        is FirCallableSymbol -> {
+            // callable fake overrides have use-site FirModuleData
+            dispatchReceiverClassOrNull()?.toFirRegularClassSymbol(resolveState.rootModuleSession) ?: this
+        }
        else -> this
    }
-    val moduleData = target.moduleData
-    requireIsInstance<KtModuleBasedModuleData>(moduleData)
-    return moduleData.module
+    return when (val moduleData = target.firModuleData) {
+        is LLFirKtModuleBasedModuleData -> moduleData.ktModule
+        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 =
@@ -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)
@@ -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()
+    }
+}
@@ -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()
-    }
-}
@@ -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)

@@ -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]
@@ -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, stdlibModule)
+
+        val builtinsAndCloneableSession = LLFirSessionFactory.createBuiltinsAndCloneableSession(project, builtinTypes, rootModule)
        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) {
-            is KtSourceModule -> moduleInfo.createModuleWithoutDependenciesOutOfBlockModificationTracker(project)
+        when (val ktModule = firSession.firKtModuleBasedModuleData.ktModule) {
+            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
+}