[FIR/NI] Refactor type variable gathering from lambda types

Motivation:
- drop getArguments from type context as a duplicate of getArgumentList
- reduce the number of collection allocations in getAllDeeplyRelatedTypeVariables

Additional minor improvements, test data fixes

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/ConeConstraintSystemUtilContext.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.resolve.inference.model.ConeFixVariableConstrain
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemUtilContext
 import org.jetbrains.kotlin.resolve.calls.inference.components.PostponedArgumentInputTypesResolver
 import org.jetbrains.kotlin.resolve.calls.inference.model.ArgumentConstraintPosition
@@ -56,7 +57,8 @@ object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
         return ConeFixVariableConstraintPosition(variable) as FixVariableConstraintPosition<T>
     }
 
-    override fun extractParameterTypesFromDeclaration(declaration: PostponedAtomWithRevisableExpectedType): List<ConeKotlinType?>? {
+    @OptIn(ExperimentalStdlibApi::class)
+    override fun extractLambdaParameterTypesFromDeclaration(declaration: PostponedAtomWithRevisableExpectedType): List<ConeKotlinType?>? {
         require(declaration is PostponedResolvedAtom)
         return when (declaration) {
             is LambdaWithTypeVariableAsExpectedTypeAtom -> {
@@ -66,7 +68,7 @@ object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
                         atom.collectDeclaredValueParameterTypes()
                     else null
                 } else { // function expression - all types are explicit, shouldn't return null
-                    mutableListOf<ConeKotlinType?>().apply {
+                    buildList {
                         atom.receiverTypeRef?.coneType?.let { add(it) }
                         addAll(atom.collectDeclaredValueParameterTypes())
                     }
@@ -77,11 +79,7 @@ object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
     }
 
     private fun FirAnonymousFunction.collectDeclaredValueParameterTypes(): List<ConeKotlinType?> =
-        valueParameters.map {
+        valueParameters.map { it.returnTypeRef.coneTypeSafe() }
-            if (it.returnTypeRef !is FirImplicitTypeRef)
-                it.returnTypeRef.coneType
-            else null
-        }
 
     override fun PostponedAtomWithRevisableExpectedType.isAnonymousFunction(): Boolean {
         require(this is PostponedResolvedAtom)
@@ -102,7 +100,7 @@ object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
         index: Int
     ): TypeVariableMarker {
         return ConeTypeVariableForPostponedAtom(
-            "${PostponedArgumentInputTypesResolver.TYPE_VARIABLE_NAME_PREFIX_FOR_LAMBDA_PARAMETER_TYPE}$index"
+            PostponedArgumentInputTypesResolver.TYPE_VARIABLE_NAME_PREFIX_FOR_LAMBDA_PARAMETER_TYPE + index
         )
     }
 
@@ -111,7 +109,7 @@ object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
         index: Int
     ): TypeVariableMarker {
         return ConeTypeVariableForPostponedAtom(
-            "${PostponedArgumentInputTypesResolver.TYPE_VARIABLE_NAME_PREFIX_FOR_CR_PARAMETER_TYPE}$index"
+            PostponedArgumentInputTypesResolver.TYPE_VARIABLE_NAME_PREFIX_FOR_CR_PARAMETER_TYPE + index
         )
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/ConstraintSystemCompleter.kt

@@ -158,10 +158,8 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
             argument.revisedExpectedType?.takeIf { it.isFunctionOrKFunctionWithAnySuspendability() }?.cast() ?: return false
 
         when (argument) {
-            is ResolvedCallableReferenceAtom -> {
+            is ResolvedCallableReferenceAtom ->
                 argument.reviseExpectedType(revisedExpectedType)
-            }
-
             is LambdaWithTypeVariableAsExpectedTypeAtom ->
                 argument.transformToResolvedLambda(c.getBuilder(), resolutionContext, revisedExpectedType, null /*TODO()*/)
             else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
@@ -188,7 +186,7 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
     }
 
     // Avoiding smart cast from filterIsInstanceOrNull looks dirty
-    private fun findPostponedArgumentWithRevisableExpectedType(postponedArguments: List<PostponedResolvedAtom>) =
+    private fun findPostponedArgumentWithRevisableExpectedType(postponedArguments: List<PostponedResolvedAtom>): PostponedResolvedAtom? =
         postponedArguments.firstOrNull { argument -> argument is PostponedAtomWithRevisableExpectedType }
 
     private fun ConstraintSystemCompletionContext.fixVariablesOrReportNotEnoughInformation(
@@ -263,9 +261,14 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
 
         // TODO: non-top-level variables?
         fun PostponedAtomWithRevisableExpectedType.collectNotFixedVariables() {
-            revisedExpectedType?.getArguments()?.map { it.getType().typeConstructor() }
+            revisedExpectedType?.lowerBoundIfFlexible()?.asArgumentList()?.let { typeArgumentList ->
-                ?.filterIsInstance<ConeTypeVariableTypeConstructor>().orEmpty()
+                for (typeArgument in typeArgumentList) {
-                .mapNotNullTo(result) { it.takeIf { it in notFixedTypeVariables } }
+                    val constructor = typeArgument.getType().typeConstructor()
+                    if (constructor in notFixedTypeVariables) {
+                        result.add(constructor)
+                    }
+                }
+            }
         }
 
         fun FirStatement.collectAllTypeVariables() {

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

@@ -214,8 +214,7 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
     override fun KotlinTypeMarker.isBuiltinFunctionalTypeOrSubtype(): Boolean {
         require(this is ConeKotlinType)
         return this.isTypeOrSubtypeOf {
-            it.lowerBoundIfFlexible().safeAs<ConeKotlinType>()?.isBuiltinFunctionalType(session)
+            (it.lowerBoundIfFlexible() as ConeKotlinType).isBuiltinFunctionalType(session)
-                ?: error("Unexpected lower bound for type: ${it.render()}")
         }
     }
 
@@ -402,7 +401,8 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
             this,
             {
                 // FIXME supertypes of type constructor contain unsubstituted arguments
-                it.typeConstructor().supertypes().cast<Collection<ConeKotlinType>>()
+                @Suppress("UNCHECKED_CAST")
+                it.typeConstructor().supertypes() as Collection<ConeKotlinType>
             },
             DFS.VisitedWithSet(),
             object : DFS.AbstractNodeHandler<ConeKotlinType, Boolean>() {
@@ -423,8 +423,7 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
     override fun KotlinTypeMarker.isSuspendFunctionTypeOrSubtype(): Boolean {
         require(this is ConeKotlinType)
         return isTypeOrSubtypeOf {
-            it.lowerBoundIfFlexible().safeAs<ConeKotlinType>()?.isSuspendFunctionType(session)
+            (it.lowerBoundIfFlexible() as ConeKotlinType).isSuspendFunctionType(session)
-                ?: error("Unexpected lower bound for type: ${it.render()}")
         }
     }
 
@@ -433,7 +432,7 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
         return this.lowerBoundIfFlexible().safeAs<ConeKotlinType>()?.isExtensionFunctionType(session) == true
     }
 
-    @ExperimentalStdlibApi
+    @OptIn(ExperimentalStdlibApi::class)
     override fun KotlinTypeMarker.extractArgumentsForFunctionalTypeOrSubtype(): List<KotlinTypeMarker> {
         val builtInFunctionalType = getFunctionalTypeFromSupertypes().cast<ConeKotlinType>()
         return buildList {
@@ -452,7 +451,7 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
 
         return fullyExpandedType(session).let {
             val simpleType = it.lowerBoundIfFlexible()
-            if (simpleType.safeAs<ConeKotlinType>()?.isBuiltinFunctionalType(session) == true)
+            if ((simpleType as ConeKotlinType).isBuiltinFunctionalType(session))
                 this
             else {
                 var functionalSupertype: KotlinTypeMarker? = null

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeTypeContext.kt

@@ -160,12 +160,6 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
             ?: session.builtinTypes.anyType.type//StandardClassIds.Any(session.firSymbolProvider).constructType(emptyArray(), false) // TODO wtf
     }
 
-    override fun KotlinTypeMarker.getArguments(): List<TypeArgumentMarker> {
-        require(this is ConeKotlinType)
-
-        return this.typeArguments.toList()
-    }
-
     override fun KotlinTypeMarker.asTypeArgument(): TypeArgumentMarker {
         require(this is ConeKotlinType)
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/types/IrTypeSystemContext.kt

@@ -94,12 +94,6 @@ interface IrTypeSystemContext : TypeSystemContext, TypeSystemCommonSuperTypesCon
             else -> error("Type $this has no arguments")
         }
 
-    override fun KotlinTypeMarker.getArguments(): List<TypeArgumentMarker> =
-        when (this) {
-            is IrSimpleType -> arguments
-            else -> error("Type $this has no arguments")
-        }
-
     override fun KotlinTypeMarker.asTypeArgument() = this as IrTypeArgument
 
     override fun CapturedTypeMarker.lowerType(): KotlinTypeMarker? = error("Captured Type is not valid for IrTypes")

compiler/resolution.common/src/org/jetbrains/kotlin/resolve/calls/inference/components/ConstraintSystemUtilContext.kt

@@ -27,7 +27,7 @@ interface ConstraintSystemUtilContext {
     // PostponedArgumentInputTypesResolver
     fun <T> createArgumentConstraintPosition(argument: T): ArgumentConstraintPosition<T>
     fun <T> createFixVariableConstraintPosition(variable: TypeVariableMarker, atom: T): FixVariableConstraintPosition<T>
-    fun extractParameterTypesFromDeclaration(declaration: PostponedAtomWithRevisableExpectedType): List<KotlinTypeMarker?>?
+    fun extractLambdaParameterTypesFromDeclaration(declaration: PostponedAtomWithRevisableExpectedType): List<KotlinTypeMarker?>?
     fun PostponedAtomWithRevisableExpectedType.isAnonymousFunction(): Boolean
     fun PostponedAtomWithRevisableExpectedType.isFunctionExpressionWithReceiver(): Boolean
     fun createTypeVariableForLambdaReturnType(): TypeVariableMarker

compiler/resolution.common/src/org/jetbrains/kotlin/resolve/calls/inference/components/PostponedArgumentInputTypesResolver.kt

@@ -332,7 +332,7 @@ class PostponedArgumentInputTypesResolver(
         for (argument in postponedArguments) {
             if (argument !is LambdaWithTypeVariableAsExpectedTypeMarker) continue
             if (argument.parameterTypesFromDeclaration != null) continue
-            argument.updateParameterTypesFromDeclaration(extractParameterTypesFromDeclaration(argument))
+            argument.updateParameterTypesFromDeclaration(extractLambdaParameterTypesFromDeclaration(argument))
         }
 
         return postponedArguments.any { argument ->
@@ -359,21 +359,33 @@ class PostponedArgumentInputTypesResolver(
 
     private fun Context.getAllDeeplyRelatedTypeVariables(
         type: KotlinTypeMarker,
-        variableDependencyProvider: TypeVariableDependencyInformationProvider
+        variableDependencyProvider: TypeVariableDependencyInformationProvider,
-    ): List<TypeVariableTypeConstructorMarker> {
+    ): Collection<TypeVariableTypeConstructorMarker> {
+        val collectedVariables = mutableSetOf<TypeVariableTypeConstructorMarker>()
+        getAllDeeplyRelatedTypeVariables(type, variableDependencyProvider, collectedVariables)
+        return collectedVariables
+    }
+
+    private fun Context.getAllDeeplyRelatedTypeVariables(
+        type: KotlinTypeMarker,
+        variableDependencyProvider: TypeVariableDependencyInformationProvider,
+        typeVariableCollector: MutableSet<TypeVariableTypeConstructorMarker>
+    ) {
         val typeConstructor = type.typeConstructor()
 
-        return when {
+        when {
             typeConstructor is TypeVariableTypeConstructorMarker -> {
                 val relatedVariables = variableDependencyProvider.getDeeplyDependentVariables(typeConstructor).orEmpty()
-                listOf(typeConstructor) + relatedVariables.filterIsInstance<TypeVariableTypeConstructorMarker>()
+                typeVariableCollector.add(typeConstructor)
+                typeVariableCollector.addAll(relatedVariables.filterIsInstance<TypeVariableTypeConstructorMarker>())
             }
             type.argumentsCount() > 0 -> {
-                type.getArguments().flatMap {
+                for (typeArgument in type.lowerBoundIfFlexible().asArgumentList()) {
-                    if (it.isStarProjection()) emptyList() else getAllDeeplyRelatedTypeVariables(it.getType(), variableDependencyProvider)
+                    if (!typeArgument.isStarProjection()) {
+                        getAllDeeplyRelatedTypeVariables(typeArgument.getType(), variableDependencyProvider, typeVariableCollector)
+                    }
                 }
             }
-            else -> emptyList()
         }
     }
 

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/components/ClassicConstraintSystemUtilContext.kt

@@ -63,7 +63,7 @@ class ClassicConstraintSystemUtilContext(
         return FixVariableConstraintPositionImpl(variable, atom) as FixVariableConstraintPosition<T>
     }
 
-    override fun extractParameterTypesFromDeclaration(declaration: PostponedAtomWithRevisableExpectedType): List<KotlinTypeMarker?>? {
+    override fun extractLambdaParameterTypesFromDeclaration(declaration: PostponedAtomWithRevisableExpectedType): List<KotlinTypeMarker?>? {
         require(declaration is ResolvedAtom)
         return when (val atom = declaration.atom) {
             is FunctionExpression -> {

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/components/KotlinConstraintSystemCompleter.kt

@@ -20,10 +20,8 @@ import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 class KotlinConstraintSystemCompleter(
     private val resultTypeResolver: ResultTypeResolver,
     val variableFixationFinder: VariableFixationFinder,
-    private val ctx: ConstraintSystemUtilContext,
+    private val postponedArgumentInputTypesResolver: PostponedArgumentInputTypesResolver,
 ) {
-    private val postponedArgumentInputTypesResolver = PostponedArgumentInputTypesResolver(resultTypeResolver, variableFixationFinder, ctx)
-
     fun runCompletion(
         c: ConstraintSystemCompletionContext,
         completionMode: ConstraintSystemCompletionMode,

compiler/testData/codegen/box/closures/capturedVarsOptimization/capturedInInlineOnlyIndexedCAO.kt

@@ -1,3 +1,4 @@
+// IGNORE_BACKEND_FIR: JVM_IR
 // WITH_RUNTIME
 
 class Host(var value: String) {

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/extensionsWithNonValuableConstraints_1_2.fir.kt

@@ -0,0 +1,17 @@
+// !LANGUAGE: -ExperimentalBuilderInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+
+interface Base
+
+interface Controller<T> : Base {
+    suspend fun yield(t: T) {}
+}
+
+fun <S> generate(g: suspend Controller<S>.() -> Unit): S = TODO()
+
+suspend fun Base.baseExtension() {}
+
+val test1 = generate {
+    <!INAPPLICABLE_CANDIDATE!>yield<!>("foo")
+    baseExtension()
+}

compiler/testData/diagnostics/testsWithStdLib/coroutines/inference/extensionsWithNonValuableConstraints_1_2.kt

@@ -1,4 +1,3 @@
-// FIR_IDENTICAL
 // !LANGUAGE: -ExperimentalBuilderInference
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 
@@ -13,6 +12,6 @@ fun <S> generate(g: suspend Controller<S>.() -> Unit): S = TODO()
 suspend fun Base.baseExtension() {}
 
 val test1 = generate {
-    <!INAPPLICABLE_CANDIDATE!>yield<!>("foo")
+    yield("foo")
     baseExtension()
 }

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

@@ -246,7 +246,6 @@ interface TypeSystemContext : TypeSystemOptimizationContext {
 
     fun KotlinTypeMarker.argumentsCount(): Int
     fun KotlinTypeMarker.getArgument(index: Int): TypeArgumentMarker
-    fun KotlinTypeMarker.getArguments(): List<TypeArgumentMarker>
 
     fun SimpleTypeMarker.getArgumentOrNull(index: Int): TypeArgumentMarker? {
         if (index in 0 until argumentsCount()) return getArgument(index)
@@ -341,6 +340,18 @@ interface TypeSystemContext : TypeSystemOptimizationContext {
         }
     }
 
+    operator fun TypeArgumentListMarker.iterator() = object : Iterator<TypeArgumentMarker> {
+        private var argumentIndex: Int = 0
+
+        override fun hasNext(): Boolean = argumentIndex < size()
+
+        override fun next(): TypeArgumentMarker {
+            val argument = get(argumentIndex)
+            argumentIndex += 1
+            return argument
+        }
+    }
+
     fun TypeConstructorMarker.isAnyConstructor(): Boolean
     fun TypeConstructorMarker.isNothingConstructor(): Boolean
 

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

@@ -159,11 +159,6 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
         return this.arguments[index]
     }
 
-    override fun KotlinTypeMarker.getArguments(): List<TypeArgumentMarker> {
-        require(this is KotlinType, this::errorMessage)
-        return this.arguments
-    }
-
     override fun TypeArgumentMarker.isStarProjection(): Boolean {
         require(this is TypeProjection, this::errorMessage)
         return this.isStarProjection