FIR Java: record Java types with flexible nullability

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/JavaUtils.kt

@@ -197,7 +197,11 @@ private fun ClassId.toConeFlexibleType(
     typeArgumentsForUpper: Array<ConeTypeProjection> = typeArguments,
     attributes: ConeAttributes = ConeAttributes.Empty
 ) = ConeFlexibleType(
-    toConeKotlinType(typeArguments, isNullable = false, attributes),
+    toConeKotlinType(
+        typeArguments,
+        isNullable = false,
+        attributes.intersectUnless(ConeAttributes.WithFlexibleNullability) { it.hasEnhancedNullability }
+    ),
     toConeKotlinType(typeArgumentsForUpper, isNullable = true, attributes)
 )
 
@@ -233,7 +237,15 @@ private fun JavaClassifierType.toConeKotlinTypeWithoutEnhancement(
             attributes = attributes
         )
 
-    return if (isRaw) ConeRawType(lowerBound, upperBound) else ConeFlexibleType(lowerBound, upperBound)
+    return if (isRaw)
+        ConeRawType(lowerBound, upperBound)
+    else
+        ConeFlexibleType(
+            lowerBound.withAttributes(
+                lowerBound.attributes.intersectUnless(ConeAttributes.WithFlexibleNullability) { it.hasEnhancedNullability }
+            ),
+            upperBound
+        )
 }
 
 private fun computeRawProjection(
@@ -303,7 +315,9 @@ private fun getErasedVersionOfFirstUpperBound(
             if (firstUpperBound.upperBound is ConeTypeParameterType) {
                 // Avoid exponential complexity
                 ConeFlexibleType(
-                    lowerBound,
+                    lowerBound.withAttributes(
+                        lowerBound.attributes.intersectUnless(ConeAttributes.WithFlexibleNullability) { it.hasEnhancedNullability }
+                    ),
                     lowerBound.withNullability(ConeNullability.NULLABLE)
                 )
             } else {

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/enhancement/javaTypeUtils.kt

@@ -88,22 +88,29 @@ private fun ConeKotlinType.enhanceConeKotlinType(
 ): ConeKotlinType {
     return when (this) {
         is ConeFlexibleType -> {
+            val needsFlexibleNullabilityAttribute = lowerBound.nullability != upperBound.nullability && !lowerBound.hasEnhancedNullability
             val lowerResult = lowerBound.enhanceInflexibleType(
-                session, TypeComponentPosition.FLEXIBLE_LOWER, qualifiers, index
+                session, TypeComponentPosition.FLEXIBLE_LOWER, qualifiers, index,
+                attributes = if (needsFlexibleNullabilityAttribute)
+                    lowerBound.attributes.intersect(ConeAttributes.WithFlexibleNullability)
+                else
+                    lowerBound.attributes
             )
             val upperResult = upperBound.enhanceInflexibleType(
-                session, TypeComponentPosition.FLEXIBLE_UPPER, qualifiers, index
+                session, TypeComponentPosition.FLEXIBLE_UPPER, qualifiers, index, upperBound.attributes
             )
 
             when {
-                lowerResult === lowerBound && upperResult === upperBound -> this
+                !needsFlexibleNullabilityAttribute && lowerResult === lowerBound && upperResult === upperBound -> this
                 this is ConeRawType -> ConeRawType(lowerResult, upperResult)
                 else -> coneFlexibleOrSimpleType(
                     session, lowerResult, upperResult, isNotNullTypeParameter = qualifiers(index).isNotNullTypeParameter
                 )
             }
         }
-        is ConeSimpleKotlinType -> enhanceInflexibleType(session, TypeComponentPosition.INFLEXIBLE, qualifiers, index)
+        is ConeSimpleKotlinType -> enhanceInflexibleType(
+            session, TypeComponentPosition.INFLEXIBLE, qualifiers, index, attributes
+        )
         else -> this
     }
 }
@@ -158,7 +165,8 @@ private fun ConeKotlinType.enhanceInflexibleType(
     session: FirSession,
     position: TypeComponentPosition,
     qualifiers: IndexedJavaTypeQualifiers,
-    index: Int
+    index: Int,
+    attributes: ConeAttributes = this.attributes
 ): ConeKotlinType {
     require(this !is ConeFlexibleType) {
         "$this should not be flexible"
@@ -198,7 +206,7 @@ private fun ConeKotlinType.enhanceInflexibleType(
 
     if (!wereChangesInArgs && originalTag == enhancedTag && enhancedNullability == isNullable) return this
 
-    val enhancedType = enhancedTag.constructType(enhancedArguments, enhancedNullability)
+    val enhancedType = enhancedTag.constructType(enhancedArguments, enhancedNullability, attributes)
 
     // TODO: why all of these is needed
 //    val enhancement = if (effectiveQualifiers.isNotNullTypeParameter) NotNullTypeParameter(enhancedType) else enhancedType

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/scopes/JavaClassMembersEnhancementScope.kt

@@ -104,7 +104,14 @@ class JavaClassMembersEnhancementScope(
                                 emptyArray(), valueParameter.returnTypeRef.isMarkedNullable == true
                             )
                             if (valueParameter.returnTypeRef.coneType is ConeFlexibleType) {
-                                ConeFlexibleType(type, type.withNullability(ConeNullability.NULLABLE))
+                                ConeFlexibleType(
+                                    type.withAttributes(
+                                        type.attributes.intersectUnless(ConeAttributes.WithFlexibleNullability) {
+                                            it.hasEnhancedNullability
+                                        }
+                                    ),
+                                    type.withNullability(ConeNullability.NULLABLE)
+                                )
                             } else {
                                 type
                             }
@@ -129,8 +136,8 @@ class JavaClassMembersEnhancementScope(
                 newParameterTypes = valueParameters.zip(newParameterTypes).map { (valueParameter, newType) ->
                     newType ?: valueParameter.returnTypeRef.coneType
                 },
-                    newDispatchReceiverType = dispatchReceiverType,
-                )
+                newDispatchReceiverType = dispatchReceiverType,
+            )
 
         }
         return this