From 21f2b3fa2b1a1cfc42822e7f102626a480459f16 Mon Sep 17 00:00:00 2001
From: pyos <pyos@google.com>
Date: Mon, 17 May 2021 13:06:28 +0200
Subject: [PATCH] JVM: expect a continuation parameter in default suspend
 references

---
 .../kotlin/codegen/inline/PsiInlineCodegen.kt | 23 ++++++++++++++-----
 .../backend/jvm/codegen/IrInlineCodegen.kt    | 14 ++++++++---
 .../defaultInlineReference.kt                 |  3 ---
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/PsiInlineCodegen.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/PsiInlineCodegen.kt
index 81907aedcb0..118f13c6d2d 100644
--- a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/PsiInlineCodegen.kt
+++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/PsiInlineCodegen.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.codegen.inline
 
+import org.jetbrains.kotlin.builtins.isSuspendFunctionType
 import org.jetbrains.kotlin.builtins.isSuspendFunctionTypeOrSubtype
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.DescriptorAsmUtil.getMethodAsmFlags
@@ -352,13 +353,23 @@ class PsiDefaultLambda(
         get() = invokeMethodDescriptor.returnType
 
     override fun mapAsmMethod(sourceCompiler: SourceCompilerForInline, isPropertyReference: Boolean): Method {
-        invokeMethodDescriptor = parameterDescriptor.type.memberScope
+        val substitutedDescriptor = parameterDescriptor.type.memberScope
             .getContributedFunctions(OperatorNameConventions.INVOKE, NoLookupLocation.FROM_BACKEND)
-            .single().let {
-                // Property references only have an erased `get` method, while function references and lambdas
-                // have a non-erased `invoke` with an erased synthetic bridge, and we want to inline the former.
-                if (isPropertyReference) it.original else it
-            }
+            .single()
+        invokeMethodDescriptor = when {
+            // Property references: `(A) -> B` => `get(Any?): Any?`
+            isPropertyReference -> substitutedDescriptor.original
+            // Suspend function references: `suspend (A) -> B` => `invoke(A, Continuation<B>): Any?`
+            // TODO: default suspend lambdas are currently uninlinable
+            parameterDescriptor.type.isSuspendFunctionType ->
+                getOrCreateJvmSuspendFunctionView(
+                    substitutedDescriptor,
+                    sourceCompiler.state.languageVersionSettings.isReleaseCoroutines(),
+                    sourceCompiler.state.bindingContext
+                )
+            // Non-suspend function references and lambdas: `(A) -> B` => `invoke(A): B`
+            else -> substitutedDescriptor
+        }
         return sourceCompiler.state.typeMapper.mapSignatureSkipGeneric(invokeMethodDescriptor).asmMethod
     }
 }
diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt
index 0fe25404627..4525b18f8ad 100644
--- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt
+++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt
@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.IrTypeProjection
+import org.jetbrains.kotlin.ir.types.typeWith
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
@@ -251,12 +252,19 @@ class IrDefaultLambda(
     override fun mapAsmMethod(sourceCompiler: SourceCompilerForInline, isPropertyReference: Boolean): Method {
         val context = (sourceCompiler as IrSourceCompilerForInline).codegen.context
         typeArguments = (irValueParameter.type as IrSimpleType).arguments.let {
-            // Property references only have an erased `get` method, while function references and lambdas
-            // have a non-erased `invoke` with an erased synthetic bridge, and we want to inline the former.
             if (isPropertyReference) {
+                // Property references: `(A) -> B` => `get(Any?): Any?`
                 List(it.size) { context.irBuiltIns.anyNType }
             } else {
-                it.map { argument -> (argument as IrTypeProjection).type }
+                // Non-suspend function references and lambdas: `(A) -> B` => `invoke(A): B`
+                // Suspend function references: `suspend (A) -> B` => `invoke(A, Continuation<B>): Any?`
+                // TODO: default suspend lambdas are currently uninlinable
+                it.mapTo(mutableListOf()) { argument -> (argument as IrTypeProjection).type }.apply {
+                    if (irValueParameter.type.isSuspendFunction()) {
+                        set(size - 1, context.ir.symbols.continuationClass.typeWith(get(size - 1)))
+                        add(context.irBuiltIns.anyNType)
+                    }
+                }
             }
         }
         val base = if (isPropertyReference) OperatorNameConventions.GET.asString() else OperatorNameConventions.INVOKE.asString()
diff --git a/compiler/testData/codegen/boxInline/suspend/defaultParameter/defaultInlineReference.kt b/compiler/testData/codegen/boxInline/suspend/defaultParameter/defaultInlineReference.kt
index 76863b47f5e..31edbe8bbae 100644
--- a/compiler/testData/codegen/boxInline/suspend/defaultParameter/defaultInlineReference.kt
+++ b/compiler/testData/codegen/boxInline/suspend/defaultParameter/defaultInlineReference.kt
@@ -1,9 +1,6 @@
 // SKIP_INLINE_CHECK_IN: bar$default
 // WITH_RUNTIME
 // WITH_COROUTINES
-// IGNORE_BACKEND: JVM, JVM_IR
-// IGNORE_BACKEND_FIR: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM, JVM_IR, JVM_MULTI_MODULE_OLD_AGAINST_IR, JVM_MULTI_MODULE_IR_AGAINST_OLD
 // FILE: 1.kt
 package test