From 2a08923c511b0602abdc7418001592ddc140db23 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Fri, 23 Apr 2021 09:56:32 +0300
Subject: [PATCH] [FIR] Add more utilities for builtin functional types

---
 .../fir/resolve/inference/InferenceUtils.kt   | 29 +++++++++++--------
 .../builtins/functions/FunctionClassKind.kt   | 15 ++++++----
 2 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/InferenceUtils.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/InferenceUtils.kt
index afe9c783316..89c62e03e76 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/InferenceUtils.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/InferenceUtils.kt
@@ -47,7 +47,7 @@ fun ConeKotlinType.isKMutableProperty(session: FirSession): Boolean {
             classId.shortClassName.identifier.startsWith("KMutableProperty")
 }
 
-private fun ConeKotlinType.functionClassKind(session: FirSession): FunctionClassKind? {
+fun ConeKotlinType.functionClassKind(session: FirSession): FunctionClassKind? {
     return classId(session)?.toFunctionClassKind()
 }
 
@@ -55,29 +55,34 @@ private fun ClassId.toFunctionClassKind(): FunctionClassKind? {
     return FunctionClassKind.byClassNamePrefix(packageFqName, relativeClassName.asString())
 }
 
+// Function, SuspendFunction, KFunction, KSuspendFunction
 fun ConeKotlinType.isBuiltinFunctionalType(session: FirSession): Boolean {
     return functionClassKind(session) != null
 }
 
-fun ConeKotlinType.isFunctionalType(session: FirSession): Boolean {
-    val kind = functionClassKind(session) ?: return false
-    return kind == FunctionClassKind.Function
-}
-
+// Function, SuspendFunction, KFunction, KSuspendFunction
 fun ConeClassLikeLookupTag.isBuiltinFunctionalType(): Boolean {
     return classId.toFunctionClassKind() != null
 }
 
-fun ConeKotlinType.isSuspendFunctionType(session: FirSession): Boolean {
+inline fun ConeKotlinType.isFunctionalType(session: FirSession, predicate: (FunctionClassKind) -> Boolean): Boolean {
     val kind = functionClassKind(session) ?: return false
-    return kind == FunctionClassKind.SuspendFunction ||
-            kind == FunctionClassKind.KSuspendFunction
+    return predicate(kind)
 }
 
+// Function
+fun ConeKotlinType.isFunctionalType(session: FirSession): Boolean {
+    return isFunctionalType(session) { it == FunctionClassKind.Function }
+}
+
+// SuspendFunction, KSuspendFunction
+fun ConeKotlinType.isSuspendFunctionType(session: FirSession): Boolean {
+    return isFunctionalType(session) { it.isSuspendType }
+}
+
+// KFunction, KSuspendFunction
 fun ConeKotlinType.isKFunctionType(session: FirSession): Boolean {
-    val kind = functionClassKind(session) ?: return false
-    return kind == FunctionClassKind.KFunction ||
-            kind == FunctionClassKind.KSuspendFunction
+    return isFunctionalType(session) { it.isReflectType }
 }
 
 fun ConeKotlinType.kFunctionTypeToFunctionType(session: FirSession): ConeClassLikeType {
diff --git a/core/compiler.common/src/org/jetbrains/kotlin/builtins/functions/FunctionClassKind.kt b/core/compiler.common/src/org/jetbrains/kotlin/builtins/functions/FunctionClassKind.kt
index f96053a3cda..e98182b82ac 100644
--- a/core/compiler.common/src/org/jetbrains/kotlin/builtins/functions/FunctionClassKind.kt
+++ b/core/compiler.common/src/org/jetbrains/kotlin/builtins/functions/FunctionClassKind.kt
@@ -9,11 +9,16 @@ import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
-enum class FunctionClassKind(val packageFqName: FqName, val classNamePrefix: String) {
-    Function(StandardNames.BUILT_INS_PACKAGE_FQ_NAME, "Function"),
-    SuspendFunction(StandardNames.COROUTINES_PACKAGE_FQ_NAME_RELEASE, "SuspendFunction"),
-    KFunction(StandardNames.KOTLIN_REFLECT_FQ_NAME, "KFunction"),
-    KSuspendFunction(StandardNames.KOTLIN_REFLECT_FQ_NAME, "KSuspendFunction");
+enum class FunctionClassKind(
+    val packageFqName: FqName,
+    val classNamePrefix: String,
+    val isSuspendType: Boolean,
+    val isReflectType: Boolean
+) {
+    Function(StandardNames.BUILT_INS_PACKAGE_FQ_NAME, "Function", isSuspendType = false, isReflectType = false),
+    SuspendFunction(StandardNames.COROUTINES_PACKAGE_FQ_NAME_RELEASE, "SuspendFunction", isSuspendType = true, isReflectType = false),
+    KFunction(StandardNames.KOTLIN_REFLECT_FQ_NAME, "KFunction", isSuspendType = false, isReflectType = true),
+    KSuspendFunction(StandardNames.KOTLIN_REFLECT_FQ_NAME, "KSuspendFunction", isSuspendType = true, isReflectType = true);
 
     fun numberedClassName(arity: Int) = Name.identifier("$classNamePrefix$arity")