JVM: push down implicit coercion to Unit in IR

This is basically a workaround for a slightly different IR generated by fir2ir vs psi2ir. Simplified, psi2ir generates something like this for the sample from KT-59218: TRY type=Unit try: BLOCK type=Unit VAR methodHandle [...] TYPE_OP type=Unit origin=IMPLICIT_COERCION_TO_UNIT CALL invokeExact [...] While fir2ir generates the following: TYPE_OP type=Unit origin=IMPLICIT_COERCION_TO_UNIT TRY type=Any? try: BLOCK type=Any? VAR methodHandle [...] CALL invokeExact [...] The lowering relies on the fact that a polymorphic call (`invokeExact` in this case) is a direct argument to the TYPE_OP, to determine the correct return type (Unit in this case) to be generated in the bytecode. The solution here is to push the type coercion "through" all the block-like structures (`try`, `when`, container expression) so that if the last statement in the block is a polymorphic call, it gets properly converted even if the whole block is under a type coercion operation, as it happens in fir2ir. We achieve that by using the "data" parameter of the IR transformer: appropriate immediate children of IrTypeOperatorCall/IrTry/IrWhen/IrContainerExpression get the type that the expression needs to be coerced to, and all the other expressions ignore that type and set it to null when transforming their children. A proper solution would be to ensure fir2ir generates exactly the same IR as psi2ir (KT-59781), but since PolymorphicSignatureLowering is the only lowering affected so far, and polymorphic calls occur very rarely, it seems safe to workaround it in the lowering for now. #KT-59218 Fixed
2023-06-15 16:04:19 +02:00
parent 69059362b8
commit 1cb1420e43
12 changed files with 323 additions and 17 deletions
@@ -35683,6 +35683,30 @@ public class FirLightTreeBlackBoxCodegenTestGenerated extends AbstractFirLightTr
            runTest("compiler/testData/codegen/box/polymorphicSignature/anonymousSubclass.kt");
        }

+        @Test
+        @TestMetadata("insideComplexExpression.kt")
+        public void testInsideComplexExpression() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideComplexExpression.kt");
+        }
+
+        @Test
+        @TestMetadata("insideIf.kt")
+        public void testInsideIf() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideIf.kt");
+        }
+
+        @Test
+        @TestMetadata("insideTry.kt")
+        public void testInsideTry() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideTry.kt");
+        }
+
+        @Test
+        @TestMetadata("insideWhen.kt")
+        public void testInsideWhen() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideWhen.kt");
+        }
+
        @Test
        @TestMetadata("invoke.kt")
        public void testInvoke() throws Exception {
@@ -35683,6 +35683,30 @@ public class FirPsiBlackBoxCodegenTestGenerated extends AbstractFirPsiBlackBoxCo
            runTest("compiler/testData/codegen/box/polymorphicSignature/anonymousSubclass.kt");
        }

+        @Test
+        @TestMetadata("insideComplexExpression.kt")
+        public void testInsideComplexExpression() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideComplexExpression.kt");
+        }
+
+        @Test
+        @TestMetadata("insideIf.kt")
+        public void testInsideIf() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideIf.kt");
+        }
+
+        @Test
+        @TestMetadata("insideTry.kt")
+        public void testInsideTry() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideTry.kt");
+        }
+
+        @Test
+        @TestMetadata("insideWhen.kt")
+        public void testInsideWhen() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideWhen.kt");
+        }
+
        @Test
        @TestMetadata("invoke.kt")
        public void testInvoke() throws Exception {
@@ -10,10 +10,11 @@ import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.config.LanguageFeature
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
 import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.types.IrType
@@ -21,7 +22,8 @@ import org.jetbrains.kotlin.ir.types.isNullableAny
 import org.jetbrains.kotlin.ir.util.copyTypeParametersFrom
 import org.jetbrains.kotlin.ir.util.dump
 import org.jetbrains.kotlin.ir.util.hasAnnotation
-import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.util.transformInPlace
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
 import org.jetbrains.kotlin.resolve.jvm.checkers.PolymorphicSignatureCallChecker

 internal val polymorphicSignaturePhase = makeIrFilePhase(
@@ -30,28 +32,76 @@ internal val polymorphicSignaturePhase = makeIrFilePhase(
    description = "Replace polymorphic methods with fake ones according to types at the call site"
 )

-class PolymorphicSignatureLowering(val context: JvmBackendContext) : IrElementTransformerVoid(), FileLoweringPass {
+private class PolymorphicSignatureLowering(val context: JvmBackendContext) : IrElementTransformer<PolymorphicSignatureLowering.Data>,
+    FileLoweringPass {
    override fun lower(irFile: IrFile) {
        if (context.state.languageVersionSettings.supportsFeature(LanguageFeature.PolymorphicSignature))
-            irFile.transformChildrenVoid()
+            irFile.transformChildren(this, Data(null))
+    }
+
+    class Data(val coerceToType: IrType?) {
+        companion object {
+            val NO_COERCION = Data(null)
+        }
    }

    private fun IrTypeOperatorCall.isCast(): Boolean =
        operator != IrTypeOperator.INSTANCEOF && operator != IrTypeOperator.NOT_INSTANCEOF

+    override fun visitElement(element: IrElement, data: Data): IrElement {
+        element.transformChildren(this, Data.NO_COERCION)
+        return element
+    }
+
    // If the return type is Any?, then it is also polymorphic (e.g. MethodHandle.invokeExact
    // has polymorphic return type, while VarHandle.compareAndSet does not).
-    override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression =
-        (expression.argument as? IrCall)?.takeIf { expression.isCast() && it.type.isNullableAny() }?.transform(expression.typeOperand)
-            ?: super.visitTypeOperator(expression)
+    override fun visitTypeOperator(expression: IrTypeOperatorCall, data: Data): IrExpression {
+        val argument = expression.argument
+        if (expression.isCast()) {
+            val result = argument.transform(this, Data(expression.typeOperand))
+            if (argument is IrCall && argument.isPolymorphicCall()) return result
+            return expression.apply {
+                this.argument = result
+            }
+        }
+        return super.visitTypeOperator(expression, data)
+    }

-    override fun visitCall(expression: IrCall): IrExpression =
-        expression.transform(null) ?: super.visitCall(expression)
+    override fun visitTry(aTry: IrTry, data: Data): IrExpression {
+        aTry.tryResult = aTry.tryResult.transform(this, data)
+        aTry.catches.transformInPlace(this, data)

-    private fun IrCall.transform(castReturnType: IrType?): IrCall? {
-        val function = symbol.owner as? IrSimpleFunction ?: return null
-        if (!function.hasAnnotation(PolymorphicSignatureCallChecker.polymorphicSignatureFqName))
-            return null
+        // If the try-catch-finally expression is under a type coercion, it needs to be pushed down only to the try and catch blocks,
+        // NOT to the finally block, because the finally block is not an expression.
+        aTry.finallyExpression = aTry.finallyExpression?.transform(this, Data.NO_COERCION)
+
+        return aTry
+    }
+
+    override fun visitWhen(expression: IrWhen, data: Data): IrExpression {
+        expression.branches.transformInPlace(this, data)
+        return expression
+    }
+
+    override fun visitContainerExpression(expression: IrContainerExpression, data: Data): IrExpression {
+        val statements = expression.statements
+        for (i in 0 until statements.size) {
+            val newData = if (i == statements.lastIndex) data else Data.NO_COERCION
+            statements[i] = statements[i].transform(this, newData) as IrStatement
+        }
+        return expression
+    }
+
+    override fun visitCall(expression: IrCall, data: Data): IrElement =
+        if (expression.isPolymorphicCall()) {
+            expression.transform(data.coerceToType)
+        } else super.visitCall(expression, Data.NO_COERCION)
+
+    private fun IrCall.isPolymorphicCall(): Boolean =
+        symbol.owner.hasAnnotation(PolymorphicSignatureCallChecker.polymorphicSignatureFqName)
+
+    private fun IrCall.transform(castReturnType: IrType?): IrCall {
+        val function = symbol.owner
        assert(function.valueParameters.singleOrNull()?.varargElementType != null) {
            "@PolymorphicSignature methods should only have a single vararg argument: ${dump()}"
        }
@@ -67,7 +117,7 @@ class PolymorphicSignatureLowering(val context: JvmBackendContext) : IrElementTr
            updateFrom(function)
            name = function.name
            origin = JvmLoweredDeclarationOrigin.POLYMORPHIC_SIGNATURE_INSTANTIATION
-            returnType = castReturnType ?: function.returnType
+            returnType = if (function.returnType.isNullableAny()) castReturnType ?: function.returnType else function.returnType
        }.apply {
            parent = function.parent
            copyTypeParametersFrom(function)
@@ -85,7 +135,7 @@ class PolymorphicSignatureLowering(val context: JvmBackendContext) : IrElementTr
            dispatchReceiver = this@transform.dispatchReceiver
            extensionReceiver = this@transform.extensionReceiver
            values.forEachIndexed(::putValueArgument)
-            transformChildrenVoid()
+            transformChildren(this@PolymorphicSignatureLowering, Data.NO_COERCION)
        }
    }
 }
@@ -35,7 +35,7 @@ internal val removeDuplicatedInlinedLocalClasses = makeIrFilePhase(
    prerequisite = setOf(functionInliningPhase, localDeclarationsPhase)
 )

-data class Data(
+private data class Data(
    var classDeclaredOnCallSiteOrIsDefaultLambda: Boolean = false,
    var insideInlineBlock: Boolean = false,
    var modifyTree: Boolean = true,
@@ -48,7 +48,7 @@ data class Data(
 // This lambda will not exist after inline, so we copy declaration into new temporary inline call `singleArgumentInlineFunction`.
 // 3. MUST NOT BE created at all because will be created at callee site.
 // This lowering drops declarations that correspond to second and third type.
-class RemoveDuplicatedInlinedLocalClassesLowering(val context: JvmBackendContext) : IrElementTransformer<Data>, FileLoweringPass {
+private class RemoveDuplicatedInlinedLocalClassesLowering(val context: JvmBackendContext) : IrElementTransformer<Data>, FileLoweringPass {
    private val visited = mutableSetOf<IrElement>()
    private val capturedConstructors = context.mapping.capturedConstructors

@@ -0,0 +1,53 @@
+// TARGET_BACKEND: JVM
+// WITH_STDLIB
+// FULL_JDK
+
+import java.lang.invoke.MethodHandles
+import java.lang.invoke.MethodType
+
+object Counter {
+    var result = 0
+    var exception = false
+
+    fun inc() {
+        result++
+        if (exception) throw AssertionError()
+    }
+}
+
+val mh = MethodHandles.lookup().findVirtual(Counter::class.java, "inc", MethodType.methodType(Void.TYPE))
+
+fun f(x: Int) {
+    when (x) {
+        1 -> try {
+            throw Exception()
+        } catch (e: Exception) {
+            if (x > 42) {} else mh.invokeExact(Counter)
+        }
+        2 -> if (x < 42) {
+            if (x < 4242) mh.invokeExact(Counter) else {}
+        }
+        3 -> try {
+            try {
+                mh.invokeExact(Counter)
+            } finally {
+                mh.invokeExact(Counter)
+            }
+        } catch (e: AssertionError) {
+            Counter.exception = false
+            mh.invokeExact(Counter)
+        } finally {
+            mh.invokeExact(Counter)
+        }
+    }
+}
+
+fun box(): String {
+    f(1)
+    f(2)
+
+    Counter.exception = true
+    f(3)
+
+    return if (Counter.result == 6) "OK" else "Fail: ${Counter.result}"
+}
@@ -0,0 +1,20 @@
+// TARGET_BACKEND: JVM
+// WITH_STDLIB
+// FULL_JDK
+
+import java.lang.invoke.MethodHandles
+import java.lang.invoke.MethodType
+
+object O {
+    fun main() {}
+}
+
+fun f() = true
+
+fun box(): String {
+    if (f()) {
+        val mh = MethodHandles.lookup().findVirtual(O::class.java, "main", MethodType.methodType(Void.TYPE))
+        mh.invokeExact(O)
+    } else {}
+    return "OK"
+}
@@ -0,0 +1,18 @@
+// TARGET_BACKEND: JVM
+// WITH_STDLIB
+// FULL_JDK
+
+import java.lang.invoke.MethodHandles
+import java.lang.invoke.MethodType
+
+object O {
+    fun main() {}
+}
+
+fun box(): String {
+    try {
+        val mh = MethodHandles.lookup().findVirtual(O::class.java, "main", MethodType.methodType(Void.TYPE))
+        mh.invokeExact(O)
+    } finally {}
+    return "OK"
+}
@@ -0,0 +1,25 @@
+// TARGET_BACKEND: JVM
+// WITH_STDLIB
+// FULL_JDK
+
+import java.lang.invoke.MethodHandles
+import java.lang.invoke.MethodType
+
+object O {
+    fun main() {}
+}
+
+fun f() = 2
+
+fun box(): String {
+    when (f()) {
+        1 -> {}
+        2 -> {
+            val mh = MethodHandles.lookup().findVirtual(O::class.java, "main", MethodType.methodType(Void.TYPE))
+            mh.invokeExact(O)
+        }
+        3 -> {}
+        else -> {}
+    }
+    return "OK"
+}
@@ -33973,6 +33973,30 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
            runTest("compiler/testData/codegen/box/polymorphicSignature/anonymousSubclass.kt");
        }

+        @Test
+        @TestMetadata("insideComplexExpression.kt")
+        public void testInsideComplexExpression() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideComplexExpression.kt");
+        }
+
+        @Test
+        @TestMetadata("insideIf.kt")
+        public void testInsideIf() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideIf.kt");
+        }
+
+        @Test
+        @TestMetadata("insideTry.kt")
+        public void testInsideTry() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideTry.kt");
+        }
+
+        @Test
+        @TestMetadata("insideWhen.kt")
+        public void testInsideWhen() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideWhen.kt");
+        }
+
        @Test
        @TestMetadata("invoke.kt")
        public void testInvoke() throws Exception {
@@ -35683,6 +35683,30 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
            runTest("compiler/testData/codegen/box/polymorphicSignature/anonymousSubclass.kt");
        }

+        @Test
+        @TestMetadata("insideComplexExpression.kt")
+        public void testInsideComplexExpression() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideComplexExpression.kt");
+        }
+
+        @Test
+        @TestMetadata("insideIf.kt")
+        public void testInsideIf() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideIf.kt");
+        }
+
+        @Test
+        @TestMetadata("insideTry.kt")
+        public void testInsideTry() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideTry.kt");
+        }
+
+        @Test
+        @TestMetadata("insideWhen.kt")
+        public void testInsideWhen() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideWhen.kt");
+        }
+
        @Test
        @TestMetadata("invoke.kt")
        public void testInvoke() throws Exception {
@@ -35683,6 +35683,30 @@ public class IrBlackBoxCodegenWithIrInlinerTestGenerated extends AbstractIrBlack
            runTest("compiler/testData/codegen/box/polymorphicSignature/anonymousSubclass.kt");
        }

+        @Test
+        @TestMetadata("insideComplexExpression.kt")
+        public void testInsideComplexExpression() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideComplexExpression.kt");
+        }
+
+        @Test
+        @TestMetadata("insideIf.kt")
+        public void testInsideIf() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideIf.kt");
+        }
+
+        @Test
+        @TestMetadata("insideTry.kt")
+        public void testInsideTry() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideTry.kt");
+        }
+
+        @Test
+        @TestMetadata("insideWhen.kt")
+        public void testInsideWhen() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideWhen.kt");
+        }
+
        @Test
        @TestMetadata("invoke.kt")
        public void testInvoke() throws Exception {
@@ -30447,6 +30447,26 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
            runTest("compiler/testData/codegen/box/polymorphicSignature/anonymousSubclass.kt");
        }

+        @TestMetadata("insideComplexExpression.kt")
+        public void testInsideComplexExpression() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideComplexExpression.kt");
+        }
+
+        @TestMetadata("insideIf.kt")
+        public void testInsideIf() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideIf.kt");
+        }
+
+        @TestMetadata("insideTry.kt")
+        public void testInsideTry() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideTry.kt");
+        }
+
+        @TestMetadata("insideWhen.kt")
+        public void testInsideWhen() throws Exception {
+            runTest("compiler/testData/codegen/box/polymorphicSignature/insideWhen.kt");
+        }
+
        @TestMetadata("invoke.kt")
        public void testInvoke() throws Exception {
            runTest("compiler/testData/codegen/box/polymorphicSignature/invoke.kt");