[FIR] Properly implement FirFunctionalTypeKindService

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

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/LLFirBuiltinsSessionFactory.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
 import java.util.concurrent.ConcurrentHashMap
 import org.jetbrains.kotlin.config.LanguageVersionSettingsImpl
 import org.jetbrains.kotlin.fir.resolve.transformers.FirDummyCompilerLazyDeclarationResolver
+import org.jetbrains.kotlin.fir.session.registerCommonComponentsAfterExtensionsAreConfigured
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.platform.isCommon
 import org.jetbrains.kotlin.platform.isJs
@@ -54,6 +55,7 @@ class LLFirBuiltinsSessionFactory(
             registerIdeComponents(project)
             register(FirLazyDeclarationResolver::class, FirDummyCompilerLazyDeclarationResolver)
             registerCommonComponents(LanguageVersionSettingsImpl.DEFAULT/*TODO*/)
+            registerCommonComponentsAfterExtensionsAreConfigured()
             registerCommonJavaComponents(JavaModuleResolver.getInstance(project))
             registerModuleData(moduleData)
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/LLFirLibrarySessionFactory.kt

@@ -22,10 +22,7 @@ import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
 import org.jetbrains.kotlin.fir.resolve.transformers.FirDummyCompilerLazyDeclarationResolver
 import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
-import org.jetbrains.kotlin.fir.session.registerCommonComponents
-import org.jetbrains.kotlin.fir.session.registerCommonJavaComponents
-import org.jetbrains.kotlin.fir.session.registerJavaSpecificResolveComponents
-import org.jetbrains.kotlin.fir.session.registerModuleData
+import org.jetbrains.kotlin.fir.session.*
 import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
 
@@ -53,6 +50,7 @@ class LLFirLibrarySessionFactory(
             registerIdeComponents(project)
             register(FirLazyDeclarationResolver::class, FirDummyCompilerLazyDeclarationResolver)
             registerCommonComponents(LanguageVersionSettingsImpl.DEFAULT/*TODO*/)
+            registerCommonComponentsAfterExtensionsAreConfigured()
             registerCommonJavaComponents(JavaModuleResolver.getInstance(project))
             registerJavaSpecificResolveComponents()
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionFactory.kt

@@ -108,6 +108,7 @@ internal object LLFirSessionFactory {
 
             registerCompilerPluginServices(contentScope, project, module)
             registerCompilerPluginExtensions(project, module)
+            registerCommonComponentsAfterExtensionsAreConfigured()
 
             val switchableExtensionDeclarationsSymbolProvider = FirSwitchableExtensionDeclarationsSymbolProvider.create(session)?.also {
                 register(FirSwitchableExtensionDeclarationsSymbolProvider::class, it)
@@ -206,6 +207,7 @@ internal object LLFirSessionFactory {
 
             registerIdeComponents(project)
             registerCommonComponents(languageVersionSettings)
+            registerCommonComponentsAfterExtensionsAreConfigured()
             registerCommonJavaComponents(JavaModuleResolver.getInstance(project))
             registerResolveComponents()
             registerJavaSpecificResolveComponents()
@@ -288,6 +290,7 @@ internal object LLFirSessionFactory {
 
             registerIdeComponents(project)
             registerCommonComponents(languageVersionSettings)
+            registerCommonComponentsAfterExtensionsAreConfigured()
             registerCommonJavaComponents(JavaModuleResolver.getInstance(project))
             registerResolveComponents()
             registerJavaSpecificResolveComponents()

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeFunctionalTypeKind.kt

@@ -85,6 +85,10 @@ abstract class ConeFunctionalTypeKind internal constructor(
 
     fun numberedClassName(arity: Int): Name = Name.identifier("$classNamePrefix$arity")
 
+    override fun toString(): String {
+        return "$packageFqName.${classNamePrefix}N"
+    }
+
     // ------------------------------------------- Builtin functional kinds -------------------------------------------
 
     object Function : ConeFunctionalTypeKind(

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)

compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/FirFunctionalTypeKindServiceImpl.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.types
+
+import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.extensions.FirFunctionalTypeKindExtension
+import org.jetbrains.kotlin.fir.extensions.extensionService
+import org.jetbrains.kotlin.fir.extensions.functionalTypeKindExtensions
+import org.jetbrains.kotlin.name.FqName
+
+class FirFunctionalTypeKindServiceImpl(session: FirSession) : FirFunctionalTypeKindService() {
+    private val knownKindsByPackageFqName = buildList {
+        add(ConeFunctionalTypeKind.Function)
+        add(ConeFunctionalTypeKind.SuspendFunction)
+        add(ConeFunctionalTypeKind.KFunction)
+        add(ConeFunctionalTypeKind.KSuspendFunction)
+
+        val registrar = object : FirFunctionalTypeKindExtension.FunctionalTypeKindRegistrar {
+            override fun registerKind(nonReflectKind: ConeFunctionalTypeKind, reflectKind: ConeFunctionalTypeKind) {
+                require(nonReflectKind.reflectKind() == reflectKind)
+                require(reflectKind.nonReflectKind() == nonReflectKind)
+                add(nonReflectKind)
+                add(reflectKind)
+            }
+        }
+
+        for (extension in session.extensionService.functionalTypeKindExtensions) {
+            with(extension) { registrar.registerKinds() }
+        }
+    }.also { kinds ->
+        val allNames = kinds.map { "${it.packageFqName}.${it.classNamePrefix}" }
+        require(allNames.distinct() == allNames) {
+            "There are clashing functional type kinds: $allNames"
+        }
+    }.groupBy { it.packageFqName }
+
+    override fun getKindByClassNamePrefix(packageFqName: FqName, className: String): ConeFunctionalTypeKind? {
+        val kinds = knownKindsByPackageFqName[packageFqName] ?: return null
+        for (kind in kinds) {
+            if (!className.startsWith(kind.classNamePrefix)) continue
+            if (!FunctionClassKind.hasArityAtTheEnd(kind.classNamePrefix, className)) continue
+            return kind
+        }
+        return null
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/FirFunctionalTypeKindExtension.kt

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.extensions
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.types.ConeFunctionalTypeKind
+import kotlin.reflect.KClass
+
+abstract class FirFunctionalTypeKindExtension(session: FirSession) : FirExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("CustomFunctionalTypeKindExtension")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension>
+        get() = FirFunctionalTypeKindExtension::class
+
+    interface FunctionalTypeKindRegistrar {
+        fun registerKind(nonReflectKind: ConeFunctionalTypeKind, reflectKind: ConeFunctionalTypeKind)
+    }
+
+    abstract fun FunctionalTypeKindRegistrar.registerKinds()
+
+    fun interface Factory : FirExtension.Factory<FirFunctionalTypeKindExtension>
+}
+
+val FirExtensionService.functionalTypeKindExtensions: List<FirFunctionalTypeKindExtension> by FirExtensionService.registeredExtensions()

core/compiler.common/src/org/jetbrains/kotlin/builtins/functions/FunctionClassKind.kt

@@ -43,6 +43,10 @@ enum class FunctionClassKind(
         fun getFunctionalClassKind(className: String, packageFqName: FqName) =
             parseClassName(className, packageFqName)?.kind
 
+        fun hasArityAtTheEnd(prefix: String, className: String): Boolean {
+            return toInt(className.substring(prefix.length)) != null
+        }
+
         private fun toInt(s: String): Int? {
             if (s.isEmpty()) return null