JVM IR: support suspend inline functions in -Xmultifile-parts-inherit mode

Support in the normal (without -Xmultifile-parts-inherit) mode is a bit
more complicated, see the added test.

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt

@@ -99,13 +99,19 @@ class JvmBackendContext(
     val suspendFunctionContinuations = mutableMapOf<IrFunction, IrClass>()
     val suspendLambdaToOriginalFunctionMap = mutableMapOf<IrClass, IrFunction>()
     val continuationClassBuilders = mutableMapOf<IrClass, ClassBuilder>()
-    val suspendFunctionViews = mutableMapOf<IrFunction, IrFunction>()
+    val suspendFunctionOriginalToView = mutableMapOf<IrFunction, IrFunction>()
+    val suspendFunctionViewToOriginal = mutableMapOf<IrFunction, IrFunction>()
     val fakeContinuation: IrExpression = createFakeContinuation(this)
 
     val staticDefaultStubs = mutableMapOf<IrFunctionSymbol, IrFunction>()
 
     val inlineClassReplacements = MemoizedInlineClassReplacements()
 
+    internal fun recordSuspendFunctionView(function: IrFunction, view: IrFunction) {
+        suspendFunctionOriginalToView[function] = view
+        suspendFunctionViewToOriginal[view] = function
+    }
+
     internal fun referenceClass(descriptor: ClassDescriptor): IrClassSymbol =
         symbolTable.referenceClass(descriptor)
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt

@@ -126,7 +126,8 @@ internal fun IrFunction.shouldNotContainSuspendMarkers(context: JvmBackendContex
 // the result is called 'view', just to be consistent with old backend.
 internal fun IrFunction.getOrCreateSuspendFunctionViewIfNeeded(context: JvmBackendContext): IrFunction {
     if (!isSuspend || origin == JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW) return this
-    return if (isSuspend) context.suspendFunctionViews.getOrElse(this) { suspendFunctionView(context) } else this
+    context.suspendFunctionOriginalToView[this]?.let { return it }
+    return suspendFunctionView(context)
 }
 
 private fun IrFunction.suspendFunctionView(context: JvmBackendContext): IrFunction {
@@ -163,9 +164,11 @@ private fun IrFunction.suspendFunctionView(context: JvmBackendContext): IrFuncti
             else context.ir.symbols.continuationClass.createType(false, listOf(makeTypeProjection(returnType, Variance.INVARIANT)))
         )
         val valueParametersMapping = explicitParameters.zip(it.explicitParameters).toMap()
+
         // Add the suspend function view to the map before transforming the body to make sure
         // that recursive suspend functions do not lead to unbounded recursion at compile time.
-        context.suspendFunctionViews.put(this, it)
+        context.recordSuspendFunctionView(this, it)
+
         it.body = body?.deepCopyWithSymbols(this)
         it.body?.transformChildrenVoid(object : IrElementTransformerVoid() {
             override fun visitGetValue(expression: IrGetValue): IrGetValue =

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
+import org.jetbrains.kotlin.ir.util.isSuspend
 import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.util.render
 import org.jetbrains.kotlin.load.java.JvmAbi
@@ -112,8 +113,17 @@ class IrSourceCompilerForInline(
         }
 
         if (parent.fileParent.fileEntry is MultifileFacadeFileEntry) {
-            return codegen.context.multifileFacadeMemberToPartMember[callee]
-                ?: error("Function from a multi-file facade without the link to the function in the part: ${callee.render()}")
+            codegen.context.multifileFacadeMemberToPartMember[callee]?.let { return it }
+
+            if (callee.isSuspend) {
+                codegen.context.suspendFunctionViewToOriginal[callee]?.let { facadeMemberOriginal ->
+                    codegen.context.multifileFacadeMemberToPartMember[facadeMemberOriginal]?.let { partMemberOriginal ->
+                        return partMemberOriginal.getOrCreateSuspendFunctionViewIfNeeded(codegen.context)
+                    }
+                }
+            }
+
+            error("Function from a multi-file facade without the link to the function in the part: ${callee.render()}")
         }
 
         return callee

compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClass.kt

@@ -0,0 +1,42 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// COMMON_COROUTINES_TEST
+// IGNORE_LIGHT_ANALYSIS
+// !INHERIT_MULTIFILE_PARTS
+// TARGET_BACKEND: JVM
+
+// FILE: test.kt
+
+@file:JvmMultifileClass
+@file:JvmName("Test")
+
+package test
+
+import helpers.*
+import COROUTINES_PACKAGE.*
+import COROUTINES_PACKAGE.intrinsics.*
+
+suspend fun foo(): String = bar("OK")
+
+suspend inline fun bar(result: String): String = suspendCoroutineUninterceptedOrReturn { x ->
+    x.resume(result)
+    COROUTINE_SUSPENDED
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+fun test(): String {
+    var result = ""
+
+    builder {
+        result = foo()
+    }
+
+    return result
+}
+
+// FILE: box.kt
+
+fun box(): String = test.test()

compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClassUnoptimized.kt

@@ -0,0 +1,44 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// COMMON_COROUTINES_TEST
+// TARGET_BACKEND: JVM
+
+// IGNORE_BACKEND: JVM_IR
+// When -Xmultifile-parts-inherit is disabled, JVM IR backend generates "bridges" that delegate into part members and puts them into
+// the multifile facade. But since the multifile facade phase happens after coroutines, continuations are not created for suspend functions.
+
+// FILE: test.kt
+
+@file:JvmMultifileClass
+@file:JvmName("Test")
+
+package test
+
+import helpers.*
+import COROUTINES_PACKAGE.*
+import COROUTINES_PACKAGE.intrinsics.*
+
+suspend fun foo(): String = bar("OK")
+
+suspend inline fun bar(result: String): String = suspendCoroutineUninterceptedOrReturn { x ->
+    x.resume(result)
+    COROUTINE_SUSPENDED
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+fun test(): String {
+    var result = ""
+
+    builder {
+        result = foo()
+    }
+
+    return result
+}
+
+// FILE: box.kt
+
+fun box(): String = test.test()

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -6064,6 +6064,26 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunInGenericClass.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("inlineFunctionInMultifileClassUnoptimized.kt")
+        public void testInlineFunctionInMultifileClassUnoptimized_1_2() throws Exception {
+            runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClassUnoptimized.kt", "kotlin.coroutines.experimental");
+        }
+
+        @TestMetadata("inlineFunctionInMultifileClassUnoptimized.kt")
+        public void testInlineFunctionInMultifileClassUnoptimized_1_3() throws Exception {
+            runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClassUnoptimized.kt", "kotlin.coroutines");
+        }
+
+        @TestMetadata("inlineFunctionInMultifileClass.kt")
+        public void testInlineFunctionInMultifileClass_1_2() throws Exception {
+            runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClass.kt", "kotlin.coroutines.experimental");
+        }
+
+        @TestMetadata("inlineFunctionInMultifileClass.kt")
+        public void testInlineFunctionInMultifileClass_1_3() throws Exception {
+            runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClass.kt", "kotlin.coroutines");
+        }
+
         @TestMetadata("inlineGenericFunCalledFromSubclass.kt")
         public void testInlineGenericFunCalledFromSubclass_1_2() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineGenericFunCalledFromSubclass.kt", "kotlin.coroutines.experimental");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -6064,6 +6064,26 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunInGenericClass.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("inlineFunctionInMultifileClassUnoptimized.kt")
+        public void testInlineFunctionInMultifileClassUnoptimized_1_2() throws Exception {
+            runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClassUnoptimized.kt", "kotlin.coroutines.experimental");
+        }
+
+        @TestMetadata("inlineFunctionInMultifileClassUnoptimized.kt")
+        public void testInlineFunctionInMultifileClassUnoptimized_1_3() throws Exception {
+            runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClassUnoptimized.kt", "kotlin.coroutines");
+        }
+
+        @TestMetadata("inlineFunctionInMultifileClass.kt")
+        public void testInlineFunctionInMultifileClass_1_2() throws Exception {
+            runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClass.kt", "kotlin.coroutines.experimental");
+        }
+
+        @TestMetadata("inlineFunctionInMultifileClass.kt")
+        public void testInlineFunctionInMultifileClass_1_3() throws Exception {
+            runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClass.kt", "kotlin.coroutines");
+        }
+
         @TestMetadata("inlineGenericFunCalledFromSubclass.kt")
         public void testInlineGenericFunCalledFromSubclass_1_2() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineGenericFunCalledFromSubclass.kt", "kotlin.coroutines.experimental");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -5899,6 +5899,16 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunInGenericClass.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("inlineFunctionInMultifileClassUnoptimized.kt")
+        public void testInlineFunctionInMultifileClassUnoptimized_1_3() throws Exception {
+            runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClassUnoptimized.kt", "kotlin.coroutines");
+        }
+
+        @TestMetadata("inlineFunctionInMultifileClass.kt")
+        public void testInlineFunctionInMultifileClass_1_3() throws Exception {
+            runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineFunctionInMultifileClass.kt", "kotlin.coroutines");
+        }
+
         @TestMetadata("inlineGenericFunCalledFromSubclass.kt")
         public void testInlineGenericFunCalledFromSubclass_1_3() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/inlineGenericFunCalledFromSubclass.kt", "kotlin.coroutines");