[FIR] Register type related extensions in libraries sessions

^KT-57140

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

@@ -75,12 +75,12 @@ fun <F> prepareJvmSessions(
                 sessionProvider,
                 libraryList.moduleDataProvider,
                 projectEnvironment,
+                extensionRegistrars,
                 librariesScope,
                 projectEnvironment.getPackagePartProvider(librariesScope),
                 configuration.languageVersionSettings,
                 registerExtraComponents = {},
             )
-
         }
     ) { moduleFiles, moduleData, sessionProvider, sessionConfigurator ->
         FirJvmSessionFactory.createModuleBasedSession(
@@ -126,6 +126,7 @@ fun <F> prepareJsSessions(
                 resolvedLibraries,
                 sessionProvider,
                 libraryList.moduleDataProvider,
+                extensionRegistrars,
                 configuration.languageVersionSettings,
                 registerExtraComponents = {},
             )
@@ -168,6 +169,7 @@ fun <F> prepareNativeSessions(
                 resolvedLibraries,
                 sessionProvider,
                 libraryList.moduleDataProvider,
+                extensionRegistrars,
                 configuration.languageVersionSettings,
                 registerExtraComponents = {},
             )
@@ -211,6 +213,7 @@ fun <F> prepareCommonSessions(
                 sessionProvider,
                 libraryList.moduleDataProvider,
                 projectEnvironment,
+                extensionRegistrars,
                 librariesScope,
                 resolvedLibraries,
                 projectEnvironment.getPackagePartProvider(librariesScope) as PackageAndMetadataPartProvider,

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.*
-import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.registerExtensions
 
 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))
+        var extensions = registeredExtensions.reduce(BunchOfRegisteredExtensions::plus)
-        session.extensionService.additionalCheckers.forEach(session.checkersComponent::register)
+        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))
     }
-}
+}

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

@@ -88,18 +88,21 @@ open class FirFrontendFacade(
 
         val (moduleDataMap, moduleDataProvider) = initializeModuleData(sortedModules)
 
+        val project = testServices.compilerConfigurationProvider.getProject(module)
+        val extensionRegistrars = FirExtensionRegistrar.getInstances(project)
         val projectEnvironment = createLibrarySession(
             module,
-            testServices.compilerConfigurationProvider.getProject(module),
+            project,
             Name.special("<${module.name}>"),
             testServices.firModuleInfoProvider.firSessionProvider,
             moduleDataProvider,
-            testServices.compilerConfigurationProvider.getCompilerConfiguration(module)
+            testServices.compilerConfigurationProvider.getCompilerConfiguration(module),
+            extensionRegistrars
         )
 
         val targetPlatform = module.targetPlatform
         val firOutputPartForDependsOnModules = sortedModules.map {
-            analyze(it, moduleDataMap[it]!!, targetPlatform, projectEnvironment)
+            analyze(it, moduleDataMap[it]!!, targetPlatform, projectEnvironment, extensionRegistrars)
         }
 
         return FirOutputArtifactImpl(firOutputPartForDependsOnModules)
@@ -182,7 +185,8 @@ open class FirFrontendFacade(
         moduleName: Name,
         sessionProvider: FirProjectSessionProvider,
         moduleDataProvider: ModuleDataProvider,
-        configuration: CompilerConfiguration
+        configuration: CompilerConfiguration,
+        extensionRegistrars: List<FirExtensionRegistrar>
     ): AbstractProjectEnvironment? {
         val compilerConfigurationProvider = testServices.compilerConfigurationProvider
         val projectEnvironment: AbstractProjectEnvironment?
@@ -202,6 +206,7 @@ open class FirFrontendFacade(
                     sessionProvider,
                     moduleDataProvider,
                     projectEnvironment,
+                    extensionRegistrars,
                     projectFileSearchScope,
                     packagePartProvider,
                     languageVersionSettings,
@@ -217,6 +222,7 @@ open class FirFrontendFacade(
                     module,
                     testServices,
                     configuration,
+                    extensionRegistrars,
                     languageVersionSettings,
                     registerExtraComponents = ::registerExtraComponents,
                 )
@@ -228,6 +234,7 @@ open class FirFrontendFacade(
                     listOf(),
                     sessionProvider,
                     moduleDataProvider,
+                    extensionRegistrars,
                     languageVersionSettings,
                     registerExtraComponents = ::registerExtraComponents,
                 )
@@ -241,7 +248,8 @@ open class FirFrontendFacade(
         module: TestModule,
         moduleData: FirModuleData,
         targetPlatform: TargetPlatform,
-        projectEnvironment: AbstractProjectEnvironment?
+        projectEnvironment: AbstractProjectEnvironment?,
+        extensionRegistrars: List<FirExtensionRegistrar>,
     ): FirOutputPartForDependsOnModule {
         val compilerConfigurationProvider = testServices.compilerConfigurationProvider
         val moduleInfoProvider = testServices.firModuleInfoProvider
@@ -258,7 +266,6 @@ open class FirFrontendFacade(
             FirParser.Psi -> testServices.sourceFileProvider.getKtFilesForSourceFiles(module.files, project).values to emptyList()
         }
 
-        val extensionRegistrars = FirExtensionRegistrar.getInstances(project)
         val sessionConfigurator: FirSessionConfigurator.() -> Unit = {
             if (FirDiagnosticsDirectives.WITH_EXTENDED_CHECKERS in module.directives) {
                 registerExtendedCommonCheckers()

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

@@ -33,6 +33,7 @@ object TestFirJsSessionFactory {
         module: TestModule,
         testServices: TestServices,
         configuration: CompilerConfiguration,
+        extensionRegistrars: List<FirExtensionRegistrar>,
         languageVersionSettings: LanguageVersionSettings,
         registerExtraComponents: ((FirSession) -> Unit),
     ): FirSession {
@@ -45,6 +46,7 @@ object TestFirJsSessionFactory {
             resolvedLibraries.map { it.library },
             sessionProvider,
             moduleDataProvider,
+            extensionRegistrars,
             languageVersionSettings,
             registerExtraComponents,
         )

plugins/fir-plugin-prototype/testData/diagnostics/checkers/dependencyWithoutAttributePlugin.fir.txt

@@ -11,13 +11,13 @@ FILE: dependencyWithoutAttributePlugin.kt
         R|org/jetbrains/kotlin/fir/plugin/consumePositiveInt|(R|<local>/someInt|)
     }
     public final fun test_2(): R|kotlin/Unit| {
-        lval x: R|@R|org/jetbrains/kotlin/fir/plugin/Positive|()  kotlin/Int| = R|org/jetbrains/kotlin/fir/plugin/producePositiveInt|()
+        lval x: R|@Positive kotlin/Int| = R|org/jetbrains/kotlin/fir/plugin/producePositiveInt|()
         R|/takePositive|(R|<local>/x|)
         R|/takeNegative|(R|<local>/x|)
         R|/takeAny|(R|<local>/x|)
     }
     public final fun test_3(): R|kotlin/Unit| {
-        lval x: R|@R|org/jetbrains/kotlin/fir/plugin/Positive|()  kotlin/Int| = R|org/jetbrains/kotlin/fir/plugin/produceBoxedPositiveInt|().R|SubstitutionOverride<org/jetbrains/kotlin/fir/plugin/Box.value: R|@R|org/jetbrains/kotlin/fir/plugin/Positive|()  kotlin/Int|>|
+        lval x: R|@Positive kotlin/Int| = R|org/jetbrains/kotlin/fir/plugin/produceBoxedPositiveInt|().R|SubstitutionOverride<org/jetbrains/kotlin/fir/plugin/Box.value: R|@Positive kotlin/Int|>|
         R|/takePositive|(R|<local>/x|)
         R|/takeNegative|(R|<local>/x|)
         R|/takeAny|(R|<local>/x|)

plugins/fir-plugin-prototype/testData/diagnostics/checkers/dependencyWithoutAttributePlugin.kt

@@ -10,20 +10,20 @@ fun test_1(
     someInt: Int
 ) {
     consumePositiveInt(positiveInt)
-    consumePositiveInt(negativeInt) // should be error
+    consumePositiveInt(<!ILLEGAL_NUMBER_SIGN!>negativeInt<!>) // should be error
-    consumePositiveInt(someInt) // should be error
+    consumePositiveInt(<!ILLEGAL_NUMBER_SIGN!>someInt<!>) // should be error
 }
 
 fun test_2() {
     val x = producePositiveInt()
-    takePositive(<!ILLEGAL_NUMBER_SIGN!>x<!>)
+    takePositive(x)
     takeNegative(<!ILLEGAL_NUMBER_SIGN!>x<!>) // should be error
     takeAny(x)
 }
 
 fun test_3() {
     val x = produceBoxedPositiveInt().value
-    takePositive(<!ILLEGAL_NUMBER_SIGN!>x<!>)
+    takePositive(x)
     takeNegative(<!ILLEGAL_NUMBER_SIGN!>x<!>) // should be error
     takeAny(x)
 }