Support @EnhancedNullability in psi2ir

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)
     }