JVM_IR handle @EnhancedNullability types in null check simplification

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -39746,6 +39746,18 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             runTest("compiler/testData/codegen/box/safeCall/safeCallOnLong.kt");
         }
 
+        @Test
+        @TestMetadata("safeCallSimplificationEnhancedNullabilityType.kt")
+        public void testSafeCallSimplificationEnhancedNullabilityType() throws Exception {
+            runTest("compiler/testData/codegen/box/safeCall/safeCallSimplificationEnhancedNullabilityType.kt");
+        }
+
+        @Test
+        @TestMetadata("safeCallSimplificationFlexibleType.kt")
+        public void testSafeCallSimplificationFlexibleType() throws Exception {
+            runTest("compiler/testData/codegen/box/safeCall/safeCallSimplificationFlexibleType.kt");
+        }
+
         @Test
         @TestMetadata("safeCallWithElvisFolding.kt")
         public void testSafeCallWithElvisFolding() throws Exception {

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

@@ -8,18 +8,18 @@ import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredStatementOrigin
-import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
-import org.jetbrains.kotlin.ir.types.IrType
-import org.jetbrains.kotlin.ir.types.isNullable
+import org.jetbrains.kotlin.ir.types.*
 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.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 
 
 val jvmSafeCallFoldingPhase = makeIrFilePhase(
@@ -119,14 +119,17 @@ class JvmSafeCallChainFoldingLowering(val context: JvmBackendContext) : FileLowe
         IrConstImpl.boolean(startOffset, endOffset, context.irBuiltIns.booleanType, false)
 
     private fun irValNotNull(startOffset: Int, endOffset: Int, irVariable: IrVariable): IrExpression =
-        if (irVariable.type.isNullable())
+        if (irVariable.type.isJvmNullable())
             IrGetValueImpl(startOffset, endOffset, irVariable.symbol).irEqEqNull().irNot()
         else
             irTrue(startOffset, endOffset)
 
+    private fun IrType.isJvmNullable(): Boolean =
+        isNullable() || hasAnnotation(JvmAnnotationNames.ENHANCED_NULLABILITY_ANNOTATION)
+    
     private fun IrType.isJvmPrimitive(): Boolean =
-        // TODO get rid of type mapper (take care of '@EnhancedNullability', maybe some other stuff).
-        AsmUtil.isPrimitive(context.typeMapper.mapType(this))
+        (isBoolean() || isByte() || isShort() || isInt() || isLong() || isChar() || isFloat() || isDouble()) &&
+                !hasAnnotation(JvmAnnotationNames.ENHANCED_NULLABILITY_ANNOTATION)
 
     private inner class Transformer : IrElementTransformerVoid() {
         override fun visitBlock(expression: IrBlock): IrExpression {
@@ -393,7 +396,7 @@ class JvmSafeCallChainFoldingLowering(val context: JvmBackendContext) : FileLowe
             if (this.origin != JvmLoweredStatementOrigin.FOLDED_SAFE_CALL) return false
             val innerWhen = this.statements[0] as? IrWhen ?: return false
             val safeCallResult = innerWhen.branches[0].result
-            return !safeCallResult.type.isNullable()
+            return !safeCallResult.type.isJvmNullable()
         }
 
         override fun visitCall(expression: IrCall): IrExpression {

compiler/testData/codegen/box/safeCall/safeCallSimplificationEnhancedNullabilityType.kt

@@ -0,0 +1,16 @@
+// TARGET_BACKEND: JVM
+// FILE: safeCallSimplificationEnhancedNullabilityType.kt
+
+fun String.zap() = "failed"
+
+fun box() =
+    J.nullString()?.zap()
+        ?: "OK"
+
+// FILE: J.java
+import org.jetbrains.annotations.NotNull;
+
+public class J {
+    @NotNull
+    public static String nullString() { return null; }
+}

compiler/testData/codegen/box/safeCall/safeCallSimplificationFlexibleType.kt

@@ -0,0 +1,13 @@
+// TARGET_BACKEND: JVM
+// FILE: safeCallSimplificationFlexibleType.kt
+
+fun String.zap() = "failed"
+
+fun box() =
+    J.nullString()?.zap()
+        ?: "OK"
+
+// FILE: J.java
+public class J {
+    public static String nullString() { return null; }
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -39590,6 +39590,18 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/safeCall/safeCallOnLong.kt");
         }
 
+        @Test
+        @TestMetadata("safeCallSimplificationEnhancedNullabilityType.kt")
+        public void testSafeCallSimplificationEnhancedNullabilityType() throws Exception {
+            runTest("compiler/testData/codegen/box/safeCall/safeCallSimplificationEnhancedNullabilityType.kt");
+        }
+
+        @Test
+        @TestMetadata("safeCallSimplificationFlexibleType.kt")
+        public void testSafeCallSimplificationFlexibleType() throws Exception {
+            runTest("compiler/testData/codegen/box/safeCall/safeCallSimplificationFlexibleType.kt");
+        }
+
         @Test
         @TestMetadata("safeCallWithElvisFolding.kt")
         public void testSafeCallWithElvisFolding() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -39746,6 +39746,18 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/safeCall/safeCallOnLong.kt");
         }
 
+        @Test
+        @TestMetadata("safeCallSimplificationEnhancedNullabilityType.kt")
+        public void testSafeCallSimplificationEnhancedNullabilityType() throws Exception {
+            runTest("compiler/testData/codegen/box/safeCall/safeCallSimplificationEnhancedNullabilityType.kt");
+        }
+
+        @Test
+        @TestMetadata("safeCallSimplificationFlexibleType.kt")
+        public void testSafeCallSimplificationFlexibleType() throws Exception {
+            runTest("compiler/testData/codegen/box/safeCall/safeCallSimplificationFlexibleType.kt");
+        }
+
         @Test
         @TestMetadata("safeCallWithElvisFolding.kt")
         public void testSafeCallWithElvisFolding() throws Exception {

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

@@ -31688,6 +31688,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/safeCall/safeCallOnLong.kt");
         }
 
+        @TestMetadata("safeCallSimplificationEnhancedNullabilityType.kt")
+        public void testSafeCallSimplificationEnhancedNullabilityType() throws Exception {
+            runTest("compiler/testData/codegen/box/safeCall/safeCallSimplificationEnhancedNullabilityType.kt");
+        }
+
+        @TestMetadata("safeCallSimplificationFlexibleType.kt")
+        public void testSafeCallSimplificationFlexibleType() throws Exception {
+            runTest("compiler/testData/codegen/box/safeCall/safeCallSimplificationFlexibleType.kt");
+        }
+
         @TestMetadata("safeCallWithElvisFolding.kt")
         public void testSafeCallWithElvisFolding() throws Exception {
             runTest("compiler/testData/codegen/box/safeCall/safeCallWithElvisFolding.kt");