[FIR] Introduce ConeFunctionalTypeKind as a replacement of FunctionClassKind

This is needed to provide an ability to extend different kinds of
  functional types

Also, cleanup and rename utilities related to functional types to avoid
  possible confusions

compiler/fir/cones/src/org/jetbrains/kotlin/fir/renderer/ConeTypeRenderer.kt

@@ -5,21 +5,20 @@
 
 package org.jetbrains.kotlin.fir.renderer
 
-import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
 import org.jetbrains.kotlin.fir.types.*
-import kotlin.contracts.ExperimentalContracts
-import kotlin.contracts.contract
 
-open class ConeTypeRenderer() {
+open class ConeTypeRenderer {
 
     lateinit var builder: StringBuilder
     lateinit var idRenderer: ConeIdRenderer
 
     open fun renderAsPossibleFunctionType(
-        type: ConeKotlinType, renderType: ConeTypeProjection.() -> Unit = { render() }
+        type: ConeKotlinType,
+        functionalClassKindExtractor: (ConeKotlinType) -> ConeFunctionalTypeKind?,
+        renderType: ConeTypeProjection.() -> Unit = { render() }
     ) {
-        val kind = type.functionTypeKind
-        if (!kind.withPrettyRender()) {
+        val kind = functionalClassKindExtractor(type)
+        if (kind?.isReflectType != false) {
             type.renderType()
             return
         }
@@ -27,8 +26,9 @@ open class ConeTypeRenderer() {
         if (type.isMarkedNullable) {
             builder.append("(")
         }
-        if (kind == FunctionClassKind.SuspendFunction) {
-            builder.append("suspend ")
+        kind.prefixForTypeRender?.let {
+            builder.append(it)
+            builder.append(" ")
         }
         val typeArguments = type.typeArguments
         val isExtension = type.isExtensionFunctionType
@@ -57,14 +57,6 @@ open class ConeTypeRenderer() {
         }
     }
 
-    @OptIn(ExperimentalContracts::class)
-    private fun FunctionClassKind?.withPrettyRender(): Boolean {
-        contract {
-            returns(true) implies (this@withPrettyRender != null)
-        }
-        return this != null && this != FunctionClassKind.KSuspendFunction && this != FunctionClassKind.KFunction
-    }
-
     fun render(type: ConeKotlinType) {
         if (type !is ConeFlexibleType && type !is ConeDefinitelyNotNullType) {
             // We don't render attributes for flexible/definitely not null types here,
@@ -216,4 +208,4 @@ open class ConeTypeRenderer() {
             }
         }
     }
-}
+}

compiler/fir/cones/src/org/jetbrains/kotlin/fir/renderer/ConeTypeRendererForDebugging.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.renderer
 
+import org.jetbrains.kotlin.fir.types.ConeFunctionalTypeKind
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeTypeProjection
 
@@ -16,9 +17,13 @@ class ConeTypeRendererForDebugging() : ConeTypeRenderer() {
         idRenderer.builder = builder
     }
 
-    override fun renderAsPossibleFunctionType(type: ConeKotlinType, renderType: ConeTypeProjection.() -> Unit) {
+    override fun renderAsPossibleFunctionType(
+        type: ConeKotlinType,
+        functionalClassKindExtractor: (ConeKotlinType) -> ConeFunctionalTypeKind?,
+        renderType: ConeTypeProjection.() -> Unit
+    ) {
         builder.append("R|")
-        super.renderAsPossibleFunctionType(type, renderType)
+        super.renderAsPossibleFunctionType(type, functionalClassKindExtractor, renderType)
         builder.append("|")
     }
-}
+}

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

@@ -0,0 +1,128 @@
+/*
+ * 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.StandardNames
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+/**
+ * [ConeFunctionalTypeKind] describes a family of various similar functional types (like kotlin.FunctionN)
+ *   All types in the same family have corresponding shape of classId: [packageFqName].[classNamePrefix]N,
+ *   where `N` is an arity of function
+ *
+ * Functional type kind may be either reflect type or non-reflect type
+ *   Each non-reflect kind should have corresponding reflect kind and vice-versa (like `kotlin.FunctionN` and `kotlin.reflection.KFunctionN`)
+ *
+ * Classes for such functional types are synthetic and generated by compiler itself
+ * [annotationOnInvokeClassId] is a classId of annotation, which will be added to `invoke` function of generated
+ *   functional interface. This argument is mandatory for non-standard functional kinds
+ *
+ * Assume we have `some.CustomFunctionN` and `some.KCustomFunctionN` families with [annotationOnInvokeClassId] = /some.Anno
+ * Classes for those kinds will look like this:
+ *
+ *   interface CustomFunctionN<P1, P2, ..., PN, R> : kotlin.Function<R> {
+ *       @some.Anno
+ *       operator fun invoke(p1: P1, p2: P2, ..., pN: PN): R
+ *   }
+ *
+ *   interface KCustomFunctionN<P1, P2, ..., PN, R> : kotlin.reflect.KFunction<R>, some.CustomFunctionN<P1, P2, ..., PN, R> {
+ *       @some.Anno
+ *       override operator fun invoke(p1: P1, p2: P2, ..., pN: PN): R
+ *   }
+ *
+ * Note that if you provide some new functional type kind it's your responsibility to handle all references to it in backend
+ *   with [IrGenerationExtension] implementation
+ */
+abstract class ConeFunctionalTypeKind internal constructor(
+    val packageFqName: FqName,
+    val classNamePrefix: String,
+    val isReflectType: Boolean,
+    val annotationOnInvokeClassId: ClassId?
+) {
+    /*
+     * This constructor is needed to enforce not nullable [annotationOnInvokeClassId] for
+     *   functional kinds provided by compiler plugins
+     */
+    constructor(
+        packageFqName: FqName,
+        classNamePrefix: String,
+        annotationOnInvokeClassId: ClassId,
+        isReflectType: Boolean
+    ) : this(packageFqName, classNamePrefix, isReflectType, annotationOnInvokeClassId)
+
+    /*
+     * Specifies how corresponding type will be rendered
+     * E.g. if `prefixForTypeRender = @Some` and type is `some.CustomFunction2<Int, String, Double>`,
+     *   then type will be rendered as `@Some (Int, String) -> Double`
+     */
+    open val prefixForTypeRender: String?
+        get() = null
+
+    /**
+     * @return corresponding non-reflect kind for reflect kind
+     * @return [this] if [isReflectType] is false
+     *
+     * Should be overridden for reflect kinds
+     */
+    open fun nonReflectKind(): ConeFunctionalTypeKind {
+        return if (isReflectType) error("Should be overridden explicitly") else this
+    }
+
+    /**
+     * @return corresponding reflect kind for non-reflect kind
+     * @return [this] if [isReflectType] is true
+     *
+     * Should be overridden for non reflect kinds
+     */
+    open fun reflectKind(): ConeFunctionalTypeKind {
+        return if (isReflectType) this else error("Should be overridden explicitly")
+    }
+
+    fun numberedClassName(arity: Int): Name = Name.identifier("$classNamePrefix$arity")
+
+    // ------------------------------------------- Builtin functional kinds -------------------------------------------
+
+    object Function : ConeFunctionalTypeKind(
+        StandardNames.BUILT_INS_PACKAGE_FQ_NAME,
+        "Function",
+        isReflectType = false,
+        annotationOnInvokeClassId = null
+    ) {
+        override fun reflectKind(): ConeFunctionalTypeKind = KFunction
+    }
+
+    object SuspendFunction : ConeFunctionalTypeKind(
+        StandardNames.COROUTINES_PACKAGE_FQ_NAME,
+        "SuspendFunction",
+        isReflectType = false,
+        annotationOnInvokeClassId = null
+    ) {
+        override val prefixForTypeRender: String
+            get() = "suspend"
+
+        override fun reflectKind(): ConeFunctionalTypeKind = KSuspendFunction
+    }
+
+    object KFunction : ConeFunctionalTypeKind(
+        StandardNames.KOTLIN_REFLECT_FQ_NAME,
+        "KFunction",
+        isReflectType = true,
+        annotationOnInvokeClassId = null
+    ) {
+        override fun nonReflectKind(): ConeFunctionalTypeKind = Function
+    }
+
+    object KSuspendFunction : ConeFunctionalTypeKind(
+        StandardNames.KOTLIN_REFLECT_FQ_NAME,
+        "KSuspendFunction",
+        isReflectType = true,
+        annotationOnInvokeClassId = null
+    ) {
+        override fun nonReflectKind(): ConeFunctionalTypeKind = SuspendFunction
+    }
+}

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

@@ -5,8 +5,10 @@
 
 package org.jetbrains.kotlin.fir.types
 
-import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
-import org.jetbrains.kotlin.fir.renderer.*
+import org.jetbrains.kotlin.fir.renderer.ConeIdRendererForDebugging
+import org.jetbrains.kotlin.fir.renderer.ConeIdShortRenderer
+import org.jetbrains.kotlin.fir.renderer.ConeTypeRendererForDebugging
+import org.jetbrains.kotlin.fir.renderer.ConeTypeRendererWithJavaFlexibleTypes
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.types.Variance
@@ -111,14 +113,6 @@ fun ConeClassLikeType.replaceArgumentsWithStarProjections(): ConeClassLikeType {
     return withArguments(newArguments)
 }
 
-val ConeKotlinType?.functionTypeKind: FunctionClassKind?
-    get() {
-        val classId = (this as? ConeClassLikeType)?.lookupTag?.classId ?: return null
-        return FunctionClassKind.getFunctionalClassKind(
-            classId.shortClassName.asString(), classId.packageFqName
-        )
-    }
-
 fun ConeKotlinType.renderForDebugging(): String {
     val builder = StringBuilder()
     ConeTypeRendererForDebugging(builder).render(this)