[FIR] Properly implement FirFunctionalTypeKindService

This implementation takes additional functional type kinds from compiler
  plugins into account

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

@@ -32,6 +32,7 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirSamConversionTransformerExtension::class,
             FirAssignExpressionAltererExtension::class,
             FirScriptConfiguratorExtension::class,
+            FirFunctionalTypeKindExtension::class,
         )
     }
 
@@ -90,6 +91,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             registerExtension(FirScriptConfiguratorExtension::class, this)
         }
 
+        @JvmName("plusFunctionalTypeKindExtension")
+        operator fun (FirFunctionalTypeKindExtension.Factory).unaryPlus() {
+            registerExtension(FirFunctionalTypeKindExtension::class, this)
+        }
+
         // ------------------ reference methods ------------------
 
         @JvmName("plusStatusTransformerExtension")
@@ -142,6 +148,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirScriptConfiguratorExtension.Factory { this.invoke(it) }.unaryPlus()
         }
 
+        @JvmName("plusFunctionalTypeKindExtension")
+        operator fun ((FirSession) -> FirFunctionalTypeKindExtension).unaryPlus() {
+            FirFunctionalTypeKindExtension.Factory { this.invoke(it) }.unaryPlus()
+        }
+
         // ------------------ utilities ------------------
 
         @JvmName("bindLeft")

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

@@ -47,6 +47,7 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirSubstitutionOverrideStorage
 import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.fir.types.FirCorrespondingSupertypesCache
 import org.jetbrains.kotlin.fir.types.FirFunctionalTypeKindService
+import org.jetbrains.kotlin.fir.types.FirFunctionalTypeKindServiceImpl
 import org.jetbrains.kotlin.fir.types.TypeComponents
 import org.jetbrains.kotlin.incremental.components.EnumWhenTracker
 import org.jetbrains.kotlin.incremental.components.LookupTracker
@@ -72,9 +73,11 @@ fun FirSession.registerCommonComponents(languageVersionSettings: LanguageVersion
     register(FirSamConstructorStorage::class, FirSamConstructorStorage(this))
     register(FirOverrideService::class, FirOverrideService(this))
     register(FirDynamicMembersStorage::class, FirDynamicMembersStorage(this))
+}
 
-    // TODO: replace with proper implementation
-    register(FirFunctionalTypeKindService::class, FirFunctionalTypeKindService.Default)
+@OptIn(SessionConfiguration::class)
+fun FirSession.registerCommonComponentsAfterExtensionsAreConfigured() {
+    register(FirFunctionalTypeKindService::class, FirFunctionalTypeKindServiceImpl(this))
 }
 
 @OptIn(SessionConfiguration::class)

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

@@ -44,6 +44,7 @@ abstract class FirAbstractSessionFactory {
 
             registerCliCompilerOnlyComponents()
             registerCommonComponents(languageVersionSettings)
+            registerCommonComponentsAfterExtensionsAreConfigured()
             registerExtraComponents(this)
 
             val kotlinScopeProvider = createKotlinScopeProvider.invoke()
@@ -105,6 +106,7 @@ abstract class FirAbstractSessionFactory {
                 }
                 init()
             }.configure()
+            registerCommonComponentsAfterExtensionsAreConfigured()
 
             val dependencyProviders = computeDependencyProviderList(moduleData)
             val generatedSymbolsProvider = FirSwitchableExtensionDeclarationsSymbolProvider.create(this)