JVM IR: Support -Xno-call-assertions

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/JvmArgumentNullabilityAssertionsLowering.kt

@@ -6,249 +6,98 @@
 package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.lower.SpecialBridgeMethods
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.config.ApiVersion
-import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.declarations.IrField
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
-import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
-import org.jetbrains.kotlin.load.java.BuiltinMethodsWithSpecialGenericSignature
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
-val jvmArgumentNullabilityAssertions =
-    makeIrFilePhase(
-        ::JvmArgumentNullabilityAssertionsLowering,
-        name = "Nullability assertions on arguments",
-        description = "Transform nullability assertions on arguments according to the compiler settings"
-    )
+val jvmArgumentNullabilityAssertions = makeIrFilePhase(
+    ::JvmArgumentNullabilityAssertionsLowering,
+    name = "ArgumentNullabilityAssertions",
+    description = "Transform nullability assertions on arguments according to the compiler settings"
+)
 
-class JvmArgumentNullabilityAssertionsLowering(context: JvmBackendContext) :
-    FileLoweringPass, IrElementVisitorVoid {
+private enum class AssertionScope {
+    Enabled, Disabled
+}
 
-    private val isWithUnifiedNullChecks =
-        context.state.languageVersionSettings.apiVersion >= ApiVersion.KOTLIN_1_4
+private class JvmArgumentNullabilityAssertionsLowering(context: JvmBackendContext) : FileLoweringPass, IrElementTransformer<AssertionScope> {
 
+    private val isWithUnifiedNullChecks = context.state.languageVersionSettings.apiVersion >= ApiVersion.KOTLIN_1_4
     private val isCallAssertionsDisabled = context.state.isCallAssertionsDisabled
     private val isReceiverAssertionsDisabled = context.state.isReceiverAssertionsDisabled
 
-    override fun lower(irFile: IrFile) {
-        irFile.acceptChildrenVoid(this)
-    }
+    private val specialBridgeMethods = SpecialBridgeMethods(context)
 
-    override fun visitElement(element: IrElement) {
-        element.acceptChildrenVoid(this)
-    }
+    override fun lower(irFile: IrFile) = irFile.transformChildren(this, AssertionScope.Enabled)
 
-    private inline fun <T : IrElement> T.transformPostfix(fn: T.() -> Unit) {
-        acceptChildrenVoid(this@JvmArgumentNullabilityAssertionsLowering)
-        fn()
-    }
+    override fun visitElement(element: IrElement, data: AssertionScope): IrElement =
+        super.visitElement(element, AssertionScope.Enabled)
 
-    override fun visitMemberAccess(expression: IrMemberAccessExpression) {
-        expression.transformPostfix {
+    override fun visitDeclaration(declaration: IrDeclaration, data: AssertionScope): IrStatement =
+        super.visitDeclaration(declaration, AssertionScope.Enabled)
 
-            // Always drop nullability assertions on dispatch receivers, assuming that it will throw NPE.
-            //
-            // NB there are some members in Kotlin built-in classes which are NOT implemented as platform method calls,
-            // and thus break this assertion - e.g., 'Array<T>.iterator()' and similar functions.
-            // See KT-30908 for more details.
-            dispatchReceiver = dispatchReceiver?.replaceImplicitNotNullWithArgument()
+    override fun visitTypeOperator(expression: IrTypeOperatorCall, data: AssertionScope): IrExpression =
+        if (expression.operator == IrTypeOperator.IMPLICIT_NOTNULL && (data == AssertionScope.Disabled || isCallAssertionsDisabled))
+            expression.argument.transform(this, data)
+        else
+            super.visitTypeOperator(expression, data)
 
-            if (isReceiverAssertionsDisabled || shouldDropNullabilityAssertionOnExtensionReceiver(expression)) {
-                extensionReceiver = extensionReceiver?.replaceImplicitNotNullWithArgument()
-            }
+    override fun visitMemberAccess(expression: IrMemberAccessExpression, data: AssertionScope): IrElement {
+        // Always drop nullability assertions on dispatch receivers, assuming that it will throw NPE.
+        //
+        // NB there are some members in Kotlin built-in classes which are NOT implemented as platform method calls,
+        // and thus break this assertion - e.g., 'Array<T>.iterator()' and similar functions.
+        // See KT-30908 for more details.
+        expression.dispatchReceiver = expression.dispatchReceiver?.transform(this, AssertionScope.Disabled)
 
-            if (isCallAssertionsDisabled || isValueArgumentForCallToMethodWithTypeCheckBarrier(expression)) {
-                for (i in 0 until expression.valueArgumentsCount) {
-                    getValueArgument(i)?.let { irArgument ->
-                        putValueArgument(i, irArgument.replaceImplicitNotNullWithArgument())
-                    }
-                }
+        val receiverAssertionScope = if (
+            isReceiverAssertionsDisabled ||
+            !isWithUnifiedNullChecks && expression.origin.isOperatorWithNoNullabilityAssertionsOnExtensionReceiver
+        ) AssertionScope.Disabled else AssertionScope.Enabled
+
+        expression.extensionReceiver = expression.extensionReceiver?.transform(this, receiverAssertionScope)
+
+        val parameterAssertionScope =
+            if (isCallToMethodWithTypeCheckBarrier(expression)) AssertionScope.Disabled else AssertionScope.Enabled
+        for (i in 0 until expression.valueArgumentsCount) {
+            expression.getValueArgument(i)?.let { irArgument ->
+                expression.putValueArgument(i, irArgument.transform(this, parameterAssertionScope))
             }
         }
+
+        return expression
     }
 
-    override fun visitContainerExpression(expression: IrContainerExpression) {
-        expression.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                val lastIndex = statements.lastIndex
-                if (lastIndex >= 0) {
-                    val lastStatement = statements[lastIndex]
-                    if (lastStatement is IrExpression) {
-                        statements[lastIndex] = lastStatement.replaceImplicitNotNullWithArgument()
-                    }
-                }
-            }
-        }
+    override fun visitGetField(expression: IrGetField, data: AssertionScope): IrExpression {
+        expression.receiver = expression.receiver?.transform(this, AssertionScope.Disabled)
+        return expression
     }
 
-    override fun visitReturn(expression: IrReturn) {
-        expression.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                value = value.replaceImplicitNotNullWithArgument()
-            }
-        }
+    override fun visitSetField(expression: IrSetField, data: AssertionScope): IrExpression {
+        expression.receiver = expression.receiver?.transform(this, AssertionScope.Disabled)
+        expression.value = expression.value.transform(this, data)
+        return expression
     }
 
-    override fun visitSetVariable(expression: IrSetVariable) {
-        expression.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                value = value.replaceImplicitNotNullWithArgument()
-            }
-        }
-    }
+    private fun isCallToMethodWithTypeCheckBarrier(expression: IrMemberAccessExpression): Boolean =
+        expression.symbol.owner.safeAs<IrSimpleFunction>()?.let { specialBridgeMethods.findSpecialWithOverride(it) != null } == true
 
-    override fun visitGetField(expression: IrGetField) {
-        expression.transformPostfix {
-            receiver = receiver?.replaceImplicitNotNullWithArgument()
-        }
-    }
-
-    override fun visitSetField(expression: IrSetField) {
-        expression.transformPostfix {
-            receiver = receiver?.replaceImplicitNotNullWithArgument()
-            if (isCallAssertionsDisabled) {
-                value = value.replaceImplicitNotNullWithArgument()
-            }
-        }
-    }
-
-    override fun visitVariable(declaration: IrVariable) {
-        declaration.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                initializer = initializer?.replaceImplicitNotNullWithArgument()
-            }
-        }
-    }
-
-    private fun IrExpressionBody.replaceImplicitNotNullWithArgument() {
-        expression = expression.replaceImplicitNotNullWithArgument()
-    }
-
-    override fun visitField(declaration: IrField) {
-        declaration.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                initializer?.replaceImplicitNotNullWithArgument()
-            }
-        }
-    }
-
-    override fun visitFunction(declaration: IrFunction) {
-        declaration.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                for (valueParameter in valueParameters) {
-                    valueParameter.defaultValue?.replaceImplicitNotNullWithArgument()
-                }
-            }
-        }
-    }
-
-    override fun visitWhen(expression: IrWhen) {
-        expression.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                for (irBranch in branches) {
-                    irBranch.condition = irBranch.condition.replaceImplicitNotNullWithArgument()
-                    irBranch.result = irBranch.result.replaceImplicitNotNullWithArgument()
-                }
-            }
-        }
-    }
-
-    override fun visitLoop(loop: IrLoop) {
-        loop.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                condition = condition.replaceImplicitNotNullWithArgument()
-            }
-        }
-    }
-
-    override fun visitThrow(expression: IrThrow) {
-        expression.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                value = value.replaceImplicitNotNullWithArgument()
-            }
-        }
-    }
-
-    override fun visitTry(aTry: IrTry) {
-        aTry.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                tryResult = tryResult.replaceImplicitNotNullWithArgument()
-                finallyExpression = finallyExpression?.replaceImplicitNotNullWithArgument()
-            }
-        }
-    }
-
-    override fun visitCatch(aCatch: IrCatch) {
-        aCatch.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                result = result.replaceImplicitNotNullWithArgument()
-            }
-        }
-    }
-
-    override fun visitTypeOperator(expression: IrTypeOperatorCall) {
-        expression.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                argument = argument.replaceImplicitNotNullWithArgument()
-            }
-        }
-    }
-
-    override fun visitVararg(expression: IrVararg) {
-        expression.transformPostfix {
-            if (isCallAssertionsDisabled) {
-                elements.forEachIndexed { index, irVarargElement ->
-                    when (irVarargElement) {
-                        is IrSpreadElement ->
-                            irVarargElement.expression = irVarargElement.expression.replaceImplicitNotNullWithArgument()
-                        is IrExpression ->
-                            putElement(index, irVarargElement.replaceImplicitNotNullWithArgument())
-                    }
-                }
-            }
-        }
-    }
-
-    private fun shouldDropNullabilityAssertionOnExtensionReceiver(expression: IrMemberAccessExpression): Boolean {
-        if (!isWithUnifiedNullChecks) {
-            if (expression.origin.isOperatorWithNoNullabilityAssertionsOnExtensionReceiver) return true
-        }
-
-        return false
-    }
+    private val IrStatementOrigin?.isOperatorWithNoNullabilityAssertionsOnExtensionReceiver
+        get() = this is IrStatementOrigin.COMPONENT_N || this in operatorsWithNoNullabilityAssertionsOnExtensionReceiver
 
     companion object {
-        private val operatorsWithNoNullabilityAssertionsOnExtensionReceiver =
-            hashSetOf(
-                IrStatementOrigin.PREFIX_INCR, IrStatementOrigin.POSTFIX_INCR,
-                IrStatementOrigin.PREFIX_DECR, IrStatementOrigin.POSTFIX_DECR
-            )
-
-        internal val IrStatementOrigin?.isOperatorWithNoNullabilityAssertionsOnExtensionReceiver
-            get() =
-                this is IrStatementOrigin.COMPONENT_N ||
-                        this in operatorsWithNoNullabilityAssertionsOnExtensionReceiver
-
-        internal fun IrExpression.replaceImplicitNotNullWithArgument(): IrExpression =
-            if (this is IrTypeOperatorCall && this.operator == IrTypeOperator.IMPLICIT_NOTNULL)
-                argument
-            else
-                this
-
-        internal fun isValueArgumentForCallToMethodWithTypeCheckBarrier(expression: IrMemberAccessExpression): Boolean {
-            val descriptor = expression.symbol.descriptor
-            if (descriptor !is CallableMemberDescriptor) return false
-
-            val specialSignatureInfo = with(BuiltinMethodsWithSpecialGenericSignature) {
-                descriptor.getSpecialSignatureInfo()
-            }
-            return specialSignatureInfo?.isObjectReplacedWithTypeParameter ?: false
-        }
+        private val operatorsWithNoNullabilityAssertionsOnExtensionReceiver = hashSetOf(
+            IrStatementOrigin.PREFIX_INCR, IrStatementOrigin.POSTFIX_INCR,
+            IrStatementOrigin.PREFIX_DECR, IrStatementOrigin.POSTFIX_DECR
+        )
     }
-}
+}

compiler/testData/codegen/boxAgainstJava/notNullAssertions/noCallAssertions.kt

@@ -1,5 +1,4 @@
 // KOTLIN_CONFIGURATION_FLAGS: +JVM.DISABLE_PARAM_ASSERTIONS, +JVM.DISABLE_CALL_ASSERTIONS
-// IGNORE_BACKEND: JVM_IR
 // FILE: noCallAssertions.kt
 
 class AssertionChecker(val nullPointerExceptionExpected: Boolean) {

compiler/testData/codegen/bytecodeText/nullCheckOptimization/nullabilityAssertionOnDispatchReceiver.kt

@@ -19,3 +19,7 @@ public class J {
 
     public static J j() { return null; }
 }
+
+// @TestKt.class:
+// 0 checkNotNullExpressionValue
+// 0 checkExpressionValueIsNotNull

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -14716,11 +14716,6 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nonNullableTypeParameter.kt");
             }
 
-            @TestMetadata("nullabilityAssertionOnDispatchReceiver.kt")
-            public void testNullabilityAssertionOnDispatchReceiver() throws Exception {
-                runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nullabilityAssertionOnDispatchReceiver.kt");
-            }
-
             @TestMetadata("nullabilityAssertionOnExtensionReceiver.kt")
             public void testNullabilityAssertionOnExtensionReceiver() throws Exception {
                 runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nullabilityAssertionOnExtensionReceiver.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java

@@ -3568,6 +3568,11 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/nullCheckOptimization/nullCheckAfterExclExcl_1_4.kt");
         }
 
+        @TestMetadata("nullabilityAssertionOnDispatchReceiver.kt")
+        public void testNullabilityAssertionOnDispatchReceiver() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/nullCheckOptimization/nullabilityAssertionOnDispatchReceiver.kt");
+        }
+
         @TestMetadata("primitiveCheck.kt")
         public void testPrimitiveCheck() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/nullCheckOptimization/primitiveCheck.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -14658,6 +14658,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
         @TestDataPath("$PROJECT_ROOT")
         @RunWith(JUnit3RunnerWithInners.class)
         public static class NotNullAssertions extends AbstractLightAnalysisModeTest {
+            @TestMetadata("paramAssertionMessage.kt")
+            public void ignoreParamAssertionMessage() throws Exception {
+                runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/paramAssertionMessage.kt");
+            }
+
             private void runTest(String testDataFilePath) throws Exception {
                 KotlinTestUtils.runTest(this::doTest, TargetBackend.JVM, testDataFilePath);
             }
@@ -14716,11 +14721,6 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nonNullableTypeParameter.kt");
             }
 
-            @TestMetadata("nullabilityAssertionOnDispatchReceiver.kt")
-            public void testNullabilityAssertionOnDispatchReceiver() throws Exception {
-                runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nullabilityAssertionOnDispatchReceiver.kt");
-            }
-
             @TestMetadata("nullabilityAssertionOnExtensionReceiver.kt")
             public void testNullabilityAssertionOnExtensionReceiver() throws Exception {
                 runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nullabilityAssertionOnExtensionReceiver.kt");
@@ -14746,11 +14746,6 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nullableTypeParameter.kt");
             }
 
-            @TestMetadata("paramAssertionMessage.kt")
-            public void testParamAssertionMessage() throws Exception {
-                runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/paramAssertionMessage.kt");
-            }
-
             @TestMetadata("staticCallErrorMessage.kt")
             public void testStaticCallErrorMessage() throws Exception {
                 runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/staticCallErrorMessage.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/FirBlackBoxCodegenTestGenerated.java

@@ -13591,11 +13591,6 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nonNullableTypeParameter.kt");
             }
 
-            @TestMetadata("nullabilityAssertionOnDispatchReceiver.kt")
-            public void testNullabilityAssertionOnDispatchReceiver() throws Exception {
-                runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nullabilityAssertionOnDispatchReceiver.kt");
-            }
-
             @TestMetadata("nullabilityAssertionOnExtensionReceiver.kt")
             public void testNullabilityAssertionOnExtensionReceiver() throws Exception {
                 runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nullabilityAssertionOnExtensionReceiver.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -13591,11 +13591,6 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nonNullableTypeParameter.kt");
             }
 
-            @TestMetadata("nullabilityAssertionOnDispatchReceiver.kt")
-            public void testNullabilityAssertionOnDispatchReceiver() throws Exception {
-                runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nullabilityAssertionOnDispatchReceiver.kt");
-            }
-
             @TestMetadata("nullabilityAssertionOnExtensionReceiver.kt")
             public void testNullabilityAssertionOnExtensionReceiver() throws Exception {
                 runTest("compiler/testData/codegen/box/javaInterop/notNullAssertions/nullabilityAssertionOnExtensionReceiver.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBytecodeTextTestGenerated.java

@@ -3613,6 +3613,11 @@ public class IrBytecodeTextTestGenerated extends AbstractIrBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/nullCheckOptimization/nullCheckAfterExclExcl_1_4.kt");
         }
 
+        @TestMetadata("nullabilityAssertionOnDispatchReceiver.kt")
+        public void testNullabilityAssertionOnDispatchReceiver() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/nullCheckOptimization/nullabilityAssertionOnDispatchReceiver.kt");
+        }
+
         @TestMetadata("primitiveCheck.kt")
         public void testPrimitiveCheck() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/nullCheckOptimization/primitiveCheck.kt");