[FIR] Make captureFromTypeParameterUpperBoundIfNeeded smarter

Before this change `ARGUMENT_TYPE_MISMATCH` would complain that
`Y` "is not a subtype of" `Inv<Y>`, because the function would only
check immediate bounds of the type parameter `Y`. `chosenSupertype`
would be `X`, not `Inv<out kotlin/String>`.

^KT-60056

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/ConeTypeCompatibilityChecker.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.analysis.checkers
 
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.fir.collectUpperBounds
 import org.jetbrains.kotlin.fir.declarations.utils.*
 import org.jetbrains.kotlin.fir.isPrimitiveType
 import org.jetbrains.kotlin.fir.languageVersionSettings

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt

@@ -34,7 +34,6 @@ import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.scopes.*
 import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
 import org.jetbrains.kotlin.fir.scopes.impl.multipleDelegatesWithTheSameSignature
-import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.symbols.impl.*
@@ -810,46 +809,6 @@ fun FirElement.isLhsOfAssignment(context: CheckerContext): Boolean {
     return lastQualified is FirVariableAssignment && lastQualified.lValue == this
 }
 
-/**
- * Collects the upper bounds as [ConeClassLikeType].
- */
-fun ConeKotlinType?.collectUpperBounds(): Set<ConeClassLikeType> {
-    if (this == null) return emptySet()
-
-    val upperBounds = mutableSetOf<ConeClassLikeType>()
-    val seen = mutableSetOf<ConeKotlinType>()
-    fun collect(type: ConeKotlinType) {
-        if (!seen.add(type)) return // Avoid infinite recursion.
-
-        when (type) {
-            is ConeErrorType -> return // Ignore error types
-            is ConeLookupTagBasedType -> when (type) {
-                is ConeClassLikeType -> upperBounds.add(type)
-                is ConeTypeParameterType -> {
-                    val symbol = type.lookupTag.typeParameterSymbol
-                    symbol.resolvedBounds.forEach { collect(it.coneType) }
-                }
-                else -> error("missing branch for ${javaClass.name}")
-            }
-            is ConeTypeVariableType -> {
-                val symbol = (type.typeConstructor.originalTypeParameter as? ConeTypeParameterLookupTag)?.typeParameterSymbol ?: return
-                symbol.resolvedBounds.forEach { collect(it.coneType) }
-            }
-            is ConeDefinitelyNotNullType -> collect(type.original)
-            is ConeIntersectionType -> type.intersectedTypes.forEach(::collect)
-            is ConeFlexibleType -> collect(type.upperBound)
-            is ConeCapturedType -> type.constructor.supertypes?.forEach(::collect)
-            is ConeIntegerConstantOperatorType -> upperBounds.add(type.getApproximatedType())
-            is ConeStubType, is ConeIntegerLiteralConstantType -> {
-                error("$type should not reach here")
-            }
-        }
-    }
-
-    collect(this)
-    return upperBounds
-}
-
 fun ConeKotlinType.leastUpperBound(session: FirSession): ConeKotlinType {
     val upperBounds = collectUpperBounds().takeIf { it.isNotEmpty() } ?: return session.builtinTypes.nullableAnyType.type
     return ConeTypeIntersector.intersectTypes(session.typeContext, upperBounds)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirFiniteBoundRestrictionChecker.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.analysis.checkers.collectUpperBounds
+import org.jetbrains.kotlin.fir.collectUpperBounds
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.typeParameterSymbols
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirNonExpansiveInheritanceRestrictionChecker.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.analysis.checkers.collectUpperBounds
+import org.jetbrains.kotlin.fir.collectUpperBounds
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirEqualityCompatibilityChecker.kt

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.KtDiagnosticFactory2
 import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.analysis.checkers.collectUpperBounds
+import org.jetbrains.kotlin.fir.collectUpperBounds
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.utils.isEnumClass

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirMissingDependencySupertypeChecker.kt

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.config.AnalysisFlags
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
-import org.jetbrains.kotlin.fir.analysis.checkers.collectUpperBounds
+import org.jetbrains.kotlin.fir.collectUpperBounds
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirBasicDeclarationChecker
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt

@@ -668,11 +668,9 @@ internal fun captureFromTypeParameterUpperBoundIfNeeded(
 ): ConeKotlinType {
     val expectedTypeClassId = expectedType.upperBoundIfFlexible().classId ?: return argumentType
     val simplifiedArgumentType = argumentType.lowerBoundIfFlexible() as? ConeTypeParameterType ?: return argumentType
-    val typeParameter = simplifiedArgumentType.lookupTag.typeParameterSymbol.fir
-
     val context = session.typeContext
 
-    val chosenSupertype = typeParameter.symbol.resolvedBounds.map { it.coneType }
+    val chosenSupertype = simplifiedArgumentType.collectUpperBounds()
         .singleOrNull { it.hasSupertypeWithGivenClassId(expectedTypeClassId, context) } ?: return argumentType
 
     val capturedType = context.captureFromExpression(chosenSupertype) ?: return argumentType

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/TypeUtils.kt

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir
+
+import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
+import org.jetbrains.kotlin.fir.types.*
+
+/**
+ * Collects the upper bounds as [ConeClassLikeType].
+ */
+fun ConeKotlinType?.collectUpperBounds(): Set<ConeClassLikeType> {
+    if (this == null) return emptySet()
+
+    val upperBounds = mutableSetOf<ConeClassLikeType>()
+    val seen = mutableSetOf<ConeKotlinType>()
+    fun collect(type: ConeKotlinType) {
+        if (!seen.add(type)) return // Avoid infinite recursion.
+
+        when (type) {
+            is ConeErrorType -> return // Ignore error types
+            is ConeLookupTagBasedType -> when (type) {
+                is ConeClassLikeType -> upperBounds.add(type)
+                is ConeTypeParameterType -> {
+                    val symbol = type.lookupTag.typeParameterSymbol
+                    symbol.resolvedBounds.forEach { collect(it.coneType) }
+                }
+                else -> error("missing branch for ${javaClass.name}")
+            }
+            is ConeTypeVariableType -> {
+                val symbol = (type.typeConstructor.originalTypeParameter as? ConeTypeParameterLookupTag)?.typeParameterSymbol ?: return
+                symbol.resolvedBounds.forEach { collect(it.coneType) }
+            }
+            is ConeDefinitelyNotNullType -> collect(type.original)
+            is ConeIntersectionType -> type.intersectedTypes.forEach(::collect)
+            is ConeFlexibleType -> collect(type.upperBound)
+            is ConeCapturedType -> type.constructor.supertypes?.forEach(::collect)
+            is ConeIntegerConstantOperatorType -> upperBounds.add(type.getApproximatedType())
+            is ConeStubType, is ConeIntegerLiteralConstantType -> {
+                error("$type should not reach here")
+            }
+        }
+    }
+
+    collect(this)
+    return upperBounds
+}

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

@@ -1,11 +0,0 @@
-interface Inv<T>
-
-fun <Y: X, X : Inv<out String>> foo(x: X, y: Y) {
-    val rX = bar(x)
-    rX.length
-
-    val rY = bar(<!ARGUMENT_TYPE_MISMATCH!>y<!>)
-    rY.<!UNRESOLVED_REFERENCE!>length<!>
-}
-
-fun <Y> bar(l: Inv<Y>): Y = TODO()

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 interface Inv<T>
 
 fun <Y: X, X : Inv<out String>> foo(x: X, y: Y) {