NI: take into account flexible types during capturing arguments from expression

^KT-40544 Fixed

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

@@ -10611,6 +10611,11 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirOldFronte
                 runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedTypeWithTypeVariableSubtyping.kt");
             }
 
+            @TestMetadata("capturingFromArgumentOfFlexibleType.kt")
+            public void testCapturingFromArgumentOfFlexibleType() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturingFromArgumentOfFlexibleType.kt");
+            }
+
             @TestMetadata("expectedTypeMismatchWithInVariance.kt")
             public void testExpectedTypeMismatchWithInVariance() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/expectedTypeMismatchWithInVariance.kt");

compiler/testData/diagnostics/tests/generics/projectionsScope/platformSuperClass.kt

@@ -10,10 +10,10 @@ public class Clazz<T> {
 // FILE: main.kt
 fun test(clazz: Clazz<*>) {
     clazz.t checkType { _<Any?>() }
-    clazz.getSuperClass() checkType { <!NI;UNRESOLVED_REFERENCE_WRONG_RECEIVER!>_<!><Clazz<*>?>() }
+    clazz.getSuperClass() checkType { _<Clazz<*>?>() }
     clazz.getSuperClass().t checkType { _<Any?>() }
 
-    clazz.superClass checkType { <!NI;UNRESOLVED_REFERENCE_WRONG_RECEIVER!>_<!><Clazz<*>?>() }
+    clazz.superClass checkType { _<Clazz<*>?>() }
     clazz.superClass.t checkType { _<Any?>() }
 
     // See KT-9294

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

@@ -0,0 +1,18 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNCHECKED_CAST
+
+// FILE: Test.java
+class Test {
+    static Foo<? extends Number> getFoo() {
+        return null;
+    }
+}
+
+// FILE: main.kt
+class Foo<T>
+
+fun <T> id(x: T) = null as T
+
+fun test() {
+    Test.getFoo()
+    <!DEBUG_INFO_EXPRESSION_TYPE("Foo<out kotlin.Number..kotlin.Number?!>")!>id(Test.getFoo())<!>
+}

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

@@ -0,0 +1,18 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNCHECKED_CAST
+
+// FILE: Test.java
+class Test {
+    static Foo<? extends Number> getFoo() {
+        return null;
+    }
+}
+
+// FILE: main.kt
+class Foo<T>
+
+fun <T> id(x: T) = null as T
+
+fun test() {
+    <!DEBUG_INFO_EXPRESSION_TYPE("(Foo<out (kotlin.Number..kotlin.Number?)>..Foo<out (kotlin.Number..kotlin.Number?)>?)")!>Test.getFoo()<!>
+    <!DEBUG_INFO_EXPRESSION_TYPE("(Foo<out (kotlin.Number..kotlin.Number?)>..Foo<out (kotlin.Number..kotlin.Number?)>?)")!>id(Test.getFoo())<!>
+}

compiler/testData/diagnostics/tests/inference/capturedTypes/capturingFromArgumentOfFlexibleType.txt

@@ -0,0 +1,21 @@
+package
+
+public fun </*0*/ T> id(/*0*/ x: T): T
+public fun test(): kotlin.Unit
+
+public final class Foo</*0*/ T> {
+    public constructor Foo</*0*/ T>()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public/*package*/ open class Test {
+    public/*package*/ constructor Test()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+
+    // Static members
+    public/*package*/ open fun getFoo(): Foo<out kotlin.Number!>!
+}

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

@@ -10618,6 +10618,11 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTestWithFirVali
                     runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedTypeWithTypeVariableSubtyping.kt");
                 }
 
+                @TestMetadata("capturingFromArgumentOfFlexibleType.kt")
+                public void testCapturingFromArgumentOfFlexibleType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturingFromArgumentOfFlexibleType.kt");
+                }
+
                 @TestMetadata("expectedTypeMismatchWithInVariance.kt")
                 public void testExpectedTypeMismatchWithInVariance() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/expectedTypeMismatchWithInVariance.kt");

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

@@ -10613,6 +10613,11 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                     runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturedTypeWithTypeVariableSubtyping.kt");
                 }
 
+                @TestMetadata("capturingFromArgumentOfFlexibleType.kt")
+                public void testCapturingFromArgumentOfFlexibleType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/capturingFromArgumentOfFlexibleType.kt");
+                }
+
                 @TestMetadata("expectedTypeMismatchWithInVariance.kt")
                 public void testExpectedTypeMismatchWithInVariance() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/capturedTypes/expectedTypeMismatchWithInVariance.kt");

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

@@ -34,10 +34,7 @@ fun prepareArgumentTypeRegardingCaptureTypes(argumentType: UnwrappedType): Unwra
     return if (argumentType is NewCapturedType) null else captureFromExpression(argumentType)
 }
 
-fun captureFromExpression(type: UnwrappedType): UnwrappedType? =
-    captureFromExpression(type.lowerIfFlexible())
-
-private fun captureFromExpression(type: SimpleType): UnwrappedType? {
+fun captureFromExpression(type: UnwrappedType): UnwrappedType? {
     val typeConstructor = type.constructor
     if (typeConstructor is IntersectionTypeConstructor) {
         var changed = false
@@ -54,10 +51,26 @@ private fun captureFromExpression(type: SimpleType): UnwrappedType? {
 }
 
 // this function suppose that input type is simple classifier type
-internal fun captureFromArguments(
-    type: SimpleType,
-    status: CaptureStatus
-): SimpleType? {
+internal fun captureFromArguments(type: SimpleType, status: CaptureStatus) =
+    captureArguments(type, status)?.let { type.replaceArguments(it) }
+
+private fun captureFromArguments(type: UnwrappedType, status: CaptureStatus): UnwrappedType? {
+    val capturedArguments = captureArguments(type, status) ?: return null
+
+    return if (type is FlexibleType) {
+        KotlinTypeFactory.flexibleType(
+            type.lowerBound.replaceArguments(capturedArguments),
+            type.upperBound.replaceArguments(capturedArguments)
+        )
+    } else {
+        type.replaceArguments(capturedArguments)
+    }
+}
+
+private fun UnwrappedType.replaceArguments(arguments: List<TypeProjection>) =
+    KotlinTypeFactory.simpleType(annotations, constructor, arguments, isMarkedNullable)
+
+private fun captureArguments(type: UnwrappedType, status: CaptureStatus): List<TypeProjection>? {
     if (type.arguments.size != type.constructor.parameters.size) return null
 
     val arguments = type.arguments
@@ -77,6 +90,7 @@ internal fun captureFromArguments(
     }
 
     val substitutor = TypeConstructorSubstitution.create(type.constructor, capturedArguments).buildSubstitutor()
+
     for (index in arguments.indices) {
         val oldProjection = arguments[index]
         val newProjection = capturedArguments[index]
@@ -94,7 +108,7 @@ internal fun captureFromArguments(
         capturedType.constructor.initializeSupertypes(capturedTypeSupertypes)
     }
 
-    return KotlinTypeFactory.simpleType(type.annotations, type.constructor, capturedArguments, type.isMarkedNullable)
+    return capturedArguments
 }
 
 /**