Reuse captured arguments for flexible type's bounds properly, by equality of type constructors modulo mutability and type argument

^KT-43630 Fixed

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/fir/FirOldFrontendDiagnosticsTestGenerated.java

@@ -10659,6 +10659,16 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirOldFronte
                 runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/captureTypeOnlyOnTopLevel.kt");
             }
 
+            @TestMetadata("capturedFlexibleIntersectionTypesWithDifferentBounds.kt")
+            public void testCapturedFlexibleIntersectionTypesWithDifferentBounds() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedFlexibleIntersectionTypesWithDifferentBounds.kt");
+            }
+
+            @TestMetadata("capturedFlexibleIntersectionTypesWithDifferentConstructors.kt")
+            public void testCapturedFlexibleIntersectionTypesWithDifferentConstructors() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedFlexibleIntersectionTypesWithDifferentConstructors.kt");
+            }
+
             @TestMetadata("capturedType.kt")
             public void testCapturedType() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedType.kt");

compiler/testData/diagnostics/tests/inference/capturedTypes/capturedFlexibleIntersectionTypesWithDifferentBounds.fir.kt

@@ -0,0 +1,23 @@
+// FILE: Bar.java
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// SKIP_TXT
+
+public class Bar<K, N> { }
+
+// FILE: Foo.java
+
+public class Foo<P> extends Bar<Integer, Integer> {
+    public static final Bar bar = null;
+}
+
+// FILE: main.kt
+
+fun <P> takeFoo(foo: Foo<P>) {}
+
+fun main(x: Foo<*>?) {
+    val y = Foo.bar
+    if (y !is Foo<*>?) return
+    if (y == null) return
+    if (x != y) return
+    takeFoo(x) // Here we capture `{Bar<Any!, Any!> & Foo<*>}..Foo<*>?`
+}

compiler/testData/diagnostics/tests/inference/capturedTypes/capturedFlexibleIntersectionTypesWithDifferentBounds.kt

@@ -0,0 +1,23 @@
+// FILE: Bar.java
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// SKIP_TXT
+
+public class Bar<K, N> { }
+
+// FILE: Foo.java
+
+public class Foo<P> extends Bar<Integer, Integer> {
+    public static final Bar bar = null;
+}
+
+// FILE: main.kt
+
+fun <P> takeFoo(foo: Foo<P>) {}
+
+fun main(x: Foo<*>?) {
+    val y = Foo.bar
+    if (y !is Foo<*>?) return
+    if (y == null) return
+    if (x != y) return
+    takeFoo(<!DEBUG_INFO_SMARTCAST!>x<!>) // Here we capture `{Bar<Any!, Any!> & Foo<*>}..Foo<*>?`
+}

compiler/testData/diagnostics/tests/inference/capturedTypes/capturedFlexibleIntersectionTypesWithDifferentConstructors.fir.kt

@@ -0,0 +1,29 @@
+// FULL_JDK
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// SKIP_TXT
+
+// FILE: Bar.java
+
+public class Bar<K, N> { }
+
+// FILE: Foo.java
+
+import java.util.List;
+
+public class Foo<P> extends Bar<Integer, Integer> {
+    public static final List<?> bar = null;
+}
+
+// FILE: main.kt
+
+fun <P> takeFoo(foo: Foo<P>) {}
+
+fun main(x: Foo<*>?) {
+    val y = Foo.bar
+    if (y !is Foo<*>?) return
+    if (y == null) return
+    if (x != y) return
+    // Here we capture `({Foo<*> & MutableList<*>}..{Foo<*>? & List<*>?})`
+    // `*` inside `MutableList` and `List` have to become the same captured type
+    takeFoo(x)
+}

compiler/testData/diagnostics/tests/inference/capturedTypes/capturedFlexibleIntersectionTypesWithDifferentConstructors.kt

@@ -0,0 +1,29 @@
+// FULL_JDK
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// SKIP_TXT
+
+// FILE: Bar.java
+
+public class Bar<K, N> { }
+
+// FILE: Foo.java
+
+import java.util.List;
+
+public class Foo<P> extends Bar<Integer, Integer> {
+    public static final List<?> bar = null;
+}
+
+// FILE: main.kt
+
+fun <P> takeFoo(foo: Foo<P>) {}
+
+fun main(x: Foo<*>?) {
+    val y = Foo.bar
+    if (y !is Foo<*>?) return
+    if (y == null) return
+    if (x != y) return
+    // Here we capture `({Foo<*> & MutableList<*>}..{Foo<*>? & List<*>?})`
+    // `*` inside `MutableList` and `List` have to become the same captured type
+    takeFoo(<!DEBUG_INFO_SMARTCAST!>x<!>)
+}

compiler/tests-gen/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -10666,6 +10666,16 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTestWithFirVali
                     runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/captureTypeOnlyOnTopLevel.kt");
                 }
 
+                @TestMetadata("capturedFlexibleIntersectionTypesWithDifferentBounds.kt")
+                public void testCapturedFlexibleIntersectionTypesWithDifferentBounds() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedFlexibleIntersectionTypesWithDifferentBounds.kt");
+                }
+
+                @TestMetadata("capturedFlexibleIntersectionTypesWithDifferentConstructors.kt")
+                public void testCapturedFlexibleIntersectionTypesWithDifferentConstructors() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedFlexibleIntersectionTypesWithDifferentConstructors.kt");
+                }
+
                 @TestMetadata("capturedType.kt")
                 public void testCapturedType() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedType.kt");

compiler/tests-gen/org/jetbrains/kotlin/checkers/javac/DiagnosticsUsingJavacTestGenerated.java

@@ -10661,6 +10661,16 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                     runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/captureTypeOnlyOnTopLevel.kt");
                 }
 
+                @TestMetadata("capturedFlexibleIntersectionTypesWithDifferentBounds.kt")
+                public void testCapturedFlexibleIntersectionTypesWithDifferentBounds() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedFlexibleIntersectionTypesWithDifferentBounds.kt");
+                }
+
+                @TestMetadata("capturedFlexibleIntersectionTypesWithDifferentConstructors.kt")
+                public void testCapturedFlexibleIntersectionTypesWithDifferentConstructors() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedFlexibleIntersectionTypesWithDifferentConstructors.kt");
+                }
+
                 @TestMetadata("capturedType.kt")
                 public void testCapturedType() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedType.kt");

core/descriptors/src/org/jetbrains/kotlin/types/checker/NewCapturedType.kt

@@ -23,12 +23,30 @@ import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.resolve.calls.inference.CapturedTypeConstructor
 import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.types.*
+import org.jetbrains.kotlin.types.FlexibleTypeBoundsChecker.areTypesMayBeLowerAndUpperBoundsOfSameFlexibleTypeByMutability
 import org.jetbrains.kotlin.types.model.CaptureStatus
 import org.jetbrains.kotlin.types.model.CapturedTypeMarker
 import org.jetbrains.kotlin.types.refinement.TypeRefinement
 import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.types.typeUtil.builtIns
 
+private class CapturedArguments(val capturedArguments: List<TypeProjection>, private val originalType: KotlinType) {
+    fun isSuitableForType(type: KotlinType): Boolean {
+        val areArgumentsMatched = type.arguments.withIndex().all { (i, typeArgumentsType) ->
+            originalType.arguments.size > i && typeArgumentsType == originalType.arguments[i]
+        }
+
+        if (!areArgumentsMatched) return false
+
+        val areConstructorsMatched = originalType.constructor == type.constructor
+                || areTypesMayBeLowerAndUpperBoundsOfSameFlexibleTypeByMutability(originalType, type)
+
+        if (!areConstructorsMatched) return false
+
+        return true
+    }
+}
+
 // null means that type should be leaved as is
 fun prepareArgumentTypeRegardingCaptureTypes(argumentType: UnwrappedType): UnwrappedType? {
     return if (argumentType is NewCapturedType) null else captureFromExpression(argumentType)
@@ -42,33 +60,44 @@ fun captureFromExpression(type: UnwrappedType): UnwrappedType? {
     }
 
     /*
-     * We capture intersection types in two stages:
+     * We capture arguments in the intersection types in specific way:
-     * capture type arguments for each component and replace it in the original type after that.
+     *  1) Firstly, we create captured arguments for all type arguments grouped by a type constructor* and a type argument's type.
-     * This is to substitute captured types into flexible types properly:
+     *      It means, that we create only one captured argument for two types `Foo<*>` and `Foo<*>?` within a flexible type, for instance.
-     * we should have the same captured types both for lower bound and for upper one.
+     *      * In addition to grouping by type constructors, we look at possibility locating of two types in different bounds of the same flexible type.
-     *
+     *        This is necessary in order to create the same captured arguments,
-     * Example:
+     *        for example, for `MutableList` in the lower bound of the flexible type and for `List` in the upper one.
-     *  The original type: ({Comparable<*> & java.io.Serializable}..{Comparable<*>? & java.io.Serializable?})
+     *        Example: MutableList<*>..List<*>? -> MutableList<Captured1(*)>..List<Captured2(*)>?, Captured1(*) and Captured2(*) are the same.
-     *  Result of capturing arguments by components: [[CapturedType(*)], null]
+     *  2) Secondly, we replace type arguments with captured arguments by given a type constructor and type arguments.
-     *  The resulting type: ({Comparable<CapturedType(*)> & java.io.Serializable}..{Comparable<CapturedType(*)>? & java.io.Serializable?})
      */
-    val capturedArgumentsByComponents = captureArgumentsForIntersectionType(typeConstructor) ?: return null
+    val capturedArgumentsByComponents = captureArgumentsForIntersectionType(type) ?: return null
 
-    fun replaceArgumentsByComponents(typeToReplace: UnwrappedType) =
+    // We reuse `TypeToCapture` for some types, suitability to reuse defines by `isSuitableForType`
-        typeToReplace.constructor.supertypes.mapIndexed { i, componentType ->
+    fun findCorrespondingCapturedArgumentsForType(type: KotlinType) =
-            val capturedArguments = capturedArgumentsByComponents[i] ?: return@mapIndexed componentType.asSimpleType()
+        capturedArgumentsByComponents.find { typeToCapture -> typeToCapture.isSuitableForType(type) }?.capturedArguments
-            componentType.unwrap().replaceArguments(capturedArguments)
+
+    fun replaceArgumentsWithCapturedArgumentsByIntersectionComponents(typeToReplace: UnwrappedType): List<SimpleType> {
+        return if (typeToReplace.constructor is IntersectionTypeConstructor) {
+            typeToReplace.constructor.supertypes.map { componentType ->
+                val capturedArguments = findCorrespondingCapturedArgumentsForType(componentType)
+                    ?: return@map componentType.asSimpleType()
+                componentType.unwrap().replaceArguments(capturedArguments)
+            }
+        } else {
+            val capturedArguments = findCorrespondingCapturedArgumentsForType(typeToReplace)
+                ?: return listOf(typeToReplace.asSimpleType())
+            listOf(typeToReplace.unwrap().replaceArguments(capturedArguments))
         }
+    }
 
     return if (type is FlexibleType) {
-        val lowerIntersectedType =
+        val lowerIntersectedType = intersectTypes(replaceArgumentsWithCapturedArgumentsByIntersectionComponents(type.lowerBound))
-            intersectTypes(replaceArgumentsByComponents(type.lowerBound)).makeNullableAsSpecified(type.lowerBound.isMarkedNullable)
+            .makeNullableAsSpecified(type.lowerBound.isMarkedNullable)
-        val upperIntersectedType =
+        val upperIntersectedType = intersectTypes(replaceArgumentsWithCapturedArgumentsByIntersectionComponents(type.upperBound))
-            intersectTypes(replaceArgumentsByComponents(type.upperBound)).makeNullableAsSpecified(type.upperBound.isMarkedNullable)
+            .makeNullableAsSpecified(type.upperBound.isMarkedNullable)
 
         KotlinTypeFactory.flexibleType(lowerIntersectedType, upperIntersectedType)
     } else {
-        intersectTypes(replaceArgumentsByComponents(type)).makeNullableAsSpecified(type.isMarkedNullable)
+        intersectTypes(replaceArgumentsWithCapturedArgumentsByIntersectionComponents(type)).makeNullableAsSpecified(type.isMarkedNullable)
     }
 }
 
@@ -76,15 +105,29 @@ fun captureFromExpression(type: UnwrappedType): UnwrappedType? {
 internal fun captureFromArguments(type: SimpleType, status: CaptureStatus) =
     captureArguments(type, status)?.let { type.replaceArguments(it) }
 
-private fun captureArgumentsForIntersectionType(typeConstructor: TypeConstructor): List<List<TypeProjection>?>? {
+private fun captureArgumentsForIntersectionType(type: KotlinType): List<CapturedArguments>? {
+    // It's possible to have one of the bounds as non-intersection type
+    fun getTypesToCapture(type: KotlinType) =
+        if (type.constructor is IntersectionTypeConstructor) type.constructor.supertypes else listOf(type)
+
+    val filteredTypesToCapture =
+        if (type is FlexibleType) {
+            val typesToCapture = getTypesToCapture(type.lowerBound) + getTypesToCapture(type.upperBound)
+            typesToCapture.distinctBy { (FlexibleTypeBoundsChecker.getBaseBoundFqNameByMutability(it) ?: it.constructor) to it.arguments }
+        } else type.constructor.supertypes
+
     var changed = false
-    val capturedArgumentsByComponents = typeConstructor.supertypes.map { supertype ->
+
-        captureArguments(supertype.unwrap(), CaptureStatus.FROM_EXPRESSION)?.apply { changed = true }
+    val capturedArgumentsByTypes = filteredTypesToCapture.mapNotNull { typeToCapture ->
+        val capturedArguments = captureArguments(typeToCapture.unwrap(), CaptureStatus.FROM_EXPRESSION)
+            ?: return@mapNotNull null
+        changed = true
+        CapturedArguments(capturedArguments, originalType = typeToCapture)
     }
 
     if (!changed) return null
 
-    return capturedArgumentsByComponents
+    return capturedArgumentsByTypes
 }
 
 private fun captureFromArguments(type: UnwrappedType, status: CaptureStatus): UnwrappedType? {