JVM: remove LambdaInfo.invokeMethodDescriptor

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/LambdaInfo.kt

@@ -47,7 +47,9 @@ abstract class LambdaInfo(@JvmField val isCrossInline: Boolean) : FunctionalArgu
 
     abstract val invokeMethod: Method
 
-    abstract val invokeMethodDescriptor: FunctionDescriptor
+    abstract val invokeMethodParameters: List<KotlinType?>
+
+    abstract val invokeMethodReturnType: KotlinType?
 
     abstract val capturedVars: List<CapturedParamDesc>
 
@@ -98,14 +100,16 @@ class PsiDefaultLambda(
     offset: Int,
     needReification: Boolean
 ) : DefaultLambda(lambdaClassType, capturedArgs, parameterDescriptor, offset, needReification) {
-    override fun mapAsmSignature(sourceCompiler: SourceCompilerForInline): Method {
+    override fun mapAsmSignature(sourceCompiler: SourceCompilerForInline, descriptor: FunctionDescriptor): Method =
-        return sourceCompiler.state.typeMapper.mapSignatureSkipGeneric(invokeMethodDescriptor).asmMethod
+        sourceCompiler.state.typeMapper.mapSignatureSkipGeneric(descriptor).asmMethod
-    }
 
-    override fun findInvokeMethodDescriptor(): FunctionDescriptor =
+    override fun findInvokeMethodDescriptor(isPropertyReference: Boolean): FunctionDescriptor =
         parameterDescriptor.type.memberScope
             .getContributedFunctions(OperatorNameConventions.INVOKE, NoLookupLocation.FROM_BACKEND)
-            .single()
+            .single().let {
+                // property reference generates erased 'get' method
+                if (isPropertyReference) it.original else it
+            }
 }
 
 abstract class DefaultLambda(
@@ -124,7 +128,13 @@ abstract class DefaultLambda(
     final override lateinit var invokeMethod: Method
         private set
 
-    override lateinit var invokeMethodDescriptor: FunctionDescriptor
+    private lateinit var invokeMethodDescriptor: FunctionDescriptor
+
+    override val invokeMethodParameters: List<KotlinType?>
+        get() = invokeMethodDescriptor.valueParameters.map { it.returnType } // should be FunctionN, so no extension receiver
+
+    override val invokeMethodReturnType: KotlinType?
+        get() = invokeMethodDescriptor.returnType
 
     final override lateinit var capturedVars: List<CapturedParamDesc>
         private set
@@ -155,10 +165,7 @@ abstract class DefaultLambda(
             }
         }, ClassReader.SKIP_CODE or ClassReader.SKIP_FRAMES or ClassReader.SKIP_DEBUG)
 
-        invokeMethodDescriptor = findInvokeMethodDescriptor().let {
+        invokeMethodDescriptor = findInvokeMethodDescriptor(isPropertyReference)
-            //property reference generates erased 'get' method
-            if (isPropertyReference) it.original else it
-        }
 
         val descriptor = Type.getMethodDescriptor(Type.VOID_TYPE, *capturedArgs)
         val constructor = getMethodNode(classBytes, "<init>", descriptor, lambdaClassType)?.node
@@ -182,7 +189,7 @@ abstract class DefaultLambda(
 
         val methodName = (if (isPropertyReference) OperatorNameConventions.GET else OperatorNameConventions.INVOKE).asString()
 
-        val signature = mapAsmSignature(sourceCompiler)
+        val signature = mapAsmSignature(sourceCompiler, invokeMethodDescriptor)
 
         node = getMethodNode(classBytes, methodName, signature.descriptor, lambdaClassType, signatureAmbiguity = true)
             ?: error("Can't find method '$methodName$signature' in '${classReader.className}'")
@@ -195,9 +202,10 @@ abstract class DefaultLambda(
         }
     }
 
-    protected abstract fun mapAsmSignature(sourceCompiler: SourceCompilerForInline): Method
+    protected abstract fun mapAsmSignature(sourceCompiler: SourceCompilerForInline, descriptor: FunctionDescriptor): Method
 
-    protected abstract fun findInvokeMethodDescriptor(): FunctionDescriptor
+    // TODO: get rid of this; descriptors should *only* be used by PsiDefaultLambda
+    protected abstract fun findInvokeMethodDescriptor(isPropertyReference: Boolean): FunctionDescriptor
 
     private companion object {
         val PROPERTY_REFERENCE_SUPER_CLASSES =
@@ -221,7 +229,6 @@ abstract class ExpressionLambda(isCrossInline: Boolean) : LambdaInfo(isCrossInli
         node.node.preprocessSuspendMarkers(forInline = true, keepFakeContinuation = false)
     }
 
-    abstract fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor
     abstract fun isCapturedSuspend(desc: CapturedParamDesc): Boolean
 }
 
@@ -237,7 +244,24 @@ class PsiExpressionLambda(
 
     override val invokeMethod: Method
 
-    override val invokeMethodDescriptor: FunctionDescriptor
+    val invokeMethodDescriptor: FunctionDescriptor
+
+    override val invokeMethodParameters: List<KotlinType?>
+        get() {
+            val actualInvokeDescriptor = if (isSuspend)
+                getOrCreateJvmSuspendFunctionView(
+                    invokeMethodDescriptor,
+                    languageVersionSettings.isReleaseCoroutines(),
+                    typeMapper.bindingContext
+                )
+            else
+                invokeMethodDescriptor
+            val valueParameters = actualInvokeDescriptor.valueParameters.map { it.returnType }
+            return actualInvokeDescriptor.extensionReceiverParameter?.let { listOf(it.returnType) + valueParameters } ?: valueParameters
+        }
+
+    override val invokeMethodReturnType: KotlinType?
+        get() = invokeMethodDescriptor.returnType
 
     val classDescriptor: ClassDescriptor
 
@@ -325,14 +349,6 @@ class PsiExpressionLambda(
     val isPropertyReference: Boolean
         get() = propertyReferenceInfo != null
 
-    override fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor {
-        return getOrCreateJvmSuspendFunctionView(
-            invokeMethodDescriptor,
-            languageVersionSettings.isReleaseCoroutines(),
-            typeMapper.bindingContext
-        )
-    }
-
     override fun isCapturedSuspend(desc: CapturedParamDesc): Boolean =
         isCapturedSuspendLambda(closure, desc.fieldName, typeMapper.bindingContext)
 }

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/MethodInliner.kt

@@ -23,13 +23,13 @@ import org.jetbrains.kotlin.codegen.optimization.fixStack.top
 import org.jetbrains.kotlin.codegen.optimization.nullCheck.isCheckParameterIsNotNull
 import org.jetbrains.kotlin.codegen.pseudoInsns.PseudoInsn
 import org.jetbrains.kotlin.config.LanguageFeature
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
-import org.jetbrains.kotlin.descriptors.ParameterDescriptor
-import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.load.java.JvmAbi
+import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes.OBJECT_TYPE
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.SmartList
 import org.jetbrains.kotlin.utils.SmartSet
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
@@ -233,38 +233,24 @@ class MethodInliner(
                         return
                     }
 
-                    // in case of inlining suspend lambda reference as ordinary parameter of inline function:
+                    val valueParameters = info.invokeMethodParameters
-                    //   suspend fun foo (...) ...
+                    // If suspend lambda has no continuation at the end, add it here. (In the IR backend, suspend lambdas
-                    //   inline fun inlineMe(c: (...) -> ...) ...
+                    // always come pre-lowered with the continuation parameter.)
-                    //   builder {
-                    //     inlineMe(::foo)
-                    //   }
-                    // we should create additional parameter for continuation.
                     var coroutineDesc = desc
-                    val actualInvokeDescriptor: FunctionDescriptor
+                    if (!inliningContext.root.state.isIrBackend && info.isSuspend) {
-                    if (info.isSuspend) {
-                        actualInvokeDescriptor = (info as ExpressionLambda).getInlineSuspendLambdaViewDescriptor()
-                        val parametersSize = actualInvokeDescriptor.valueParameters.size +
-                                (if (actualInvokeDescriptor.extensionReceiverParameter != null) 1 else 0)
-                        // And here we expect invoke(...Ljava/lang/Object;) be replaced with invoke(...Lkotlin/coroutines/Continuation;)
-                        // if this does not happen, insert fake continuation, since we could not have one yet.
                         val argumentTypes = Type.getArgumentTypes(desc)
-                        if (argumentTypes.size != parametersSize &&
+                        if (argumentTypes.size != valueParameters.size) {
-                            // But do not add it in IR. In IR we already have lowered lambdas with additional parameter, while in Old BE we don't.
-                            !inliningContext.root.state.isIrBackend
-                        ) {
                             addFakeContinuationMarker(this)
                             coroutineDesc = Type.getMethodDescriptor(Type.getReturnType(desc), *argumentTypes, AsmTypes.OBJECT_TYPE)
                         }
-                    } else {
-                        actualInvokeDescriptor = info.invokeMethodDescriptor
                     }
 
-                    val valueParameters =
+                    // TODO: this is a weird way to get N instances of `Any?`
-                        listOfNotNull(actualInvokeDescriptor.extensionReceiverParameter) + actualInvokeDescriptor.valueParameters
+                    val moduleBuiltIns = inliningContext.root.sourceCompilerForInline.compilationContextDescriptor.builtIns
-
+                    val erasedFunctionType = moduleBuiltIns.getFunction(valueParameters.size).defaultType
-                    val erasedInvokeFunction = ClosureCodegen.getErasedInvokeFunction(actualInvokeDescriptor)
+                    val erasedInvokeFunction = erasedFunctionType.memberScope
-                    val invokeParameters = erasedInvokeFunction.valueParameters
+                        .getContributedFunctions(OperatorNameConventions.INVOKE, NoLookupLocation.FROM_BACKEND).first()
+                    val invokeParameters = erasedInvokeFunction.valueParameters.map { it.returnType }
 
                     val valueParamShift = max(nextLocalIndex, markerShift)//NB: don't inline cause it changes
                     val parameterTypesFromDesc = info.invokeMethod.argumentTypes
@@ -305,7 +291,7 @@ class MethodInliner(
                     result.reifiedTypeParametersUsages.mergeAll(lambdaResult.reifiedTypeParametersUsages)
 
                     StackValue
-                        .onStack(info.invokeMethod.returnType, info.invokeMethodDescriptor.returnType)
+                        .onStack(info.invokeMethod.returnType, info.invokeMethodReturnType)
                         .put(OBJECT_TYPE, erasedInvokeFunction.returnType, this)
                     setLambdaInlining(false)
                     addInlineMarker(this, false)
@@ -1075,8 +1061,8 @@ class MethodInliner(
 
         private fun putStackValuesIntoLocalsForLambdaOnInvoke(
             directOrder: List<Type>,
-            directOrderOfArguments: List<ParameterDescriptor>,
+            directOrderOfArguments: List<KotlinType?>,
-            directOrderOfInvokeParameters: List<ValueParameterDescriptor>,
+            directOrderOfInvokeParameters: List<KotlinType?>,
             shift: Int,
             iv: InstructionAdapter,
             descriptor: String
@@ -1099,8 +1085,8 @@ class MethodInliner(
                 val argumentKotlinType: KotlinType?
                 val invokeParameterKotlinType: KotlinType?
                 if (safeToUseArgumentKotlinType) {
-                    argumentKotlinType = directOrderOfArguments[index].type
+                    argumentKotlinType = directOrderOfArguments[index]
-                    invokeParameterKotlinType = directOrderOfInvokeParameters[index].type
+                    invokeParameterKotlinType = directOrderOfInvokeParameters[index]
                 } else {
                     argumentKotlinType = null
                     invokeParameterKotlinType = null

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

@@ -20,11 +20,11 @@ import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.types.IrType
-import org.jetbrains.kotlin.ir.types.classOrNull
 import org.jetbrains.kotlin.ir.types.classifierOrFail
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
+import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.commons.Method
@@ -201,27 +201,30 @@ class IrExpressionLambdaImpl(
 
     private val loweredMethod = codegen.methodSignatureMapper.mapAsmMethod(function)
 
-    val capturedParamsInDesc: List<Type> = if (isBoundCallableReference) {
+    val capturedParamsInDesc: List<Type> = when {
-        loweredMethod.argumentTypes.take(1)
+        isBoundCallableReference -> loweredMethod.argumentTypes.take(1)
-    } else loweredMethod.argumentTypes.drop(if (isExtensionLambda) 1 else 0).take(capturedVars.size)
+        isExtensionLambda -> loweredMethod.argumentTypes.drop(1).take(capturedVars.size)
-
+        else -> loweredMethod.argumentTypes.take(capturedVars.size)
-    override val invokeMethod: Method = loweredMethod.let {
-        Method(
-            it.name,
-            it.returnType,
-            (if (isBoundCallableReference) it.argumentTypes.drop(1)
-            else (if (isExtensionLambda) it.argumentTypes.take(1) else emptyList()) +
-                    it.argumentTypes.drop((if (isExtensionLambda) 1 else 0) + capturedVars.size)).toTypedArray()
-        )
     }
 
-    // Need the descriptor without captured parameters here.
+    override val invokeMethod: Method = loweredMethod.let {
-    override val invokeMethodDescriptor: FunctionDescriptor = function.originalFunction.toIrBasedDescriptor()
+        val nonCapturedParameters = when {
+            isBoundCallableReference -> it.argumentTypes.drop(1)
+            isExtensionLambda -> it.argumentTypes.take(1) + it.argumentTypes.drop(capturedVars.size + 1)
+            else -> it.argumentTypes.drop(capturedVars.size)
+        }.toTypedArray()
+        Method(it.name, it.returnType, nonCapturedParameters)
+    }
+
+    // TODO: no extension receiver if bound
+    override val invokeMethodParameters: List<KotlinType?>
+        get() = function.originalFunction.explicitParameters.map { it.type.toIrBasedKotlinType() }
+
+    override val invokeMethodReturnType: KotlinType
+        get() = function.originalFunction.returnType.toIrBasedKotlinType() // not including COROUTINE_SUSPENDED
 
     override val hasDispatchReceiver: Boolean = false
 
-    override fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor = function.toIrBasedDescriptor()
-
     override fun isCapturedSuspend(desc: CapturedParamDesc): Boolean =
         capturedParameters[desc]?.let { it.isInlineParameter() && it.type.isSuspendFunctionTypeOrSubtype() } == true
 }
@@ -229,23 +232,20 @@ class IrExpressionLambdaImpl(
 class IrDefaultLambda(
     lambdaClassType: Type,
     capturedArgs: Array<Type>,
-    private val irValueParameter: IrValueParameter,
+    irValueParameter: IrValueParameter,
     offset: Int,
     needReification: Boolean
 ) : DefaultLambda(
     lambdaClassType, capturedArgs, irValueParameter.toIrBasedDescriptor() as ValueParameterDescriptor, offset, needReification
 ) {
+    private val invoke =
+        (irValueParameter.type.classifierOrFail.owner as IrClass).functions.single { it.name == OperatorNameConventions.INVOKE }
 
-    override fun mapAsmSignature(sourceCompiler: SourceCompilerForInline): Method {
+    override fun mapAsmSignature(sourceCompiler: SourceCompilerForInline, descriptor: FunctionDescriptor): Method =
-        val invoke =
+        (sourceCompiler as IrSourceCompilerForInline).codegen.context.methodSignatureMapper.mapSignatureSkipGeneric(invoke).asmMethod
-            irValueParameter.type.classOrNull!!.owner.declarations.filterIsInstance<IrFunction>().single { it.name.asString() == "invoke" }
-        return (sourceCompiler as IrSourceCompilerForInline).codegen.context.methodSignatureMapper.mapSignatureSkipGeneric(invoke).asmMethod
-    }
 
-    override fun findInvokeMethodDescriptor(): FunctionDescriptor =
+    override fun findInvokeMethodDescriptor(isPropertyReference: Boolean): FunctionDescriptor =
-        (irValueParameter.type.classifierOrFail.owner as IrClass).functions.single {
+        invoke.toIrBasedDescriptor()
-            it.name == OperatorNameConventions.INVOKE
-        }.toIrBasedDescriptor()
 }
 
 fun IrExpression.isInlineIrExpression() =