[FIR] Store mapped jvm declarations in session component, share it between MPP modules

^KT-61640 Fixed

This change doesn't affect any tests in the moment, because it's part of
  bigger refactoring of proper storing IR declarations during FIR2IR
  conversion (KT-61637)

Without this change, there is one test break in branch for KT-61637:
- FirPsiBlackBoxCodegenTestGenerated.Multiplatform.K2.testJavaMethodWithTypeParameter

compiler/cli/src/org/jetbrains/kotlin/cli/common/FirSessionConstructionUtils.kt

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.analysis.checkers.OptInLanguageVersionSettingsCh
 import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
-import org.jetbrains.kotlin.fir.java.enhancement.FirEnhancedSymbolsStorage
 import org.jetbrains.kotlin.fir.session.*
 import org.jetbrains.kotlin.fir.session.environment.AbstractProjectEnvironment
 import org.jetbrains.kotlin.fir.session.environment.AbstractProjectFileSearchScope
@@ -30,7 +29,6 @@ import org.jetbrains.kotlin.load.kotlin.PackageAndMetadataPartProvider
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.CommonPlatforms
 import org.jetbrains.kotlin.platform.TargetPlatform
-import org.jetbrains.kotlin.platform.isCommon
 import org.jetbrains.kotlin.platform.js.JsPlatforms
 import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.platform.konan.NativePlatforms
@@ -73,7 +71,7 @@ fun <F> prepareJvmSessions(
     createProviderAndScopeForIncrementalCompilation: (List<F>) -> IncrementalCompilationContext?,
 ): List<SessionWithSources<F>> {
     val javaSourcesScope = projectEnvironment.getSearchScopeForProjectJavaSources()
-    val enhancementSymbolsStorage = FirEnhancedSymbolsStorage(firCachesFactoryForCliMode)
+    val predefinedJavaComponents = FirSharableJavaComponents(firCachesFactoryForCliMode)
     return prepareSessions(
         files, configuration, rootModuleName, JvmPlatforms.unspecifiedJvmPlatform,
         JvmPlatformAnalyzerServices, metadataCompilationMode = false, libraryList, isCommonSource, fileBelongsToModule,
@@ -87,7 +85,7 @@ fun <F> prepareJvmSessions(
                 librariesScope,
                 projectEnvironment.getPackagePartProvider(librariesScope),
                 configuration.languageVersionSettings,
-                predefinedEnhancementStorage = enhancementSymbolsStorage,
+                predefinedJavaComponents = predefinedJavaComponents,
                 registerExtraComponents = {},
             )
         }
@@ -103,7 +101,7 @@ fun <F> prepareJvmSessions(
             configuration.get(CommonConfigurationKeys.LOOKUP_TRACKER),
             configuration.get(CommonConfigurationKeys.ENUM_WHEN_TRACKER),
             configuration.get(CommonConfigurationKeys.IMPORT_TRACKER),
-            predefinedEnhancementStorage = enhancementSymbolsStorage,
+            predefinedJavaComponents = predefinedJavaComponents,
             needRegisterJavaElementFinder = true,
             registerExtraComponents = {},
             sessionConfigurator,

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt

@@ -49,6 +49,7 @@ import org.jetbrains.kotlin.fir.scopes.FirOverrideService
 import org.jetbrains.kotlin.fir.scopes.FirPlatformClassMapper
 import org.jetbrains.kotlin.fir.scopes.PlatformSpecificOverridabilityRules
 import org.jetbrains.kotlin.fir.scopes.impl.*
+import org.jetbrains.kotlin.fir.scopes.jvm.JvmMappedScope.FirMappedSymbolStorage
 import org.jetbrains.kotlin.fir.serialization.FirProvidedDeclarationsForMetadataService
 import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.fir.types.*
@@ -104,14 +105,25 @@ fun FirSession.registerCliCompilerOnlyComponents() {
     register(FirPredicateBasedProvider::class, FirPredicateBasedProviderImpl(this))
 }
 
+class FirSharableJavaComponents(
+    val enhancementStorage: FirEnhancedSymbolsStorage,
+    val mappedStorage: FirMappedSymbolStorage
+) {
+    constructor(cachesFactory: FirCachesFactory) : this(
+        FirEnhancedSymbolsStorage(cachesFactory),
+        FirMappedSymbolStorage(cachesFactory)
+    )
+}
+
 @OptIn(SessionConfiguration::class)
 fun FirSession.registerCommonJavaComponents(
     javaModuleResolver: JavaModuleResolver,
-    predefinedEnhancementStorage: FirEnhancedSymbolsStorage? = null,
+    predefinedComponents: FirSharableJavaComponents? = null,
 ) {
     val jsr305State = languageVersionSettings.getFlag(JvmAnalysisFlags.javaTypeEnhancementState)
     register(FirAnnotationTypeQualifierResolver::class, FirAnnotationTypeQualifierResolver(this, jsr305State, javaModuleResolver))
-    register(FirEnhancedSymbolsStorage::class, predefinedEnhancementStorage ?: FirEnhancedSymbolsStorage(this))
+    register(FirEnhancedSymbolsStorage::class, predefinedComponents?.enhancementStorage ?: FirEnhancedSymbolsStorage(this))
+    register(FirMappedSymbolStorage::class, predefinedComponents?.mappedStorage ?: FirMappedSymbolStorage(this))
     register(FirSyntheticPropertiesStorage::class, FirSyntheticPropertiesStorage(this))
     register(
         FirJvmDefaultModeComponent::class,

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/FirJvmSessionFactory.kt

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
 import org.jetbrains.kotlin.fir.java.deserialization.JvmClassFileBasedSymbolProvider
 import org.jetbrains.kotlin.fir.java.deserialization.OptionalAnnotationClassesProvider
-import org.jetbrains.kotlin.fir.java.enhancement.FirEnhancedSymbolsStorage
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirBuiltinSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCloneableSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
@@ -42,7 +41,7 @@ object FirJvmSessionFactory : FirAbstractSessionFactory() {
         scope: AbstractProjectFileSearchScope,
         packagePartProvider: PackagePartProvider,
         languageVersionSettings: LanguageVersionSettings,
-        predefinedEnhancementStorage: FirEnhancedSymbolsStorage? = null,
+        predefinedJavaComponents: FirSharableJavaComponents? = null,
         registerExtraComponents: ((FirSession) -> Unit),
     ): FirSession {
         return createLibrarySession(
@@ -52,7 +51,7 @@ object FirJvmSessionFactory : FirAbstractSessionFactory() {
             languageVersionSettings,
             extensionRegistrars,
             registerExtraComponents = {
-                it.registerCommonJavaComponents(projectEnvironment.getJavaModuleResolver(), predefinedEnhancementStorage)
+                it.registerCommonJavaComponents(projectEnvironment.getJavaModuleResolver(), predefinedJavaComponents)
                 registerExtraComponents(it)
             },
             createKotlinScopeProvider = { FirKotlinScopeProvider(::wrapScopeWithJvmMapped) },
@@ -92,7 +91,7 @@ object FirJvmSessionFactory : FirAbstractSessionFactory() {
         lookupTracker: LookupTracker? = null,
         enumWhenTracker: EnumWhenTracker? = null,
         importTracker: ImportTracker? = null,
-        predefinedEnhancementStorage: FirEnhancedSymbolsStorage? = null,
+        predefinedJavaComponents: FirSharableJavaComponents? = null,
         needRegisterJavaElementFinder: Boolean,
         registerExtraComponents: ((FirSession) -> Unit) = {},
         init: FirSessionConfigurator.() -> Unit = {}
@@ -107,7 +106,7 @@ object FirJvmSessionFactory : FirAbstractSessionFactory() {
             importTracker,
             init,
             registerExtraComponents = {
-                it.registerCommonJavaComponents(projectEnvironment.getJavaModuleResolver(), predefinedEnhancementStorage)
+                it.registerCommonJavaComponents(projectEnvironment.getJavaModuleResolver(), predefinedJavaComponents)
                 it.registerJavaSpecificResolveComponents()
                 registerExtraComponents(it)
             },

compiler/fir/java/src/org/jetbrains/kotlin/fir/scopes/jvm/JvmMappedScope.kt

@@ -7,13 +7,11 @@ package org.jetbrains.kotlin.fir.scopes.jvm
 
 import org.jetbrains.kotlin.builtins.jvm.JavaToKotlinClassMap
 import org.jetbrains.kotlin.builtins.jvm.JvmBuiltInsSignatures
-import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.containingClassLookupTag
+import org.jetbrains.kotlin.fir.caches.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.classId
 import org.jetbrains.kotlin.fir.declarations.utils.isFinal
-import org.jetbrains.kotlin.fir.dispatchReceiverClassLookupTagOrNull
-import org.jetbrains.kotlin.fir.isSubstitutionOrIntersectionOverride
 import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.lookupSuperTypes
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
@@ -51,9 +49,7 @@ class JvmMappedScope(
     private val declaredMemberScope: FirContainingNamesAwareScope,
     private val javaMappedClassUseSiteScope: FirTypeScope,
 ) : FirTypeScope() {
-    private val functionsCache = mutableMapOf<FirNamedFunctionSymbol, FirNamedFunctionSymbol>()
+    private val mappedSymbolCache = session.mappedSymbolStorage.cacheByOwner.getValue(firKotlinClass.symbol)
-
-    private val constructorsCache = mutableMapOf<FirConstructorSymbol, FirConstructorSymbol>()
 
     private val overrideChecker = FirStandardOverrideChecker(session)
 
@@ -137,7 +133,7 @@ class JvmMappedScope(
             // hidden methods in final class can't be overridden or called with 'super'
             if (jdkMemberStatus == JDKMemberStatus.HIDDEN && firKotlinClass.isFinal) return@processor
 
-            val newSymbol = getOrCreateSubstitutedCopy(symbol, jdkMemberStatus)
+            val newSymbol = mappedSymbolCache.mappedFunctions.getValue(symbol, this to jdkMemberStatus)
             processor(newSymbol)
         }
     }
@@ -206,7 +202,7 @@ class JvmMappedScope(
         return JDKMemberStatus.HIDDEN
     }
 
-    private enum class JDKMemberStatus {
+    internal enum class JDKMemberStatus {
         HIDDEN, VISIBLE, DROP
     }
 
@@ -214,26 +210,24 @@ class JvmMappedScope(
         declaredMemberScope.processPropertiesByName(name, processor)
     }
 
-    private fun getOrCreateSubstitutedCopy(symbol: FirNamedFunctionSymbol, jdkMemberStatus: JDKMemberStatus): FirNamedFunctionSymbol {
+    private fun createMappedFunction(symbol: FirNamedFunctionSymbol, jdkMemberStatus: JDKMemberStatus): FirNamedFunctionSymbol {
-        return functionsCache.getOrPut(symbol) {
+        val oldFunction = symbol.fir
-            val oldFunction = symbol.fir
+        val newSymbol = FirNamedFunctionSymbol(CallableId(firKotlinClass.classId, symbol.callableId.callableName))
-            val newSymbol = FirNamedFunctionSymbol(CallableId(firKotlinClass.classId, symbol.callableId.callableName))
+        FirFakeOverrideGenerator.createCopyForFirFunction(
-            FirFakeOverrideGenerator.createCopyForFirFunction(
+            newSymbol,
-                newSymbol,
+            baseFunction = symbol.fir,
-                baseFunction = symbol.fir,
+            derivedClassLookupTag = firKotlinClass.symbol.toLookupTag(),
-                derivedClassLookupTag = firKotlinClass.symbol.toLookupTag(),
+            session,
-                session,
+            symbol.fir.origin,
-                symbol.fir.origin,
+            newDispatchReceiverType = kotlinDispatchReceiverType,
-                newDispatchReceiverType = kotlinDispatchReceiverType,
+            newParameterTypes = oldFunction.valueParameters.map { substitutor.substituteOrSelf(it.returnTypeRef.coneType) },
-                newParameterTypes = oldFunction.valueParameters.map { substitutor.substituteOrSelf(it.returnTypeRef.coneType) },
+            newReturnType = substitutor.substituteOrSelf(oldFunction.returnTypeRef.coneType),
-                newReturnType = substitutor.substituteOrSelf(oldFunction.returnTypeRef.coneType),
+        ).apply {
-            ).apply {
+            if (jdkMemberStatus == JDKMemberStatus.HIDDEN) {
-                if (jdkMemberStatus == JDKMemberStatus.HIDDEN) {
+                isHiddenEverywhereBesideSuperCalls = true
-                    isHiddenEverywhereBesideSuperCalls = true
-                }
             }
-            newSymbol
         }
+        return newSymbol
     }
 
     override fun processDirectOverriddenFunctionsWithBaseScope(
@@ -275,7 +269,7 @@ class JvmMappedScope(
             if (javaCtor.isTrivialCopyConstructor()) return@processor
             if (firKotlinClassConstructors.any { javaCtor.isShadowedBy(it) }) return@processor
 
-            val newSymbol = getOrCreateCopy(javaCtorSymbol)
+            val newSymbol = mappedSymbolCache.mappedConstructors.getValue(javaCtorSymbol, this)
             processor(newSymbol)
         }
 
@@ -284,27 +278,25 @@ class JvmMappedScope(
 
     private fun FirDeclaration.isDeprecated(): Boolean = symbol.getDeprecation(session, callSite = null) != null
 
-    private fun getOrCreateCopy(symbol: FirConstructorSymbol): FirConstructorSymbol {
+    private fun createMappedConstructor(symbol: FirConstructorSymbol): FirConstructorSymbol {
-        return constructorsCache.getOrPut(symbol) {
+        val oldConstructor = symbol.fir
-            val oldConstructor = symbol.fir
+        val classId = firKotlinClass.classId
-            val classId = firKotlinClass.classId
+        val newSymbol = FirConstructorSymbol(CallableId(classId, classId.shortClassName))
-            val newSymbol = FirConstructorSymbol(CallableId(classId, classId.shortClassName))
+        FirFakeOverrideGenerator.createCopyForFirConstructor(
-            FirFakeOverrideGenerator.createCopyForFirConstructor(
+            newSymbol,
-                newSymbol,
+            session,
-                session,
+            oldConstructor,
-                oldConstructor,
+            derivedClassLookupTag = firKotlinClass.symbol.toLookupTag(),
-                derivedClassLookupTag = firKotlinClass.symbol.toLookupTag(),
+            symbol.fir.origin,
-                symbol.fir.origin,
+            newDispatchReceiverType = null,
-                newDispatchReceiverType = null,
+            newReturnType = substitutor.substituteOrSelf(oldConstructor.returnTypeRef.coneType),
-                newReturnType = substitutor.substituteOrSelf(oldConstructor.returnTypeRef.coneType),
+            newParameterTypes = oldConstructor.valueParameters.map { substitutor.substituteOrSelf(it.returnTypeRef.coneType) },
-                newParameterTypes = oldConstructor.valueParameters.map { substitutor.substituteOrSelf(it.returnTypeRef.coneType) },
+            newTypeParameters = null,
-                newTypeParameters = null,
+            newContextReceiverTypes = emptyList(),
-                newContextReceiverTypes = emptyList(),
+            isExpect = false,
-                isExpect = false,
+            fakeOverrideSubstitution = null
-                fakeOverrideSubstitution = null
+        )
-            )
+        return newSymbol
-            newSymbol
-        }
     }
 
     override fun processDirectOverriddenPropertiesWithBaseScope(
@@ -322,6 +314,26 @@ class JvmMappedScope(
         return declaredMemberScope.getClassifierNames()
     }
 
+    class FirMappedSymbolStorage(private val cachesFactory: FirCachesFactory) : FirSessionComponent {
+        constructor(session: FirSession) : this(session.firCachesFactory)
+
+        // Key is the kotlin class
+        internal val cacheByOwner: FirCache<FirRegularClassSymbol, MappedSymbolsCache, Nothing?> =
+            cachesFactory.createCache { _ -> MappedSymbolsCache(cachesFactory) }
+
+        internal class MappedSymbolsCache(cachesFactory: FirCachesFactory) {
+            val mappedFunctions: FirCache<FirNamedFunctionSymbol, FirNamedFunctionSymbol, Pair<JvmMappedScope, JDKMemberStatus>> =
+                cachesFactory.createCache { symbol, (scope, jdkMemberStatus) ->
+                    scope.createMappedFunction(symbol, jdkMemberStatus)
+                }
+
+            val mappedConstructors: FirCache<FirConstructorSymbol, FirConstructorSymbol, JvmMappedScope> =
+                cachesFactory.createCache { symbol, scope ->
+                    scope.createMappedConstructor(symbol)
+                }
+        }
+    }
+
     companion object {
         /**
          * For fromClass=A<T1, T2>, toClass=B<F1, F1> classes
@@ -343,3 +355,5 @@ class JvmMappedScope(
         return "JVM mapped scope for ${firKotlinClass.classId}"
     }
 }
+
+private val FirSession.mappedSymbolStorage: JvmMappedScope.FirMappedSymbolStorage by FirSession.sessionComponentAccessor()

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/FirFrontendFacade.kt

@@ -25,7 +25,6 @@ import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
 import org.jetbrains.kotlin.fir.deserialization.ModuleDataProvider
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
-import org.jetbrains.kotlin.fir.java.enhancement.FirEnhancedSymbolsStorage
 import org.jetbrains.kotlin.fir.session.*
 import org.jetbrains.kotlin.fir.session.environment.AbstractProjectEnvironment
 import org.jetbrains.kotlin.js.config.JSConfigurationKeys
@@ -97,7 +96,9 @@ open class FirFrontendFacade(
         val project = testServices.compilerConfigurationProvider.getProject(module)
         val extensionRegistrars = FirExtensionRegistrar.getInstances(project)
         val targetPlatform = module.targetPlatform
-        val enhancedSymbolsStorage = runIf(targetPlatform.isJvm()) { FirEnhancedSymbolsStorage(firCachesFactoryForCliMode) }
+        val predefinedJavaComponents = runIf(targetPlatform.isJvm()) {
+            FirSharableJavaComponents(firCachesFactoryForCliMode)
+        }
         val projectEnvironment = createLibrarySession(
             module,
             project,
@@ -106,11 +107,11 @@ open class FirFrontendFacade(
             moduleDataProvider,
             testServices.compilerConfigurationProvider.getCompilerConfiguration(module),
             extensionRegistrars,
-            enhancedSymbolsStorage
+            predefinedJavaComponents
         )
 
         val firOutputPartForDependsOnModules = sortedModules.map {
-            analyze(it, moduleDataMap[it]!!, targetPlatform, projectEnvironment, extensionRegistrars, enhancedSymbolsStorage)
+            analyze(it, moduleDataMap[it]!!, targetPlatform, projectEnvironment, extensionRegistrars, predefinedJavaComponents)
         }
 
         return FirOutputArtifactImpl(firOutputPartForDependsOnModules)
@@ -172,7 +173,7 @@ open class FirFrontendFacade(
         moduleDataProvider: ModuleDataProvider,
         configuration: CompilerConfiguration,
         extensionRegistrars: List<FirExtensionRegistrar>,
-        enhancedSymbolsStorage: FirEnhancedSymbolsStorage?
+        predefinedJavaComponents: FirSharableJavaComponents?
     ): AbstractProjectEnvironment? {
         val compilerConfigurationProvider = testServices.compilerConfigurationProvider
         val projectEnvironment: AbstractProjectEnvironment?
@@ -210,7 +211,7 @@ open class FirFrontendFacade(
                         projectFileSearchScope,
                         packagePartProvider,
                         languageVersionSettings,
-                        predefinedEnhancementStorage = enhancedSymbolsStorage,
+                        predefinedJavaComponents = predefinedJavaComponents,
                         registerExtraComponents = ::registerExtraComponents,
                     )
                 }
@@ -267,7 +268,7 @@ open class FirFrontendFacade(
         targetPlatform: TargetPlatform,
         projectEnvironment: AbstractProjectEnvironment?,
         extensionRegistrars: List<FirExtensionRegistrar>,
-        enhancedSymbolsStorage: FirEnhancedSymbolsStorage?,
+        predefinedJavaComponents: FirSharableJavaComponents?,
     ): FirOutputPartForDependsOnModule {
         val compilerConfigurationProvider = testServices.compilerConfigurationProvider
         val moduleInfoProvider = testServices.firModuleInfoProvider
@@ -301,7 +302,7 @@ open class FirFrontendFacade(
             projectEnvironment,
             extensionRegistrars,
             sessionConfigurator,
-            enhancedSymbolsStorage,
+            predefinedJavaComponents,
             project,
             ktFiles
         )
@@ -330,7 +331,7 @@ open class FirFrontendFacade(
         projectEnvironment: AbstractProjectEnvironment?,
         extensionRegistrars: List<FirExtensionRegistrar>,
         sessionConfigurator: FirSessionConfigurator.() -> Unit,
-        enhancedSymbolsStorage: FirEnhancedSymbolsStorage?,
+        predefinedJavaComponents: FirSharableJavaComponents?,
         project: Project,
         ktFiles: Collection<KtFile>,
     ): FirSession {
@@ -357,7 +358,7 @@ open class FirFrontendFacade(
                     incrementalCompilationContext = null,
                     extensionRegistrars,
                     languageVersionSettings,
-                    predefinedEnhancementStorage = enhancedSymbolsStorage,
+                    predefinedJavaComponents = predefinedJavaComponents,
                     needRegisterJavaElementFinder = true,
                     registerExtraComponents = ::registerExtraComponents,
                     init = sessionConfigurator,