[FIR] Register type related extensions in libraries sessions

^KT-57140

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt

@@ -35,6 +35,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirFunctionTypeKindExtension::class,
             FirDeclarationsForMetadataProviderExtension::class,
         )
+
+        internal val ALLOWED_EXTENSIONS_FOR_LIBRARY_SESSION = listOf(
+            FirTypeAttributeExtension::class,
+            FirFunctionTypeKindExtension::class,
+        )
     }
 
     protected abstract fun ExtensionRegistrarContext.configurePlugin()

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

@@ -34,6 +34,7 @@ abstract class FirAbstractSessionFactory {
         sessionProvider: FirProjectSessionProvider,
         moduleDataProvider: ModuleDataProvider,
         languageVersionSettings: LanguageVersionSettings,
+        extensionRegistrars: List<FirExtensionRegistrar>,
         registerExtraComponents: ((FirSession) -> Unit),
         createKotlinScopeProvider: () -> FirKotlinScopeProvider,
         createProviders: (FirSession, FirModuleData, FirKotlinScopeProvider) -> List<FirSymbolProvider>
@@ -59,6 +60,12 @@ abstract class FirAbstractSessionFactory {
             )
             builtinsModuleData.bindSession(this)
 
+            FirSessionConfigurator(this).apply {
+                for (extensionRegistrar in extensionRegistrars) {
+                    registerExtensions(extensionRegistrar.configure())
+                }
+            }.configure()
+
             val providers = createProviders(this, builtinsModuleData, kotlinScopeProvider)
 
             val symbolProvider = FirCachingCompositeSymbolProvider(this, providers)

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

@@ -37,6 +37,7 @@ object FirCommonSessionFactory : FirAbstractSessionFactory() {
         sessionProvider: FirProjectSessionProvider,
         moduleDataProvider: ModuleDataProvider,
         projectEnvironment: AbstractProjectEnvironment,
+        extensionRegistrars: List<FirExtensionRegistrar>,
         librariesScope: AbstractProjectFileSearchScope,
         resolvedKLibs: List<KotlinResolvedLibrary>,
         packageAndMetadataPartProvider: PackageAndMetadataPartProvider,
@@ -48,6 +49,7 @@ object FirCommonSessionFactory : FirAbstractSessionFactory() {
             sessionProvider,
             moduleDataProvider,
             languageVersionSettings,
+            extensionRegistrars,
             registerExtraComponents = {
                 registerExtraComponents(it)
             },

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

@@ -66,13 +66,15 @@ object FirJsSessionFactory : FirAbstractSessionFactory() {
         resolvedLibraries: List<KotlinLibrary>,
         sessionProvider: FirProjectSessionProvider,
         moduleDataProvider: ModuleDataProvider,
+        extensionRegistrars: List<FirExtensionRegistrar>,
         languageVersionSettings: LanguageVersionSettings = LanguageVersionSettingsImpl.DEFAULT,
         registerExtraComponents: ((FirSession) -> Unit),
-    ) = createLibrarySession(
+    ): FirSession = createLibrarySession(
         mainModuleName,
         sessionProvider,
         moduleDataProvider,
         languageVersionSettings,
+        extensionRegistrars,
         registerExtraComponents = {
             it.registerJsSpecificResolveComponents()
             registerExtraComponents(it)
@@ -83,7 +85,7 @@ object FirJsSessionFactory : FirAbstractSessionFactory() {
                 KlibBasedSymbolProvider(session, moduleDataProvider, kotlinScopeProvider, resolvedLibraries),
                 // (Most) builtins should be taken from the dependencies in JS compilation, therefore builtins provider is the last one
                 // TODO: consider using "poisoning" provider for builtins to ensure that proper ones are taken from dependencies
-                // NOTE: it requires precise filtering for true nuiltins, like Function*
+                // NOTE: it requires precise filtering for true builtins, like Function*
                 FirBuiltinSymbolProvider(session, builtinsModuleData, kotlinScopeProvider),
                 FirExtensionSyntheticFunctionInterfaceProvider.createIfNeeded(session, builtinsModuleData, kotlinScopeProvider),
             )

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

@@ -37,6 +37,7 @@ object FirJvmSessionFactory : FirAbstractSessionFactory() {
         sessionProvider: FirProjectSessionProvider,
         moduleDataProvider: ModuleDataProvider,
         projectEnvironment: AbstractProjectEnvironment,
+        extensionRegistrars: List<FirExtensionRegistrar>,
         scope: AbstractProjectFileSearchScope,
         packagePartProvider: PackagePartProvider,
         languageVersionSettings: LanguageVersionSettings,
@@ -47,6 +48,7 @@ object FirJvmSessionFactory : FirAbstractSessionFactory() {
             sessionProvider,
             moduleDataProvider,
             languageVersionSettings,
+            extensionRegistrars,
             registerExtraComponents = {
                 it.registerCommonJavaComponents(projectEnvironment.getJavaModuleResolver())
                 registerExtraComponents(it)

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

@@ -27,6 +27,7 @@ object FirNativeSessionFactory : FirAbstractSessionFactory() {
         resolvedLibraries: List<KotlinResolvedLibrary>,
         sessionProvider: FirProjectSessionProvider,
         moduleDataProvider: ModuleDataProvider,
+        extensionRegistrars: List<FirExtensionRegistrar>,
         languageVersionSettings: LanguageVersionSettings,
         registerExtraComponents: ((FirSession) -> Unit) = {},
     ): FirSession {
@@ -35,6 +36,7 @@ object FirNativeSessionFactory : FirAbstractSessionFactory() {
             sessionProvider,
             moduleDataProvider,
             languageVersionSettings,
+            extensionRegistrars,
             registerExtraComponents,
             createKotlinScopeProvider = { FirKotlinScopeProvider { _, declaredMemberScope, _, _, _ -> declaredMemberScope } },
             createProviders = { session, builtinsModuleData, kotlinScopeProvider ->

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

@@ -12,9 +12,7 @@ import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.type.TypeCheckers
 import org.jetbrains.kotlin.fir.analysis.checkersComponent
 import org.jetbrains.kotlin.fir.analysis.extensions.additionalCheckers
-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.extensions.*
 
 class FirSessionConfigurator(private val session: FirSession) {
     private val registeredExtensions: MutableList<BunchOfRegisteredExtensions> = mutableListOf(BunchOfRegisteredExtensions.empty())
@@ -38,9 +36,17 @@ class FirSessionConfigurator(private val session: FirSession) {
         session.checkersComponent.register(checkers)
     }
 
+    @OptIn(PluginServicesInitialization::class)
     @SessionConfiguration
     fun configure() {
-        session.extensionService.registerExtensions(registeredExtensions.reduce(BunchOfRegisteredExtensions::plus))
-        session.extensionService.additionalCheckers.forEach(session.checkersComponent::register)
+        var extensions = registeredExtensions.reduce(BunchOfRegisteredExtensions::plus)
+        if (session.kind == FirSession.Kind.Library) {
+            val filteredExtensions = extensions.extensions.filterKeys { it in FirExtensionRegistrar.ALLOWED_EXTENSIONS_FOR_LIBRARY_SESSION }
+            extensions = BunchOfRegisteredExtensions(filteredExtensions)
+        }
+        session.extensionService.registerExtensions(extensions)
+        if (session.kind == FirSession.Kind.Source) {
+            session.extensionService.additionalCheckers.forEach(session.checkersComponent::register)
+        }
     }
-}
+}

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

@@ -51,6 +51,7 @@ object FirSessionFactoryHelper {
             sessionProvider,
             dependencyList.moduleDataProvider,
             projectEnvironment,
+            extensionRegistrars,
             librariesScope,
             packagePartProvider,
             languageVersionSettings,
@@ -129,4 +130,4 @@ object FirSessionFactoryHelper {
         register(FirOverridesBackwardCompatibilityHelper::class, FirOverridesBackwardCompatibilityHelper.Default())
         register(FirEnumEntriesSupport::class, FirEnumEntriesSupport(this))
     }
-}
+}