Support @EnhancedNullability in psi2ir

compiler/fir/fir2ir/tests/org/jetbrains/kotlin/fir/Fir2IrTextTestGenerated.java

@@ -1670,6 +1670,11 @@ public class Fir2IrTextTestGenerated extends AbstractFir2IrTextTest {
             runTest("compiler/testData/ir/irText/types/asOnPlatformType.kt");
         }
 
+        @TestMetadata("enhancedNullability.kt")
+        public void testEnhancedNullability() throws Exception {
+            runTest("compiler/testData/ir/irText/types/enhancedNullability.kt");
+        }
+
         @TestMetadata("explicitEqualsAndCompareToCallsOnPlatformTypeReceiver.kt")
         public void testExplicitEqualsAndCompareToCallsOnPlatformTypeReceiver() throws Exception {
             runTest("compiler/testData/ir/irText/types/explicitEqualsAndCompareToCallsOnPlatformTypeReceiver.kt");

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmGeneratorExtensions.kt

@@ -10,14 +10,17 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.Visibility
+import org.jetbrains.kotlin.descriptors.annotations.FilteredAnnotations
 import org.jetbrains.kotlin.ir.builders.declarations.buildClass
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 import org.jetbrains.kotlin.load.java.descriptors.JavaCallableMemberDescriptor
 import org.jetbrains.kotlin.load.java.descriptors.JavaClassDescriptor
 import org.jetbrains.kotlin.load.java.sam.SamAdapterDescriptor
 import org.jetbrains.kotlin.load.java.sam.SamConstructorDescriptor
 import org.jetbrains.kotlin.load.java.sam.SingleAbstractMethodUtils
+import org.jetbrains.kotlin.load.java.typeEnhancement.hasEnhancedNullability
 import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
 import org.jetbrains.kotlin.psi2ir.generators.GeneratorExtensions
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
@@ -27,6 +30,7 @@ import org.jetbrains.kotlin.synthetic.SamAdapterExtensionFunctionDescriptor
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeSubstitutor
 import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.types.typeUtil.replaceAnnotations
 
 class JvmGeneratorExtensions(private val generateFacades: Boolean = true) : GeneratorExtensions() {
     val classNameOverride = mutableMapOf<IrClass, JvmClassName>()
@@ -88,4 +92,24 @@ class JvmGeneratorExtensions(private val generateFacades: Boolean = true) : Gene
 
     override fun isPropertyWithPlatformField(descriptor: PropertyDescriptor): Boolean =
         descriptor.hasJvmFieldAnnotation()
+
+    override val enhancedNullability: EnhancedNullability
+        get() = JvmEnhancedNullability
+
+    open class JvmEnhancedNullability : EnhancedNullability() {
+        override fun hasEnhancedNullability(kotlinType: KotlinType): Boolean =
+            kotlinType.hasEnhancedNullability()
+
+        override fun stripEnhancedNullability(kotlinType: KotlinType): KotlinType =
+            if (kotlinType.hasEnhancedNullability())
+                kotlinType.replaceAnnotations(
+                    FilteredAnnotations(kotlinType.annotations, true) {
+                        it != JvmAnnotationNames.ENHANCED_NULLABILITY_ANNOTATION
+                    }
+                )
+            else
+                kotlinType
+
+        companion object Instance : JvmEnhancedNullability()
+    }
 }

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/GeneratorExtensions.kt

@@ -30,4 +30,16 @@ open class GeneratorExtensions : StubGeneratorExtensions() {
     }
 
     open fun computeFieldVisibility(descriptor: PropertyDescriptor): Visibility? = null
+
+    open val enhancedNullability: EnhancedNullability
+        get() = EnhancedNullability
+
+    open class EnhancedNullability {
+        open fun hasEnhancedNullability(kotlinType: KotlinType): Boolean = false
+
+        open fun stripEnhancedNullability(kotlinType: KotlinType): KotlinType = kotlinType
+
+        companion object Instance : EnhancedNullability()
+    }
+
 }

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/transformations/InsertImplicitCasts.kt

@@ -50,7 +50,7 @@ fun insertImplicitCasts(element: IrElement, context: GeneratorContext) {
             context.irBuiltIns,
             context.typeTranslator,
             context.callToSubstitutedDescriptorMap,
-            context.extensions.samConversion
+            context.extensions
         ),
         null
     )
@@ -61,7 +61,7 @@ open class InsertImplicitCasts(
     private val irBuiltIns: IrBuiltIns,
     private val typeTranslator: TypeTranslator,
     private val callToSubstitutedDescriptorMap: Map<IrMemberAccessExpression, CallableDescriptor>,
-    private val samConversion: GeneratorExtensions.SamConversion
+    private val generatorExtensions: GeneratorExtensions
 ) : IrElementTransformerVoid() {
 
     private fun KotlinType.toIrType() = typeTranslator.translateType(this)
@@ -191,10 +191,13 @@ open class InsertImplicitCasts(
             finallyExpression = finallyExpression?.coerceToUnit()
         }
 
+    private fun KotlinType.getSubstitutedFunctionTypeForSamType() =
+        generatorExtensions.samConversion.getSubstitutedFunctionTypeForSamType(this)
+
     override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression =
         when (expression.operator) {
             IrTypeOperator.SAM_CONVERSION -> expression.transformPostfix {
-                argument = argument.cast(samConversion.getSubstitutedFunctionTypeForSamType(typeOperand.originalKotlinType!!))
+                argument = argument.cast(typeOperand.originalKotlinType!!.getSubstitutedFunctionTypeForSamType())
             }
 
             IrTypeOperator.IMPLICIT_CAST -> {
@@ -251,7 +254,8 @@ open class InsertImplicitCasts(
                 else
                     implicitCast(expectedType, IrTypeOperator.IMPLICIT_DYNAMIC_CAST)
 
-            valueType.isNullabilityFlexible() && valueType.containsNull() && !expectedType.containsNull() ->
+            (valueType.isNullabilityFlexible() && valueType.containsNull() || valueType.hasEnhancedNullability()) &&
+                    !expectedType.containsNull() ->
                 implicitNonNull(valueType, expectedType)
 
             KotlinTypeChecker.DEFAULT.isSubtypeOf(valueType, expectedType.makeNullable()) ->
@@ -267,8 +271,12 @@ open class InsertImplicitCasts(
         }
     }
 
+    private fun KotlinType.hasEnhancedNullability() =
+        generatorExtensions.enhancedNullability.hasEnhancedNullability(this)
+
     private fun IrExpression.implicitNonNull(valueType: KotlinType, expectedType: KotlinType): IrExpression {
-        val nonNullValueType = valueType.upperIfFlexible().makeNotNullable()
+        val nonNullFlexibleType = valueType.upperIfFlexible().makeNotNullable()
+        val nonNullValueType = generatorExtensions.enhancedNullability.stripEnhancedNullability(nonNullFlexibleType)
         return implicitCast(nonNullValueType, IrTypeOperator.IMPLICIT_NOTNULL).cast(expectedType)
     }
 

compiler/testData/codegen/box/nullCheckOptimization/javaNullCheckThrowsNpe_1_4.kt

@@ -1,7 +1,6 @@
 // !API_VERSION: LATEST
 // IGNORE_BACKEND_FIR: JVM_IR
 // TARGET_BACKEND: JVM
-// IGNORE_BACKEND: JVM_IR
 // FILE: A.java
 
 import org.jetbrains.annotations.NotNull;

compiler/testData/ir/irText/types/enhancedNullability.fir.txt

@@ -0,0 +1,31 @@
+FILE fqName:<root> fileName:/enhancedNullability.kt
+  FUN name:use visibility:public modality:FINAL <> (s:kotlin.String) returnType:kotlin.Unit
+    VALUE_PARAMETER name:s index:0 type:kotlin.String
+    BLOCK_BODY
+  FUN name:testUse visibility:public modality:FINAL <> () returnType:kotlin.Unit
+    BLOCK_BODY
+      CALL 'public final fun use (s: kotlin.String): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+        s: CALL 'public open fun s (): kotlin.String [operator] declared in <root>.J' type=kotlin.String origin=null
+  FUN name:testLocalVal visibility:public modality:FINAL <> () returnType:kotlin.Unit
+    BLOCK_BODY
+      VAR name:local type:kotlin.String [val]
+        CALL 'public open fun s (): kotlin.String [operator] declared in <root>.J' type=kotlin.String origin=null
+  FUN name:testReturnValue visibility:public modality:FINAL <> () returnType:kotlin.String
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun testReturnValue (): kotlin.String declared in <root>'
+        CALL 'public open fun s (): kotlin.String [operator] declared in <root>.J' type=kotlin.String origin=null
+  PROPERTY name:testGlobalVal visibility:public modality:FINAL [val]
+    FIELD PROPERTY_BACKING_FIELD name:testGlobalVal type:kotlin.String visibility:private [final,static]
+      EXPRESSION_BODY
+        CALL 'public open fun s (): kotlin.String [operator] declared in <root>.J' type=kotlin.String origin=null
+    FUN DEFAULT_PROPERTY_ACCESSOR name:<get-testGlobalVal> visibility:public modality:FINAL <> () returnType:kotlin.String
+      correspondingProperty: PROPERTY name:testGlobalVal visibility:public modality:FINAL [val]
+      BLOCK_BODY
+        RETURN type=kotlin.Nothing from='public final fun <get-testGlobalVal> (): kotlin.String declared in <root>'
+          GET_FIELD 'FIELD PROPERTY_BACKING_FIELD name:testGlobalVal type:kotlin.String visibility:private [final,static]' type=kotlin.String origin=null
+  PROPERTY name:testGlobalValGetter visibility:public modality:FINAL [val]
+    FUN name:<get-testGlobalValGetter> visibility:public modality:FINAL <> () returnType:kotlin.String
+      correspondingProperty: PROPERTY name:testGlobalValGetter visibility:public modality:FINAL [val]
+      BLOCK_BODY
+        RETURN type=kotlin.Nothing from='public final fun <get-testGlobalValGetter> (): kotlin.String declared in <root>'
+          CALL 'public open fun s (): kotlin.String [operator] declared in <root>.J' type=kotlin.String origin=null

compiler/testData/ir/irText/types/enhancedNullability.kt

@@ -0,0 +1,23 @@
+// FILE: enhancedNullability.kt
+fun use(s: String) {}
+
+fun testUse() {
+    use(J.s())
+}
+
+fun testLocalVal() {
+    val local = J.s()
+}
+
+fun testReturnValue() = J.s()
+
+val testGlobalVal = J.s()
+
+val testGlobalValGetter get() = J.s()
+
+// FILE: J.java
+import org.jetbrains.annotations.*
+
+public class J {
+    public static @NotNull String s() { return null; }
+}

compiler/testData/ir/irText/types/enhancedNullability.txt

@@ -0,0 +1,36 @@
+FILE fqName:<root> fileName:/enhancedNullability.kt
+  FUN name:use visibility:public modality:FINAL <> (s:kotlin.String) returnType:kotlin.Unit
+    VALUE_PARAMETER name:s index:0 type:kotlin.String
+    BLOCK_BODY
+  FUN name:testUse visibility:public modality:FINAL <> () returnType:kotlin.Unit
+    BLOCK_BODY
+      CALL 'public final fun use (s: kotlin.String): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+        s: TYPE_OP type=kotlin.String origin=IMPLICIT_NOTNULL typeOperand=kotlin.String
+          CALL 'public open fun s (): kotlin.String declared in <root>.J' type=kotlin.String origin=null
+  FUN name:testLocalVal visibility:public modality:FINAL <> () returnType:kotlin.Unit
+    BLOCK_BODY
+      VAR name:local type:kotlin.String [val]
+        TYPE_OP type=kotlin.String origin=IMPLICIT_NOTNULL typeOperand=kotlin.String
+          CALL 'public open fun s (): kotlin.String declared in <root>.J' type=kotlin.String origin=null
+  FUN name:testReturnValue visibility:public modality:FINAL <> () returnType:kotlin.String
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun testReturnValue (): kotlin.String declared in <root>'
+        TYPE_OP type=kotlin.String origin=IMPLICIT_NOTNULL typeOperand=kotlin.String
+          CALL 'public open fun s (): kotlin.String declared in <root>.J' type=kotlin.String origin=null
+  PROPERTY name:testGlobalVal visibility:public modality:FINAL [val]
+    FIELD PROPERTY_BACKING_FIELD name:testGlobalVal type:kotlin.String visibility:private [final,static]
+      EXPRESSION_BODY
+        TYPE_OP type=kotlin.String origin=IMPLICIT_NOTNULL typeOperand=kotlin.String
+          CALL 'public open fun s (): kotlin.String declared in <root>.J' type=kotlin.String origin=null
+    FUN DEFAULT_PROPERTY_ACCESSOR name:<get-testGlobalVal> visibility:public modality:FINAL <> () returnType:kotlin.String
+      correspondingProperty: PROPERTY name:testGlobalVal visibility:public modality:FINAL [val]
+      BLOCK_BODY
+        RETURN type=kotlin.Nothing from='public final fun <get-testGlobalVal> (): kotlin.String declared in <root>'
+          GET_FIELD 'FIELD PROPERTY_BACKING_FIELD name:testGlobalVal type:kotlin.String visibility:private [final,static]' type=kotlin.String origin=null
+  PROPERTY name:testGlobalValGetter visibility:public modality:FINAL [val]
+    FUN name:<get-testGlobalValGetter> visibility:public modality:FINAL <> () returnType:kotlin.String
+      correspondingProperty: PROPERTY name:testGlobalValGetter visibility:public modality:FINAL [val]
+      BLOCK_BODY
+        RETURN type=kotlin.Nothing from='public final fun <get-testGlobalValGetter> (): kotlin.String declared in <root>'
+          TYPE_OP type=kotlin.String origin=IMPLICIT_NOTNULL typeOperand=kotlin.String
+            CALL 'public open fun s (): kotlin.String declared in <root>.J' type=kotlin.String origin=null

compiler/tests/org/jetbrains/kotlin/ir/IrTextTestCaseGenerated.java

@@ -1670,6 +1670,11 @@ public class IrTextTestCaseGenerated extends AbstractIrTextTestCase {
             runTest("compiler/testData/ir/irText/types/asOnPlatformType.kt");
         }
 
+        @TestMetadata("enhancedNullability.kt")
+        public void testEnhancedNullability() throws Exception {
+            runTest("compiler/testData/ir/irText/types/enhancedNullability.kt");
+        }
+
         @TestMetadata("explicitEqualsAndCompareToCallsOnPlatformTypeReceiver.kt")
         public void testExplicitEqualsAndCompareToCallsOnPlatformTypeReceiver() throws Exception {
             runTest("compiler/testData/ir/irText/types/explicitEqualsAndCompareToCallsOnPlatformTypeReceiver.kt");