JVM: remove LambdaInfo.invokeMethodDescriptor

2021-05-10 17:22:28 +02:00
parent 47dc524763
commit 614d529168
3 changed files with 86 additions and 84 deletions
@@ -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 {
-        return sourceCompiler.state.typeMapper.mapSignatureSkipGeneric(invokeMethodDescriptor).asmMethod
-    }
+    override fun mapAsmSignature(sourceCompiler: SourceCompilerForInline, descriptor: FunctionDescriptor): Method =
+        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 {
-            //property reference generates erased 'get' method
-            if (isPropertyReference) it.original else it
-        }
+        invokeMethodDescriptor = findInvokeMethodDescriptor(isPropertyReference)

        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)
 }
@@ -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.descriptors.ParameterDescriptor
-import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 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:
-                    //   suspend fun foo (...) ...
-                    //   inline fun inlineMe(c: (...) -> ...) ...
-                    //   builder {
-                    //     inlineMe(::foo)
-                    //   }
-                    // we should create additional parameter for continuation.
+                    val valueParameters = info.invokeMethodParameters
+                    // If suspend lambda has no continuation at the end, add it here. (In the IR backend, suspend lambdas
+                    // always come pre-lowered with the continuation parameter.)
                    var coroutineDesc = desc
-                    val actualInvokeDescriptor: FunctionDescriptor
-                    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.
+                    if (!inliningContext.root.state.isIrBackend && info.isSuspend) {
                        val argumentTypes = Type.getArgumentTypes(desc)
-                        if (argumentTypes.size != parametersSize &&
-                            // 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
-                        ) {
+                        if (argumentTypes.size != valueParameters.size) {
                            addFakeContinuationMarker(this)
                            coroutineDesc = Type.getMethodDescriptor(Type.getReturnType(desc), *argumentTypes, AsmTypes.OBJECT_TYPE)
                        }
-                    } else {
-                        actualInvokeDescriptor = info.invokeMethodDescriptor
                    }

-                    val valueParameters =
-                        listOfNotNull(actualInvokeDescriptor.extensionReceiverParameter) + actualInvokeDescriptor.valueParameters
-
-                    val erasedInvokeFunction = ClosureCodegen.getErasedInvokeFunction(actualInvokeDescriptor)
-                    val invokeParameters = erasedInvokeFunction.valueParameters
+                    // TODO: this is a weird way to get N instances of `Any?`
+                    val moduleBuiltIns = inliningContext.root.sourceCompilerForInline.compilationContextDescriptor.builtIns
+                    val erasedFunctionType = moduleBuiltIns.getFunction(valueParameters.size).defaultType
+                    val erasedInvokeFunction = erasedFunctionType.memberScope
+                        .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>,
-            directOrderOfInvokeParameters: List<ValueParameterDescriptor>,
+            directOrderOfArguments: List<KotlinType?>,
+            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
-                    invokeParameterKotlinType = directOrderOfInvokeParameters[index].type
+                    argumentKotlinType = directOrderOfArguments[index]
+                    invokeParameterKotlinType = directOrderOfInvokeParameters[index]
                } else {
                    argumentKotlinType = null
                    invokeParameterKotlinType = null
@@ -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) {
-        loweredMethod.argumentTypes.take(1)
-    } else loweredMethod.argumentTypes.drop(if (isExtensionLambda) 1 else 0).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()
-        )
+    val capturedParamsInDesc: List<Type> = when {
+        isBoundCallableReference -> loweredMethod.argumentTypes.take(1)
+        isExtensionLambda -> loweredMethod.argumentTypes.drop(1).take(capturedVars.size)
+        else -> loweredMethod.argumentTypes.take(capturedVars.size)
    }

-    // Need the descriptor without captured parameters here.
-    override val invokeMethodDescriptor: FunctionDescriptor = function.originalFunction.toIrBasedDescriptor()
+    override val invokeMethod: Method = loweredMethod.let {
+        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 {
-        val invoke =
-            irValueParameter.type.classOrNull!!.owner.declarations.filterIsInstance<IrFunction>().single { it.name.asString() == "invoke" }
-        return (sourceCompiler as IrSourceCompilerForInline).codegen.context.methodSignatureMapper.mapSignatureSkipGeneric(invoke).asmMethod
-    }
+    override fun mapAsmSignature(sourceCompiler: SourceCompilerForInline, descriptor: FunctionDescriptor): Method =
+        (sourceCompiler as IrSourceCompilerForInline).codegen.context.methodSignatureMapper.mapSignatureSkipGeneric(invoke).asmMethod

-    override fun findInvokeMethodDescriptor(): FunctionDescriptor =
-        (irValueParameter.type.classifierOrFail.owner as IrClass).functions.single {
-            it.name == OperatorNameConventions.INVOKE
-        }.toIrBasedDescriptor()
+    override fun findInvokeMethodDescriptor(isPropertyReference: Boolean): FunctionDescriptor =
+        invoke.toIrBasedDescriptor()
 }

 fun IrExpression.isInlineIrExpression() =