[FIR] Share enhancedSymbolStorage between common and platform sessions

This is needed for two reasons:
1. common and platform modules are analyzed in the same setup with same
   dependencies and configurations, so the results of the enhancement for
   any function will be completely identical. So by sharing enhanced
   symbol storage, we avoid recomputation enhanced functions and improve
   performance
2. There is a goal to have some unique key for IR declarations in FIR2IR.
   For regular declarations, the key is just Fir symbol, and for fake-overrides
   it is a pair of original symbol and fake override owners lookup tag.
   But for enhanced functions the symbol is not unique, because we create
   different symbols for the same enhanced function for different modules.
   So this change fixes this problem

^KT-60397 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 are some test breaks in branch for KT-61637:
- FirPsiBlackBoxCodegenTestGenerated.Multiplatform.testStarImportOfExpectEnumWithActualTypeAlias
- FirPsiBlackBoxCodegenTestGenerated.Multiplatform.K2.DefaultArguments.testNestedEnumEntryValue
- FirPsiBlackBoxCodegenTestGenerated.Multiplatform.K2.Basic.testEnumEntryNameCall

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

@@ -16,6 +16,7 @@ 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
@@ -72,7 +73,7 @@ fun <F> prepareJvmSessions(
     createProviderAndScopeForIncrementalCompilation: (List<F>) -> IncrementalCompilationContext?,
 ): List<SessionWithSources<F>> {
     val javaSourcesScope = projectEnvironment.getSearchScopeForProjectJavaSources()
-
+    val enhancementSymbolsStorage = FirEnhancedSymbolsStorage(firCachesFactoryForCliMode)
     return prepareSessions(
         files, configuration, rootModuleName, JvmPlatforms.unspecifiedJvmPlatform,
         JvmPlatformAnalyzerServices, metadataCompilationMode = false, libraryList, isCommonSource, fileBelongsToModule,
@@ -86,6 +87,7 @@ fun <F> prepareJvmSessions(
                 librariesScope,
                 projectEnvironment.getPackagePartProvider(librariesScope),
                 configuration.languageVersionSettings,
+                predefinedEnhancementStorage = enhancementSymbolsStorage,
                 registerExtraComponents = {},
             )
         }
@@ -101,6 +103,7 @@ fun <F> prepareJvmSessions(
             configuration.get(CommonConfigurationKeys.LOOKUP_TRACKER),
             configuration.get(CommonConfigurationKeys.ENUM_WHEN_TRACKER),
             configuration.get(CommonConfigurationKeys.IMPORT_TRACKER),
+            predefinedEnhancementStorage = enhancementSymbolsStorage,
             needRegisterJavaElementFinder = true,
             registerExtraComponents = {},
             sessionConfigurator,

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

@@ -90,9 +90,12 @@ fun FirSession.registerCommonComponentsAfterExtensionsAreConfigured() {
     register(FirProvidedDeclarationsForMetadataService::class, FirProvidedDeclarationsForMetadataService.create(this))
 }
 
+val firCachesFactoryForCliMode: FirCachesFactory
+    get() = FirThreadUnsafeCachesFactory
+
 @OptIn(SessionConfiguration::class)
 fun FirSession.registerCliCompilerOnlyComponents() {
-    register(FirCachesFactory::class, FirThreadUnsafeCachesFactory)
+    register(FirCachesFactory::class, firCachesFactoryForCliMode)
     register(SealedClassInheritorsProvider::class, SealedClassInheritorsProviderImpl)
     register(FirLazyDeclarationResolver::class, FirDummyCompilerLazyDeclarationResolver)
     register(FirExceptionHandler::class, FirCliExceptionHandler)
@@ -102,10 +105,13 @@ fun FirSession.registerCliCompilerOnlyComponents() {
 }
 
 @OptIn(SessionConfiguration::class)
-fun FirSession.registerCommonJavaComponents(javaModuleResolver: JavaModuleResolver) {
+fun FirSession.registerCommonJavaComponents(
+    javaModuleResolver: JavaModuleResolver,
+    predefinedEnhancementStorage: FirEnhancedSymbolsStorage? = null,
+) {
     val jsr305State = languageVersionSettings.getFlag(JvmAnalysisFlags.javaTypeEnhancementState)
     register(FirAnnotationTypeQualifierResolver::class, FirAnnotationTypeQualifierResolver(this, jsr305State, javaModuleResolver))
-    register(FirEnhancedSymbolsStorage::class, FirEnhancedSymbolsStorage(this))
+    register(FirEnhancedSymbolsStorage::class, predefinedEnhancementStorage ?: FirEnhancedSymbolsStorage(this))
     register(FirSyntheticPropertiesStorage::class, FirSyntheticPropertiesStorage(this))
     register(
         FirJvmDefaultModeComponent::class,

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

@@ -19,9 +19,9 @@ 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.providers.impl.FirExtensionSyntheticFunctionInterfaceProvider
 import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
 import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
 import org.jetbrains.kotlin.fir.session.environment.AbstractProjectEnvironment
@@ -42,6 +42,7 @@ object FirJvmSessionFactory : FirAbstractSessionFactory() {
         scope: AbstractProjectFileSearchScope,
         packagePartProvider: PackagePartProvider,
         languageVersionSettings: LanguageVersionSettings,
+        predefinedEnhancementStorage: FirEnhancedSymbolsStorage? = null,
         registerExtraComponents: ((FirSession) -> Unit),
     ): FirSession {
         return createLibrarySession(
@@ -51,7 +52,7 @@ object FirJvmSessionFactory : FirAbstractSessionFactory() {
             languageVersionSettings,
             extensionRegistrars,
             registerExtraComponents = {
-                it.registerCommonJavaComponents(projectEnvironment.getJavaModuleResolver())
+                it.registerCommonJavaComponents(projectEnvironment.getJavaModuleResolver(), predefinedEnhancementStorage)
                 registerExtraComponents(it)
             },
             createKotlinScopeProvider = { FirKotlinScopeProvider(::wrapScopeWithJvmMapped) },
@@ -91,6 +92,7 @@ object FirJvmSessionFactory : FirAbstractSessionFactory() {
         lookupTracker: LookupTracker? = null,
         enumWhenTracker: EnumWhenTracker? = null,
         importTracker: ImportTracker? = null,
+        predefinedEnhancementStorage: FirEnhancedSymbolsStorage? = null,
         needRegisterJavaElementFinder: Boolean,
         registerExtraComponents: ((FirSession) -> Unit) = {},
         init: FirSessionConfigurator.() -> Unit = {}
@@ -105,7 +107,7 @@ object FirJvmSessionFactory : FirAbstractSessionFactory() {
             importTracker,
             init,
             registerExtraComponents = {
-                it.registerCommonJavaComponents(projectEnvironment.getJavaModuleResolver())
+                it.registerCommonJavaComponents(projectEnvironment.getJavaModuleResolver(), predefinedEnhancementStorage)
                 it.registerJavaSpecificResolveComponents()
                 registerExtraComponents(it)
             },

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

@@ -76,9 +76,9 @@ object FirSessionFactoryHelper {
             lookupTracker,
             enumWhenTracker,
             importTracker,
-            needRegisterJavaElementFinder,
+            needRegisterJavaElementFinder = needRegisterJavaElementFinder,
             registerExtraComponents = {},
-            sessionConfigurator,
+            init = sessionConfigurator,
         )
     }
 

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/enhancement/SignatureEnhancement.kt

@@ -747,8 +747,8 @@ private class EnhancementSignatureParts(
         get() = (this as ConeTypeParameterLookupTag).symbol.fir.origin is FirDeclarationOrigin.Java
 }
 
-class FirEnhancedSymbolsStorage(val session: FirSession) : FirSessionComponent {
+class FirEnhancedSymbolsStorage(private val cachesFactory: FirCachesFactory) : FirSessionComponent {
-    private val cachesFactory = session.firCachesFactory
+    constructor(session: FirSession) : this(session.firCachesFactory)
 
     val cacheByOwner: FirCache<FirRegularClassSymbol, EnhancementSymbolsCache, Nothing?> =
         cachesFactory.createCache { _ -> EnhancementSymbolsCache(cachesFactory) }

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

@@ -25,10 +25,8 @@ 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.session.FirCommonSessionFactory
+import org.jetbrains.kotlin.fir.java.enhancement.FirEnhancedSymbolsStorage
-import org.jetbrains.kotlin.fir.session.FirJvmSessionFactory
+import org.jetbrains.kotlin.fir.session.*
-import org.jetbrains.kotlin.fir.session.FirNativeSessionFactory
-import org.jetbrains.kotlin.fir.session.FirSessionConfigurator
 import org.jetbrains.kotlin.fir.session.environment.AbstractProjectEnvironment
 import org.jetbrains.kotlin.js.config.JSConfigurationKeys
 import org.jetbrains.kotlin.js.config.WasmTarget
@@ -54,6 +52,7 @@ import org.jetbrains.kotlin.test.services.*
 import org.jetbrains.kotlin.test.services.configuration.JsEnvironmentConfigurator
 import org.jetbrains.kotlin.test.services.configuration.NativeEnvironmentConfigurator
 import org.jetbrains.kotlin.test.services.configuration.WasmEnvironmentConfigurator
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import java.nio.file.Paths
 
 open class FirFrontendFacade(
@@ -97,6 +96,8 @@ 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 projectEnvironment = createLibrarySession(
             module,
             project,
@@ -104,12 +105,12 @@ open class FirFrontendFacade(
             testServices.firModuleInfoProvider.firSessionProvider,
             moduleDataProvider,
             testServices.compilerConfigurationProvider.getCompilerConfiguration(module),
-            extensionRegistrars
+            extensionRegistrars,
+            enhancedSymbolsStorage
         )
 
-        val targetPlatform = module.targetPlatform
         val firOutputPartForDependsOnModules = sortedModules.map {
-            analyze(it, moduleDataMap[it]!!, targetPlatform, projectEnvironment, extensionRegistrars)
+            analyze(it, moduleDataMap[it]!!, targetPlatform, projectEnvironment, extensionRegistrars, enhancedSymbolsStorage)
         }
 
         return FirOutputArtifactImpl(firOutputPartForDependsOnModules)
@@ -170,7 +171,8 @@ open class FirFrontendFacade(
         sessionProvider: FirProjectSessionProvider,
         moduleDataProvider: ModuleDataProvider,
         configuration: CompilerConfiguration,
-        extensionRegistrars: List<FirExtensionRegistrar>
+        extensionRegistrars: List<FirExtensionRegistrar>,
+        enhancedSymbolsStorage: FirEnhancedSymbolsStorage?
     ): AbstractProjectEnvironment? {
         val compilerConfigurationProvider = testServices.compilerConfigurationProvider
         val projectEnvironment: AbstractProjectEnvironment?
@@ -208,6 +210,7 @@ open class FirFrontendFacade(
                         projectFileSearchScope,
                         packagePartProvider,
                         languageVersionSettings,
+                        predefinedEnhancementStorage = enhancedSymbolsStorage,
                         registerExtraComponents = ::registerExtraComponents,
                     )
                 }
@@ -264,6 +267,7 @@ open class FirFrontendFacade(
         targetPlatform: TargetPlatform,
         projectEnvironment: AbstractProjectEnvironment?,
         extensionRegistrars: List<FirExtensionRegistrar>,
+        enhancedSymbolsStorage: FirEnhancedSymbolsStorage?,
     ): FirOutputPartForDependsOnModule {
         val compilerConfigurationProvider = testServices.compilerConfigurationProvider
         val moduleInfoProvider = testServices.firModuleInfoProvider
@@ -297,6 +301,7 @@ open class FirFrontendFacade(
             projectEnvironment,
             extensionRegistrars,
             sessionConfigurator,
+            enhancedSymbolsStorage,
             project,
             ktFiles
         )
@@ -325,8 +330,9 @@ open class FirFrontendFacade(
         projectEnvironment: AbstractProjectEnvironment?,
         extensionRegistrars: List<FirExtensionRegistrar>,
         sessionConfigurator: FirSessionConfigurator.() -> Unit,
+        enhancedSymbolsStorage: FirEnhancedSymbolsStorage?,
         project: Project,
-        ktFiles: Collection<KtFile>
+        ktFiles: Collection<KtFile>,
     ): FirSession {
         val languageVersionSettings = module.languageVersionSettings
         return when {
@@ -351,6 +357,7 @@ open class FirFrontendFacade(
                     incrementalCompilationContext = null,
                     extensionRegistrars,
                     languageVersionSettings,
+                    predefinedEnhancementStorage = enhancedSymbolsStorage,
                     needRegisterJavaElementFinder = true,
                     registerExtraComponents = ::registerExtraComponents,
                     init = sessionConfigurator,