JVM: fix inlining of default lambdas of signature (Result) -> Any

They have no `invoke` bridge, and the overridden invoke expectes a boxed `Result` as an argument.
2021-06-07 11:55:51 +02:00
parent a0a14d9e25
commit 392e4fba42
15 changed files with 104 additions and 55 deletions
@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.codegen.inline
 import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes.*
 import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.org.objectweb.asm.ClassReader
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.Label
@@ -128,13 +127,14 @@ abstract class DefaultLambda(
            loadClassBytesByInternalName(sourceCompiler.state, lambdaClassType.internalName)
        }

-    protected fun loadInvoke(sourceCompiler: SourceCompilerForInline, invokeMethod: Method) {
+    // Returns whether the loaded invoke is erased, i.e. the name equals the fallback and all types are `Object`.
+    protected fun loadInvoke(sourceCompiler: SourceCompilerForInline, erasedName: String, actualMethod: Method): Boolean {
        val classBytes = loadClass(sourceCompiler)
-        val invokeNameFallback = (if (isPropertyReference) OperatorNameConventions.GET else OperatorNameConventions.INVOKE).asString()
        // TODO: `signatureAmbiguity = true` ignores the argument types from `invokeMethod` and only looks at the count.
-        node = getMethodNode(classBytes, invokeMethod.name, invokeMethod.descriptor, lambdaClassType, signatureAmbiguity = true)
-            ?: getMethodNode(classBytes, invokeNameFallback, invokeMethod.descriptor, lambdaClassType, signatureAmbiguity = true)
-                    ?: error("Can't find method '$invokeMethod' in '${lambdaClassType.internalName}'")
+        node = getMethodNode(classBytes, actualMethod.name, actualMethod.descriptor, lambdaClassType, signatureAmbiguity = true)
+            ?: getMethodNode(classBytes, erasedName, actualMethod.descriptor, lambdaClassType, signatureAmbiguity = true)
+                    ?: error("Can't find method '$actualMethod' in '${lambdaClassType.internalName}'")
+        return invokeMethod.run { name == erasedName && returnType == OBJECT_TYPE && argumentTypes.all { it == OBJECT_TYPE } }
    }

    private companion object {
@@ -13,9 +13,6 @@ import org.jetbrains.kotlin.codegen.binding.CalculatedClosure
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
 import org.jetbrains.kotlin.codegen.coroutines.getOrCreateJvmSuspendFunctionView
 import org.jetbrains.kotlin.codegen.state.GenerationState
-import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
-import org.jetbrains.kotlin.config.LanguageVersionSettings
-import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.psi.*
@@ -170,9 +167,7 @@ class PsiInlineCodegen(
            JvmCodegenUtil.getBoundCallableReferenceReceiver(resolvedCall)
        } else null

-        val lambda = PsiExpressionLambda(
-            ktLambda!!, state.typeMapper, state.languageVersionSettings, parameter.isCrossinline, boundReceiver != null
-        )
+        val lambda = PsiExpressionLambda(ktLambda!!, state, parameter.isCrossinline, boundReceiver != null)
        rememberClosure(type, parameter.index, lambda)
        if (boundReceiver != null) {
            // Has to be done immediately to preserve evaluation order.
@@ -225,8 +220,7 @@ private val FunctionDescriptor.explicitParameters

 class PsiExpressionLambda(
    expression: KtExpression,
-    private val typeMapper: KotlinTypeMapper,
-    private val languageVersionSettings: LanguageVersionSettings,
+    private val state: GenerationState,
    isCrossInline: Boolean,
    override val isBoundCallableReference: Boolean
 ) : ExpressionLambda(isCrossInline) {
@@ -239,9 +233,7 @@ class PsiExpressionLambda(
    override val invokeMethodParameters: List<KotlinType?>
        get() {
            val actualInvokeDescriptor = if (isSuspend)
-                getOrCreateJvmSuspendFunctionView(
-                    invokeMethodDescriptor, languageVersionSettings.isReleaseCoroutines(), typeMapper.bindingContext
-                )
+                getOrCreateJvmSuspendFunctionView(invokeMethodDescriptor, state)
            else
                invokeMethodDescriptor
            return actualInvokeDescriptor.explicitParameters.map { it.returnType }
@@ -263,7 +255,7 @@ class PsiExpressionLambda(
    val closure: CalculatedClosure

    init {
-        val bindingContext = typeMapper.bindingContext
+        val bindingContext = state.bindingContext
        val function = bindingContext.get(BindingContext.FUNCTION, functionWithBodyOrCallableReference)
        if (function == null && expression is KtCallableReferenceExpression) {
            val variableDescriptor =
@@ -271,7 +263,7 @@ class PsiExpressionLambda(
                    ?: throw AssertionError("Reference expression not resolved to variable descriptor with accessors: ${expression.getText()}")
            classDescriptor = bindingContext.get(CodegenBinding.CLASS_FOR_CALLABLE, variableDescriptor)
                ?: throw IllegalStateException("Class for callable not found: $variableDescriptor\n${expression.text}")
-            lambdaClassType = typeMapper.mapClass(classDescriptor)
+            lambdaClassType = state.typeMapper.mapClass(classDescriptor)
            val getFunction = PropertyReferenceCodegen.findGetFunction(variableDescriptor)
            invokeMethodDescriptor = PropertyReferenceCodegen.createFakeOpenDescriptor(getFunction, classDescriptor)
            val resolvedCall = expression.callableReference.getResolvedCallWithAssert(bindingContext)
@@ -287,7 +279,7 @@ class PsiExpressionLambda(
        closure = bindingContext.get(CodegenBinding.CLOSURE, classDescriptor)
            ?: throw AssertionError("null closure for lambda ${expression.text}")
        returnLabels = getDeclarationLabels(expression, invokeMethodDescriptor).associateWith { null }
-        invokeMethod = typeMapper.mapAsmMethod(invokeMethodDescriptor)
+        invokeMethod = state.typeMapper.mapAsmMethod(invokeMethodDescriptor)
        isSuspend = invokeMethodDescriptor.isSuspend
    }

@@ -296,13 +288,13 @@ class PsiExpressionLambda(
        arrayListOf<CapturedParamDesc>().apply {
            val captureThis = closure.capturedOuterClassDescriptor
            if (captureThis != null) {
-                add(capturedParamDesc(AsmUtil.CAPTURED_THIS_FIELD, typeMapper.mapType(captureThis.defaultType), isSuspend = false))
+                add(capturedParamDesc(AsmUtil.CAPTURED_THIS_FIELD, state.typeMapper.mapType(captureThis.defaultType), isSuspend = false))
            }

            val capturedReceiver = closure.capturedReceiverFromOuterContext
            if (capturedReceiver != null) {
-                val fieldName = closure.getCapturedReceiverFieldName(typeMapper.bindingContext, languageVersionSettings)
-                val type = typeMapper.mapType(capturedReceiver).let {
+                val fieldName = closure.getCapturedReceiverFieldName(state.typeMapper.bindingContext, state.languageVersionSettings)
+                val type = state.typeMapper.mapType(capturedReceiver).let {
                    if (isBoundCallableReference) AsmUtil.boxType(it) else it
                }
                add(capturedParamDesc(fieldName, type, isSuspend = false))
@@ -336,23 +328,14 @@ class PsiDefaultLambda(
        get() = invokeMethodDescriptor.returnType

    init {
-        val substitutedDescriptor = parameterDescriptor.type.memberScope
+        val name = if (isPropertyReference) OperatorNameConventions.GET else OperatorNameConventions.INVOKE
+        val descriptor = parameterDescriptor.type.memberScope
            .getContributedFunctions(OperatorNameConventions.INVOKE, NoLookupLocation.FROM_BACKEND)
            .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
-        }
-        loadInvoke(sourceCompiler, sourceCompiler.state.typeMapper.mapSignatureSkipGeneric(invokeMethodDescriptor).asmMethod)
+            .let { if (parameterDescriptor.type.isSuspendFunctionType) getOrCreateJvmSuspendFunctionView(it, sourceCompiler.state) else it }
+        // This is technically wrong as it always uses `invoke`, but `loadInvoke` will fall back to `get` which is never mangled...
+        val asmMethod = sourceCompiler.state.typeMapper.mapAsmMethod(descriptor)
+        val invokeIsErased = loadInvoke(sourceCompiler, name.asString(), asmMethod)
+        invokeMethodDescriptor = if (invokeIsErased) descriptor.original else descriptor
    }
 }
@@ -90,7 +90,7 @@ class PsiSourceCompilerForInline(
        get() = codegen.parentCodegen.orCreateSourceMapper

    override fun generateLambdaBody(lambdaInfo: ExpressionLambda, reifiedTypeParameters: ReifiedTypeParametersUsages): SMAPAndMethodNode {
-        lambdaInfo as? PsiExpressionLambda ?: error("TODO")
+        require(lambdaInfo is PsiExpressionLambda)
        val invokeMethodDescriptor = lambdaInfo.invokeMethodDescriptor
        val jvmMethodSignature = state.typeMapper.mapSignatureSkipGeneric(invokeMethodDescriptor)
        val asmMethod = jvmMethodSignature.asmMethod
@@ -1765,6 +1765,12 @@ public class FirBlackBoxInlineCodegenTestGenerated extends AbstractFirBlackBoxIn
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/kt26636.kt");
            }

+            @Test
+            @TestMetadata("lambdaTakesResult.kt")
+            public void testLambdaTakesResult() throws Exception {
+                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/lambdaTakesResult.kt");
+            }
+
            @Test
            @TestMetadata("noInline.kt")
            public void testNoInline() throws Exception {
@@ -203,23 +203,16 @@ class IrDefaultLambda(
    needReification: Boolean,
    sourceCompiler: IrSourceCompilerForInline
 ) : DefaultLambda(lambdaClassType, capturedArgs, irValueParameter.isCrossinline, offset, needReification, sourceCompiler) {
-    private val typeArguments: List<IrType> = (irValueParameter.type as IrSimpleType).arguments.let {
-        val context = sourceCompiler.codegen.context
-        if (isPropertyReference) {
-            // Property references: `(A) -> B` => `get(Any?): Any?`
-            List(it.size) { context.irBuiltIns.anyNType }
-        } else {
-            // Non-suspend function references and lambdas: `(A) -> B` => `invoke(A): B`
+    private val typeArguments: MutableList<IrType> =
+        (irValueParameter.type as IrSimpleType).arguments.mapTo(mutableListOf()) { (it as IrTypeProjection).type }.apply {
            // 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)
-                }
+            // TODO: default suspend lambdas are currently uninlinable due to having a state machine
+            if (irValueParameter.type.isSuspendFunction()) {
+                val context = sourceCompiler.codegen.context
+                set(size - 1, context.ir.symbols.continuationClass.typeWith(get(size - 1)))
+                add(context.irBuiltIns.anyNType)
            }
        }
-    }

    override val invokeMethodParameters: List<KotlinType>
        get() = typeArguments.dropLast(1).map { it.toIrBasedKotlinType() }
@@ -236,7 +229,10 @@ class IrDefaultLambda(
        )?.let { "$base-$it" } ?: base
        // TODO: while technically only the number of arguments here matters right now (see `loadInvoke`),
        //       it would be better to map to a non-erased signature if not a property reference.
-        loadInvoke(sourceCompiler, Method(name, AsmTypes.OBJECT_TYPE, Array(typeArguments.size - 1) { AsmTypes.OBJECT_TYPE }))
+        if (loadInvoke(sourceCompiler, base, Method(name, AsmTypes.OBJECT_TYPE, Array(typeArguments.size - 1) { AsmTypes.OBJECT_TYPE }))) {
+            // If the loaded method is `invoke(Object, ...) -> Object`, then it expects boxed parameters and returns a boxed value.
+            typeArguments.replaceAll { sourceCompiler.codegen.context.irBuiltIns.anyNType }
+        }
    }
 }

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+// SKIP_INLINE_CHECK_IN: inlineFun$default
+// IGNORE_BACKEND: JS
+// FILE: 1.kt
+package test
+
+inline fun <T> inlineFun(v: T, x: (Result<T>) -> T = { it.getOrNull()!! }) =
+    x(Result.success(v))
+
+// FILE: 2.kt
+import test.*
+
+fun box() = inlineFun("OK")
@@ -1765,6 +1765,12 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/kt26636.kt");
            }

+            @Test
+            @TestMetadata("lambdaTakesResult.kt")
+            public void testLambdaTakesResult() throws Exception {
+                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/lambdaTakesResult.kt");
+            }
+
            @Test
            @TestMetadata("noInline.kt")
            public void testNoInline() throws Exception {
@@ -1765,6 +1765,12 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/kt26636.kt");
            }

+            @Test
+            @TestMetadata("lambdaTakesResult.kt")
+            public void testLambdaTakesResult() throws Exception {
+                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/lambdaTakesResult.kt");
+            }
+
            @Test
            @TestMetadata("noInline.kt")
            public void testNoInline() throws Exception {
@@ -1765,6 +1765,12 @@ public class IrBlackBoxInlineCodegenTestGenerated extends AbstractIrBlackBoxInli
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/kt26636.kt");
            }

+            @Test
+            @TestMetadata("lambdaTakesResult.kt")
+            public void testLambdaTakesResult() throws Exception {
+                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/lambdaTakesResult.kt");
+            }
+
            @Test
            @TestMetadata("noInline.kt")
            public void testNoInline() throws Exception {
@@ -1765,6 +1765,12 @@ public class IrCompileKotlinAgainstInlineKotlinTestGenerated extends AbstractIrC
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/kt26636.kt");
            }

+            @Test
+            @TestMetadata("lambdaTakesResult.kt")
+            public void testLambdaTakesResult() throws Exception {
+                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/lambdaTakesResult.kt");
+            }
+
            @Test
            @TestMetadata("noInline.kt")
            public void testNoInline() throws Exception {
@@ -1765,6 +1765,12 @@ public class JvmIrAgainstOldBoxInlineTestGenerated extends AbstractJvmIrAgainstO
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/kt26636.kt");
            }

+            @Test
+            @TestMetadata("lambdaTakesResult.kt")
+            public void testLambdaTakesResult() throws Exception {
+                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/lambdaTakesResult.kt");
+            }
+
            @Test
            @TestMetadata("noInline.kt")
            public void testNoInline() throws Exception {
@@ -1765,6 +1765,12 @@ public class JvmOldAgainstIrBoxInlineTestGenerated extends AbstractJvmOldAgainst
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/kt26636.kt");
            }

+            @Test
+            @TestMetadata("lambdaTakesResult.kt")
+            public void testLambdaTakesResult() throws Exception {
+                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/lambdaTakesResult.kt");
+            }
+
            @Test
            @TestMetadata("noInline.kt")
            public void testNoInline() throws Exception {
@@ -1393,6 +1393,11 @@ public class IrJsCodegenInlineES6TestGenerated extends AbstractIrJsCodegenInline
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/kt25106.kt");
            }

+            @TestMetadata("lambdaTakesResult.kt")
+            public void testLambdaTakesResult() throws Exception {
+                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/lambdaTakesResult.kt");
+            }
+
            @TestMetadata("noInline.kt")
            public void testNoInline() throws Exception {
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/noInline.kt");
@@ -1393,6 +1393,11 @@ public class IrJsCodegenInlineTestGenerated extends AbstractIrJsCodegenInlineTes
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/kt25106.kt");
            }

+            @TestMetadata("lambdaTakesResult.kt")
+            public void testLambdaTakesResult() throws Exception {
+                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/lambdaTakesResult.kt");
+            }
+
            @TestMetadata("noInline.kt")
            public void testNoInline() throws Exception {
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/noInline.kt");
@@ -1393,6 +1393,11 @@ public class JsCodegenInlineTestGenerated extends AbstractJsCodegenInlineTest {
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/kt25106.kt");
            }

+            @TestMetadata("lambdaTakesResult.kt")
+            public void testLambdaTakesResult() throws Exception {
+                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/lambdaTakesResult.kt");
+            }
+
            @TestMetadata("noInline.kt")
            public void testNoInline() throws Exception {
                runTest("compiler/testData/codegen/boxInline/defaultValues/lambdaInlining/noInline.kt");