K2: Imitate K1 behavior for case of captured types with a raw supertype

See the comments in the code, but mostly the motivation is that once
it was decided to stick with such a legacy thing as raw types,
we are ok with some corner-cases hacks for them
(if there are not too many of them) and they don't break anything
when there are no raw types in the code.

^KT-56616 Fixed

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -19396,6 +19396,18 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/j+k/rawTypeScope.kt");
             }
 
+            @Test
+            @TestMetadata("rawTypesFromCaptured.kt")
+            public void testRawTypesFromCaptured() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/j+k/rawTypesFromCaptured.kt");
+            }
+
+            @Test
+            @TestMetadata("rawTypesFromCapturedOriginal.kt")
+            public void testRawTypesFromCapturedOriginal() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/j+k/rawTypesFromCapturedOriginal.kt");
+            }
+
             @Test
             @TestMetadata("rawUpperBounds.kt")
             public void testRawUpperBounds() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java

@@ -19396,6 +19396,18 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                 runTest("compiler/testData/diagnostics/tests/j+k/rawTypeScope.kt");
             }
 
+            @Test
+            @TestMetadata("rawTypesFromCaptured.kt")
+            public void testRawTypesFromCaptured() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/j+k/rawTypesFromCaptured.kt");
+            }
+
+            @Test
+            @TestMetadata("rawTypesFromCapturedOriginal.kt")
+            public void testRawTypesFromCapturedOriginal() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/j+k/rawTypesFromCapturedOriginal.kt");
+            }
+
             @Test
             @TestMetadata("rawUpperBounds.kt")
             public void testRawUpperBounds() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java

@@ -19402,6 +19402,18 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                 runTest("compiler/testData/diagnostics/tests/j+k/rawTypeScope.kt");
             }
 
+            @Test
+            @TestMetadata("rawTypesFromCaptured.kt")
+            public void testRawTypesFromCaptured() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/j+k/rawTypesFromCaptured.kt");
+            }
+
+            @Test
+            @TestMetadata("rawTypesFromCapturedOriginal.kt")
+            public void testRawTypesFromCapturedOriginal() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/j+k/rawTypesFromCapturedOriginal.kt");
+            }
+
             @Test
             @TestMetadata("rawUpperBounds.kt")
             public void testRawUpperBounds() throws Exception {

compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/ConeInferenceContext.kt

@@ -558,6 +558,11 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
 
     override fun useRefinedBoundsForTypeVariableInFlexiblePosition(): Boolean = true
 
+    override fun KotlinTypeMarker.convertToNonRaw(): KotlinTypeMarker {
+        require(this is ConeKotlinType)
+        return this.convertToNonRawVersion()
+    }
+
     override fun createSubstitutorForSuperTypes(baseType: KotlinTypeMarker): TypeSubstitutorMarker? =
         if (baseType is ConeLookupTagBasedType) createSubstitutionForSupertype(baseType, session) else null
 

compiler/resolution.common/src/org/jetbrains/kotlin/types/AbstractTypeApproximator.kt

@@ -297,6 +297,18 @@ abstract class AbstractTypeApproximator(
             type.isMarkedNullable() -> baseResult.withNullability(true)
             type.isProjectionNotNull() -> baseResult.withNullability(false)
             else -> baseResult
+        }.let {
+            when {
+                // This is just a hack that is necessary to preserve compatibility with K1 where return type of the calls
+                // if they contain a captured types with RAW supertype would be approximated to a regular non-raw flexible type
+                // See CapturedTypeApproximationKt.approximateCapturedTypes and especially the comment
+                // "// tod*: dynamic & raw type?" before it :)
+                // If we don't repeat that behavior, we would stumble upon KT-56616 with hardly having any workarounds.
+                isK2 && conf.convertToNonRawVersionAfterApproximationInK2 && it.isRawType() -> {
+                    it.convertToNonRaw()
+                }
+                else -> it
+            }
         }
     }
 

compiler/resolution.common/src/org/jetbrains/kotlin/types/TypeApproximatorConfiguration.kt

@@ -26,6 +26,12 @@ open class TypeApproximatorConfiguration {
     open val intersectionTypesInContravariantPositions = false
     open val localTypes = false
 
+    /**
+     * Is only expected to be true for FinalApproximationAfterResolutionAndInference
+     * But it's only used for K2 to reproduce K1 behavior for the approximation of resolved calls
+     */
+    open val convertToNonRawVersionAfterApproximationInK2 get() = false
+
     /**
      * Whether to approximate anonymous type. This flag does not have any effect if `localTypes` is true because all anonymous types are
      * local.
@@ -88,6 +94,8 @@ open class TypeApproximatorConfiguration {
         AbstractCapturedTypesApproximation(CaptureStatus.FROM_EXPRESSION) {
         override val integerLiteralConstantType: Boolean get() = true
         override val intersectionTypesInContravariantPositions: Boolean get() = true
+
+        override val convertToNonRawVersionAfterApproximationInK2: Boolean get() = true
     }
 
     object TypeArgumentApproximation : AbstractCapturedTypesApproximation(null) {

compiler/testData/diagnostics/tests/j+k/rawTypesFromCaptured.fir.kt

@@ -0,0 +1,25 @@
+// ISSUE: KT-56616
+
+// FILE: StubElement.java
+
+public interface StubElement<T> {
+    <E> E bar(E v);
+}
+
+// FILE: StubBasedPsiElement.java
+
+public interface StubBasedPsiElement<T extends StubElement> {
+    T foo1();
+
+    StubElement foo2();
+}
+
+// FILE: test.kt
+
+fun StubBasedPsiElement<*>.foo(): String? {
+    if ("".hashCode() == 0) {
+        return foo1().bar("")
+    }
+
+    return <!RETURN_TYPE_MISMATCH!>foo2().bar("")<!>
+}

compiler/testData/diagnostics/tests/j+k/rawTypesFromCaptured.kt

@@ -0,0 +1,25 @@
+// ISSUE: KT-56616
+
+// FILE: StubElement.java
+
+public interface StubElement<T> {
+    <E> E bar(E v);
+}
+
+// FILE: StubBasedPsiElement.java
+
+public interface StubBasedPsiElement<T extends StubElement> {
+    T foo1();
+
+    StubElement foo2();
+}
+
+// FILE: test.kt
+
+fun StubBasedPsiElement<*>.foo(): String? {
+    if ("".hashCode() == 0) {
+        return foo1().bar("")
+    }
+
+    return <!TYPE_MISMATCH!>foo2().<!TYPE_MISMATCH!>bar("")<!><!>
+}

compiler/testData/diagnostics/tests/j+k/rawTypesFromCapturedOriginal.kt

@@ -0,0 +1,37 @@
+// FIR_IDENTICAL
+// WITH_STDLIB
+// ISSUE: KT-56616
+
+// FILE: StubElement.java
+
+import org.jetbrains.annotations.NotNull;
+
+public interface StubElement<T extends PsiElement> {
+    <E extends PsiElement> E @NotNull [] getChildrenByType(@NotNull String filter, final E[] array);
+}
+
+// FILE: PsiElement.java
+
+public interface PsiElement
+
+// FILE: StubBasedPsiElement.java
+
+public interface StubBasedPsiElement<Stub extends StubElement> extends PsiElement {
+    Stub getStub();
+}
+
+// FILE: test.kt
+
+private val STRING_TEMPLATE_EMPTY_ARRAY = emptyArray<KtStringTemplateExpression>()
+
+open class KtStringTemplateExpression : PsiElement
+
+fun StubBasedPsiElement<*>.foo(): KtStringTemplateExpression? {
+    stub?.let {
+        // K1: Array<KtStringTemplateExpression>, was K2: Array<PsiElement>
+        val expressions = it.getChildrenByType("", STRING_TEMPLATE_EMPTY_ARRAY)
+        // Ok in K1, Should not be error in K2
+        return expressions.firstOrNull()
+    }
+    return null
+}

compiler/testData/diagnostics/tests/platformTypes/rawTypes/rawTypeInUpperBound.fir.kt

@@ -1,42 +0,0 @@
-// !DIAGNOSTICS: -UNUSED_VARIABLE
-// FILE: A.java
-
-import java.util.List;
-
-public class A<T> {
-    List<String> getChildrenStubs() { return null; }
-    void consume(T x) {}
-
-    T produce() { return null; }
-}
-
-// FILE: B.java
-
-public class B<E extends A> {
-    public E foo() { return null;}
-    E field;
-}
-
-// FILE: Test.java
-
-public class Test {
-    static B rawB = null;
-}
-
-// FILE: main.kt
-
-fun foo(x: B<*>) {
-    // TODO: In K1, x.foo() now is flexible type instead of raw, because of captured type approximation
-    // Works in K2 as expected: x.foo() returns raw `A`, thus it's `getChildrenStubs` has a type `MutableList<Any!>..List<*>?`
-    val q: MutableList<String> = <!INITIALIZER_TYPE_MISMATCH, TYPE_MISMATCH!>x.foo().getChildrenStubs()<!>
-
-    // Raw(B).field erased to A<Any!>..A<out Any!>?
-    Test.rawB.field = A<String>()
-    val anyA: A<Any> = Test.rawB.field
-
-    // FIR doesn't work here, because
-    // field has a type of just 'A' and it's not clear why should it accept 'String' at consume
-    // NB: some kind of BareTypeScope should be in use here
-    Test.rawB.field.consume("")
-    val y: Any = Test.rawB.field.produce()
-}

compiler/testData/diagnostics/tests/platformTypes/rawTypes/rawTypeInUpperBound.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_VARIABLE
 // FILE: A.java
 
@@ -26,8 +27,7 @@ public class Test {
 // FILE: main.kt
 
 fun foo(x: B<*>) {
-    // TODO: In K1, x.foo() now is flexible type instead of raw, because of captured type approximation
+    // x.foo()  is flexible type instead of raw, because of captured type approximation
-    // Works in K2 as expected: x.foo() returns raw `A`, thus it's `getChildrenStubs` has a type `MutableList<Any!>..List<*>?`
     val q: MutableList<String> = x.foo().getChildrenStubs()
 
     // Raw(B).field erased to A<Any!>..A<out Any!>?

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -19402,6 +19402,18 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/j+k/rawTypeScope.kt");
             }
 
+            @Test
+            @TestMetadata("rawTypesFromCaptured.kt")
+            public void testRawTypesFromCaptured() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/j+k/rawTypesFromCaptured.kt");
+            }
+
+            @Test
+            @TestMetadata("rawTypesFromCapturedOriginal.kt")
+            public void testRawTypesFromCapturedOriginal() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/j+k/rawTypesFromCapturedOriginal.kt");
+            }
+
             @Test
             @TestMetadata("rawUpperBounds.kt")
             public void testRawUpperBounds() throws Exception {

core/compiler.common/src/org/jetbrains/kotlin/types/model/TypeSystemContext.kt

@@ -312,6 +312,11 @@ interface TypeSystemInferenceExtensionContext : TypeSystemContext, TypeSystemBui
      */
     fun useRefinedBoundsForTypeVariableInFlexiblePosition(): Boolean
 
+    /**
+     * It's only relevant for K2 (and is not expected to be implemented properly in other contexts)
+     */
+    fun KotlinTypeMarker.convertToNonRaw(): KotlinTypeMarker
+
     fun createCapturedStarProjectionForSelfType(
         typeVariable: TypeVariableTypeConstructorMarker,
         typesForRecursiveTypeParameters: List<KotlinTypeMarker>,
@@ -357,6 +362,7 @@ interface TypeSystemContext : TypeSystemOptimizationContext {
     fun FlexibleTypeMarker.asDynamicType(): DynamicTypeMarker?
 
     fun KotlinTypeMarker.isRawType(): Boolean
+
     fun FlexibleTypeMarker.upperBound(): SimpleTypeMarker
 
     fun FlexibleTypeMarker.lowerBound(): SimpleTypeMarker

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

@@ -158,6 +158,10 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
         return this is RawType
     }
 
+    override fun KotlinTypeMarker.convertToNonRaw(): KotlinTypeMarker {
+        error("Is not expected to be called in K1")
+    }
+
     override fun FlexibleTypeMarker.upperBound(): SimpleTypeMarker {
         require(this is FlexibleType, this::errorMessage)
         return this.upperBound