From ae906c163d83c2f7a15a2fc85b0bc291d746a198 Mon Sep 17 00:00:00 2001
From: Igor Chevdar <igor.chevdar@jetbrains.com>
Date: Mon, 6 Apr 2020 18:50:56 +0300
Subject: [PATCH] [IR] Added possibility to customize inline functions resolver

---
 .../common/lower/inline/FunctionInlining.kt   | 51 +++++++++++--------
 1 file changed, 31 insertions(+), 20 deletions(-)

diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt
index c20e4ea4b39..b90b2e2d4b2 100644
--- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt
+++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt
@@ -32,7 +32,36 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 fun IrValueParameter.isInlineParameter(type: IrType = this.type) =
     index >= 0 && !isNoinline && !type.isNullable() && (type.isFunction() || type.isSuspendFunction())
 
-class FunctionInlining(val context: CommonBackendContext) : IrElementTransformerVoidWithContext(), BodyLoweringPass {
+interface InlineFunctionResolver {
+    fun getFunctionDeclaration(symbol: IrFunctionSymbol): IrFunction
+}
+
+open class DefaultInlineFunctionResolver(open val context: CommonBackendContext) : InlineFunctionResolver {
+    override fun getFunctionDeclaration(symbol: IrFunctionSymbol): IrFunction {
+        val descriptor = symbol.descriptor.original
+        val languageVersionSettings = context.configuration.languageVersionSettings
+        // TODO: Remove these hacks when coroutine intrinsics are fixed.
+        return when {
+            descriptor.isBuiltInIntercepted(languageVersionSettings) ->
+                error("Continuation.intercepted is not available with release coroutines")
+
+            descriptor.isBuiltInSuspendCoroutineUninterceptedOrReturn(languageVersionSettings) ->
+                context.ir.symbols.suspendCoroutineUninterceptedOrReturn.owner
+
+            symbol == context.ir.symbols.coroutineContextGetter ->
+                context.ir.symbols.coroutineGetContext.owner
+
+            else -> (symbol.owner as? IrSimpleFunction)?.resolveFakeOverride() ?: symbol.owner
+        }
+    }
+}
+
+class FunctionInlining(
+    val context: CommonBackendContext,
+    val inlineFunctionResolver: InlineFunctionResolver
+) : IrElementTransformerVoidWithContext(), BodyLoweringPass {
+
+    constructor(context: CommonBackendContext): this(context, DefaultInlineFunctionResolver(context))
 
     private var containerScope: ScopeWithIr? = null
 
@@ -61,7 +90,7 @@ class FunctionInlining(val context: CommonBackendContext) : IrElementTransformer
         if (Symbols.isTypeOfIntrinsic(callee.symbol))
             return expression
 
-        val actualCallee = getFunctionDeclaration(callee.symbol)
+        val actualCallee = inlineFunctionResolver.getFunctionDeclaration(callee.symbol)
 
         val parent = allScopes.map { it.irElement }.filterIsInstance<IrDeclarationParent>().lastOrNull()
             ?: allScopes.map { it.irElement }.filterIsInstance<IrDeclaration>().lastOrNull()?.parent
@@ -72,24 +101,6 @@ class FunctionInlining(val context: CommonBackendContext) : IrElementTransformer
         return inliner.inline()
     }
 
-    private fun getFunctionDeclaration(symbol: IrFunctionSymbol): IrFunction {
-        val descriptor = symbol.descriptor.original
-        val languageVersionSettings = context.configuration.languageVersionSettings
-        // TODO: Remove these hacks when coroutine intrinsics are fixed.
-        return when {
-            descriptor.isBuiltInIntercepted(languageVersionSettings) ->
-                error("Continuation.intercepted is not available with release coroutines")
-
-            descriptor.isBuiltInSuspendCoroutineUninterceptedOrReturn(languageVersionSettings) ->
-                context.ir.symbols.suspendCoroutineUninterceptedOrReturn.owner
-
-            symbol == context.ir.symbols.coroutineContextGetter ->
-                context.ir.symbols.coroutineGetContext.owner
-
-            else -> (symbol.owner as? IrSimpleFunction)?.resolveFakeOverride() ?: symbol.owner
-        }
-    }
-
     private val IrFunction.needsInlining get() = this.isInline && !this.isExternal
 
     private inner class Inliner(