From 24796557083040d3909d6a17cffa0fb7259a82b4 Mon Sep 17 00:00:00 2001
From: Victor Petukhov <victor.petukhov@jetbrains.com>
Date: Wed, 12 May 2021 19:13:05 +0300
Subject: [PATCH] [Jspecify] Do preserving flexibility when adding equality
 constraints by taking into account type with enhancement

---
 .../calls/components/ResolutionParts.kt        | 18 +++++++++++++-----
 .../warnMode/NonPlatformTypeParameter.fir.kt   |  8 ++++----
 .../warnMode/NonPlatformTypeParameter.kt       |  8 ++++----
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt b/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt
index 8ef5ab48de5..aa303d826fb 100644
--- a/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt
+++ b/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()) {
-        KotlinTypeFactory.flexibleType(type.makeNotNullable().lowerIfFlexible(), type.makeNullable().upperIfFlexible())
-    } else type
+    ): KotlinType {
+        fun createFlexibleType() =
+            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,
diff --git a/compiler/testData/diagnostics/foreignAnnotationsTests/java8Tests/jspecify/warnMode/NonPlatformTypeParameter.fir.kt b/compiler/testData/diagnostics/foreignAnnotationsTests/java8Tests/jspecify/warnMode/NonPlatformTypeParameter.fir.kt
index 97bf3d28136..1f5df653fd9 100644
--- a/compiler/testData/diagnostics/foreignAnnotationsTests/java8Tests/jspecify/warnMode/NonPlatformTypeParameter.fir.kt
+++ b/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}
-    a1.bar<T>(null)
+    // jspecify_nullness_mismatch
+    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}
-    a2.bar<T>(null)
+    // jspecify_nullness_mismatch
+    a2.bar<T>(<!NULLABILITY_MISMATCH_BASED_ON_JAVA_ANNOTATIONS!>null<!>)
     a2.bar<T>(x)
 }
\ No newline at end of file
diff --git a/compiler/testData/diagnostics/foreignAnnotationsTests/java8Tests/jspecify/warnMode/NonPlatformTypeParameter.kt b/compiler/testData/diagnostics/foreignAnnotationsTests/java8Tests/jspecify/warnMode/NonPlatformTypeParameter.kt
index 97bf3d28136..1f5df653fd9 100644
--- a/compiler/testData/diagnostics/foreignAnnotationsTests/java8Tests/jspecify/warnMode/NonPlatformTypeParameter.kt
+++ b/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}
-    a1.bar<T>(null)
+    // jspecify_nullness_mismatch
+    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}
-    a2.bar<T>(null)
+    // jspecify_nullness_mismatch
+    a2.bar<T>(<!NULLABILITY_MISMATCH_BASED_ON_JAVA_ANNOTATIONS!>null<!>)
     a2.bar<T>(x)
 }
\ No newline at end of file