[Jspecify] Do preserving flexibility when adding equality constraints by taking into account type with enhancement

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt

@@ -171,18 +171,26 @@ internal object CreateFreshVariablesSubstitutor : ResolutionPart() {
         }
     }
 
-    fun TypeParameterDescriptor.shouldBeFlexible(): Boolean {
+    fun TypeParameterDescriptor.shouldBeFlexible(flexibleCheck: (KotlinType) -> Boolean = { it.isFlexible() }): Boolean {
         return upperBounds.any {
-            it.isFlexible() || ((it.constructor.declarationDescriptor as? TypeParameterDescriptor)?.run { shouldBeFlexible() } ?: false)
+            flexibleCheck(it) || ((it.constructor.declarationDescriptor as? TypeParameterDescriptor)?.run { shouldBeFlexible() } ?: false)
         }
     }
 
     private fun getTypePreservingFlexibilityWrtTypeVariable(
         type: KotlinType,
         typeVariable: TypeVariableFromCallableDescriptor
-    ) = if (typeVariable.originalTypeParameter.shouldBeFlexible()) {
+    ): KotlinType {
-        KotlinTypeFactory.flexibleType(type.makeNotNullable().lowerIfFlexible(), type.makeNullable().upperIfFlexible())
+        fun createFlexibleType() =
-    } else type
+            KotlinTypeFactory.flexibleType(type.makeNotNullable().lowerIfFlexible(), type.makeNullable().upperIfFlexible())
+
+        return when {
+            typeVariable.originalTypeParameter.shouldBeFlexible { it is FlexibleTypeWithEnhancement } ->
+                createFlexibleType().wrapEnhancement(type)
+            typeVariable.originalTypeParameter.shouldBeFlexible() -> createFlexibleType()
+            else -> type
+        }
+    }
 
     private fun createKnownParametersFromFreshVariablesSubstitutor(
         freshVariableSubstitutor: FreshVariableNewTypeSubstitutor,

compiler/testData/diagnostics/foreignAnnotationsTests/java8Tests/jspecify/warnMode/NonPlatformTypeParameter.fir.kt

@@ -14,14 +14,14 @@ public class Test {}
 fun <T : Test> main(a1: NonPlatformTypeParameter<Any?>, a2: NonPlatformTypeParameter<Test>, x: T): Unit {
     a1.foo(null)
     a1.bar<Test?>(null)
-    // jspecify_nullness_mismatch{mute}
+    // jspecify_nullness_mismatch
-    a1.bar<T>(null)
+    a1.bar<T>(<!NULLABILITY_MISMATCH_BASED_ON_JAVA_ANNOTATIONS!>null<!>)
     a1.bar<T>(x)
 
     // jspecify_nullness_mismatch{mute}
     a2.foo(null)
     a2.bar<Test?>(null)
-    // jspecify_nullness_mismatch{mute}
+    // jspecify_nullness_mismatch
-    a2.bar<T>(null)
+    a2.bar<T>(<!NULLABILITY_MISMATCH_BASED_ON_JAVA_ANNOTATIONS!>null<!>)
     a2.bar<T>(x)
 }

compiler/testData/diagnostics/foreignAnnotationsTests/java8Tests/jspecify/warnMode/NonPlatformTypeParameter.kt

@@ -14,14 +14,14 @@ public class Test {}
 fun <T : Test> main(a1: NonPlatformTypeParameter<Any?>, a2: NonPlatformTypeParameter<Test>, x: T): Unit {
     a1.foo(null)
     a1.bar<Test?>(null)
-    // jspecify_nullness_mismatch{mute}
+    // jspecify_nullness_mismatch
-    a1.bar<T>(null)
+    a1.bar<T>(<!NULLABILITY_MISMATCH_BASED_ON_JAVA_ANNOTATIONS!>null<!>)
     a1.bar<T>(x)
 
     // jspecify_nullness_mismatch{mute}
     a2.foo(null)
     a2.bar<Test?>(null)
-    // jspecify_nullness_mismatch{mute}
+    // jspecify_nullness_mismatch
-    a2.bar<T>(null)
+    a2.bar<T>(<!NULLABILITY_MISMATCH_BASED_ON_JAVA_ANNOTATIONS!>null<!>)
     a2.bar<T>(x)
 }