Reuse built functional types for postponed arguments by expected types and paths from a top level type variable

^KT-42221 Fixed

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

@@ -3234,6 +3234,49 @@ public class FirOldFrontendDiagnosticsTestWithStdlibGenerated extends AbstractFi
                 public void testSuspendFunctions() throws Exception {
                     runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/suspendFunctions.kt");
                 }
+
+                @TestMetadata("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance")
+                @TestDataPath("$PROJECT_ROOT")
+                @RunWith(JUnit3RunnerWithInners.class)
+                public static class Performance extends AbstractFirOldFrontendDiagnosticsTestWithStdlib {
+                    private void runTest(String testDataFilePath) throws Exception {
+                        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+                    }
+
+                    public void testAllFilesPresentInPerformance() throws Exception {
+                        KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForIdLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForIdLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForIdLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairOfLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairOfLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairOfLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairsOfDeepLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairsOfDeepLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfDeepLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairsOfDeepMixedLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairsOfDeepMixedLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfDeepMixedLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairsOfIdLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairsOfIdLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfIdLambdas.kt");
+                    }
+                }
             }
         }
 

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

@@ -37,6 +37,7 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
         collectVariablesFromContext: Boolean = false,
         analyze: (PostponedResolvedAtom) -> Unit
     ) = with(c) {
+        val topLevelTypeVariables = candidateReturnType.extractTypeVariables()
 
         completion@ while (true) {
             val postponedArguments = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
@@ -62,7 +63,11 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
 
             // Stage 2
             val newExpectedTypeWasBuilt = postponedArgumentsInputTypesResolver.collectParameterTypesAndBuildNewExpectedTypes(
-                asConstraintSystemCompletionContext(), postponedArgumentsWithRevisableType, completionMode, dependencyProvider
+                asConstraintSystemCompletionContext(),
+                postponedArgumentsWithRevisableType,
+                completionMode,
+                dependencyProvider,
+                topLevelTypeVariables
             )
 
             if (newExpectedTypeWasBuilt)

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

@@ -17,11 +17,23 @@ interface ConstraintSystemOperation {
     fun unmarkPostponedVariable(variable: TypeVariableMarker)
     fun removePostponedVariables()
 
-    fun getRevisedVariableForParameter(expectedType: TypeVariableMarker, index: Int): TypeVariableMarker?
+    fun getBuiltFunctionalExpectedTypeForPostponedArgument(
-    fun getRevisedVariableForReturnType(expectedType: TypeVariableMarker): TypeVariableMarker?
+        topLevelVariable: TypeConstructorMarker,
+        pathToExpectedType: List<Pair<TypeConstructorMarker, Int>>
+    ): KotlinTypeMarker?
 
-    fun putRevisedVariableForParameter(expectedType: TypeVariableMarker, index: Int, newVariable: TypeVariableMarker)
+    fun getBuiltFunctionalExpectedTypeForPostponedArgument(expectedTypeVariable: TypeConstructorMarker): KotlinTypeMarker?
-    fun putRevisedVariableForReturnType(expectedType: TypeVariableMarker, newVariable: TypeVariableMarker)
+
+    fun putBuiltFunctionalExpectedTypeForPostponedArgument(
+        topLevelVariable: TypeConstructorMarker,
+        pathToExpectedType: List<Pair<TypeConstructorMarker, Int>>,
+        builtFunctionalType: KotlinTypeMarker
+    )
+
+    fun putBuiltFunctionalExpectedTypeForPostponedArgument(
+        expectedTypeVariable: TypeConstructorMarker,
+        builtFunctionalType: KotlinTypeMarker
+    )
 
     fun addSubtypeConstraint(lowerType: KotlinTypeMarker, upperType: KotlinTypeMarker, position: ConstraintPosition)
     fun addEqualityConstraint(a: KotlinTypeMarker, b: KotlinTypeMarker, position: ConstraintPosition)

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.types.model.*
 import org.jetbrains.kotlin.utils.SmartSet
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+import java.util.*
 
 private typealias Context = ConstraintSystemCompletionContext
 private typealias ResolvedAtomProvider = (TypeVariableMarker) -> Any?
@@ -133,48 +134,22 @@ class PostponedArgumentInputTypesResolver(
 
     private fun Context.createTypeVariableForReturnType(argument: PostponedAtomWithRevisableExpectedType): TypeVariableMarker =
         with(resolutionTypeSystemContext) {
-            val expectedType = argument.expectedType
-                ?: throw IllegalStateException("Postponed argument's expected type must not be null")
-
-            val variable = getBuilder().currentStorage().allTypeVariables[expectedType.typeConstructor()]
-            if (variable != null) {
-                val revisedVariableForReturnType = getBuilder().getRevisedVariableForReturnType(variable)
-                if (revisedVariableForReturnType != null) return revisedVariableForReturnType
-            }
-
             return when (argument) {
                 is LambdaWithTypeVariableAsExpectedTypeMarker -> createTypeVariableForLambdaReturnType()
                 is PostponedCallableReferenceMarker -> createTypeVariableForCallableReferenceReturnType()
                 else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
-            }.also {
+            }.also { getBuilder().registerVariable(it) }
-                if (variable != null) getBuilder().putRevisedVariableForReturnType(variable, it)
-
-                getBuilder().registerVariable(it)
-            }
         }
 
     private fun Context.createTypeVariableForParameterType(
         argument: PostponedAtomWithRevisableExpectedType,
         index: Int
     ): TypeVariableMarker = with(resolutionTypeSystemContext) {
-        val expectedType = argument.expectedType
-            ?: throw IllegalStateException("Postponed argument's expected type must not be null")
-
-        val variable = getBuilder().currentStorage().allTypeVariables[expectedType.typeConstructor()]
-        if (variable != null) {
-            val revisedVariableForParameter = getBuilder().getRevisedVariableForParameter(variable, index)
-            if (revisedVariableForParameter != null) return revisedVariableForParameter
-        }
-
         return when (argument) {
             is LambdaWithTypeVariableAsExpectedTypeMarker -> createTypeVariableForLambdaParameterType(argument, index)
             is PostponedCallableReferenceMarker -> createTypeVariableForCallableReferenceParameterType(argument, index)
             else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
-        }.also {
+        }.also { getBuilder().registerVariable(it) }
-            if (variable != null) getBuilder().putRevisedVariableForParameter(variable, index, it)
-
-            getBuilder().registerVariable(it)
-        }
     }
 
     private fun Context.createTypeVariablesForParameters(
@@ -243,14 +218,73 @@ class PostponedArgumentInputTypesResolver(
         }
     }
 
+    private fun Context.computeTypeVariablePathInsideGivenType(
+        type: KotlinTypeMarker,
+        targetVariable: TypeConstructorMarker,
+        path: Stack<Pair<TypeConstructorMarker, Int>> = Stack()
+    ): List<Pair<TypeConstructorMarker, Int>>? {
+        val typeConstructor = type.typeConstructor()
+
+        if (typeConstructor == targetVariable)
+            return emptyList()
+
+        for (i in 0 until type.argumentsCount()) {
+            val argumentType = type.getArgument(i).getType()
+
+            if (argumentType.typeConstructor() == targetVariable) {
+                return path.toList() + (typeConstructor to i)
+            } else if (argumentType.argumentsCount() != 0) {
+                path.push(typeConstructor to i)
+                computeTypeVariablePathInsideGivenType(argumentType, targetVariable, path)?.let { return it }
+                path.pop()
+            }
+        }
+
+        return null
+    }
+
+    private fun Context.selectFirstRelatedVariable(
+        variables: Set<TypeVariableTypeConstructorMarker>,
+        targetVariable: TypeConstructorMarker,
+        variableDependencyProvider: TypeVariableDependencyInformationProvider
+    ): TypeVariableTypeConstructorMarker? {
+        val relatedVariables = variableDependencyProvider.getDeeplyDependentVariables(targetVariable).orEmpty() +
+                variableDependencyProvider.getShallowlyDependentVariables(targetVariable)
+
+        return variables.firstOrNull { it in relatedVariables && it in notFixedTypeVariables }
+    }
+
     private fun Context.buildNewFunctionalExpectedType(
         argument: PostponedAtomWithRevisableExpectedType,
-        parameterTypesInfo: ParameterTypesInfo
+        parameterTypesInfo: ParameterTypesInfo,
+        variableDependencyProvider: TypeVariableDependencyInformationProvider,
+        topLevelTypeVariables: Set<TypeVariableTypeConstructorMarker>
     ): KotlinTypeMarker? = with(resolutionTypeSystemContext) {
-        val expectedType = argument.expectedType
+        val expectedType = argument.expectedType ?: return null
+        val expectedTypeConstructor = expectedType.typeConstructor()
 
-        if (expectedType == null || expectedType.typeConstructor() !in notFixedTypeVariables)
+        if (expectedTypeConstructor !in notFixedTypeVariables) return null
-            return null
+
+        val relatedTopLevelVariable = selectFirstRelatedVariable(topLevelTypeVariables, expectedTypeConstructor, variableDependencyProvider)
+        val pathFromRelatedTopLevelVariable = if (relatedTopLevelVariable != null) {
+            val constraintTypes = notFixedTypeVariables.getValue(relatedTopLevelVariable).constraints.map { it.type }.toSet()
+            val containingType = constraintTypes.find { constraintType ->
+                constraintType.contains { it.typeConstructor() == expectedTypeConstructor }
+            }
+            if (containingType != null) {
+                computeTypeVariablePathInsideGivenType(containingType, expectedTypeConstructor)
+            } else null
+        } else null
+
+        if (pathFromRelatedTopLevelVariable != null && relatedTopLevelVariable != null) {
+            // try to take from the cache of functional types by paths from a top level type variable
+            getBuilder().getBuiltFunctionalExpectedTypeForPostponedArgument(relatedTopLevelVariable, pathFromRelatedTopLevelVariable)
+                ?.let { return it }
+        } else {
+            // try to take from the cache of functional types by expected types
+            getBuilder().getBuiltFunctionalExpectedTypeForPostponedArgument(expectedTypeConstructor)
+                ?.let { return it }
+        }
 
         val parametersFromConstraints = parameterTypesInfo.parametersFromConstraints
         val parametersFromDeclaration = getDeclaredParametersConsideringExtensionFunctionsPresence(parameterTypesInfo)
@@ -319,6 +353,12 @@ class PostponedArgumentInputTypesResolver(
             createArgumentConstraintPosition(argument)
         )
 
+        if (pathFromRelatedTopLevelVariable != null && relatedTopLevelVariable != null) {
+            getBuilder().putBuiltFunctionalExpectedTypeForPostponedArgument(relatedTopLevelVariable, pathFromRelatedTopLevelVariable, newExpectedType)
+        } else {
+            getBuilder().putBuiltFunctionalExpectedTypeForPostponedArgument(expectedTypeConstructor, newExpectedType)
+        }
+
         return newExpectedType
     }
 
@@ -326,7 +366,8 @@ class PostponedArgumentInputTypesResolver(
         c: Context,
         postponedArguments: List<PostponedAtomWithRevisableExpectedType>,
         completionMode: ConstraintSystemCompletionMode,
-        dependencyProvider: TypeVariableDependencyInformationProvider
+        dependencyProvider: TypeVariableDependencyInformationProvider,
+        topLevelTypeVariables: Set<TypeVariableTypeConstructorMarker>
     ): Boolean = with(resolutionTypeSystemContext) {
         // We can collect parameter types from declaration in any mode, they can't change during completion.
         for (argument in postponedArguments) {
@@ -349,7 +390,7 @@ class PostponedArgumentInputTypesResolver(
             val parameterTypesInfo =
                 c.extractParameterTypesInfo(argument, postponedArguments, dependencyProvider) ?: return@any false
             val newExpectedType =
-                c.buildNewFunctionalExpectedType(argument, parameterTypesInfo) ?: return@any false
+                c.buildNewFunctionalExpectedType(argument, parameterTypesInfo, dependencyProvider, topLevelTypeVariables) ?: return@any false
 
             argument.reviseExpectedType(newExpectedType)
 

compiler/resolution.common/src/org/jetbrains/kotlin/resolve/calls/inference/model/ConstraintStorage.kt

@@ -39,8 +39,8 @@ interface ConstraintStorage {
     val hasContradiction: Boolean
     val fixedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker>
     val postponedTypeVariables: List<TypeVariableMarker>
-    val revisedVariablesForParameters: Map<Pair<TypeVariableMarker, Int>, TypeVariableMarker>
+    val builtFunctionalTypesForPostponedArgumentsByTopLevelTypeVariables: Map<Pair<TypeConstructorMarker, List<Pair<TypeConstructorMarker, Int>>>, KotlinTypeMarker>
-    val revisedReturnTypes: Map<TypeVariableMarker, TypeVariableMarker>
+    val builtFunctionalTypesForPostponedArgumentsByExpectedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker>
 
     object Empty : ConstraintStorage {
         override val allTypeVariables: Map<TypeConstructorMarker, TypeVariableMarker> get() = emptyMap()
@@ -51,8 +51,8 @@ interface ConstraintStorage {
         override val hasContradiction: Boolean get() = false
         override val fixedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker> get() = emptyMap()
         override val postponedTypeVariables: List<TypeVariableMarker> get() = emptyList()
-        override val revisedVariablesForParameters: Map<Pair<TypeVariableMarker, Int>, TypeVariableMarker> = emptyMap()
+        override val builtFunctionalTypesForPostponedArgumentsByTopLevelTypeVariables: Map<Pair<TypeConstructorMarker, List<Pair<TypeConstructorMarker, Int>>>, KotlinTypeMarker> = emptyMap()
-        override val revisedReturnTypes: Map<TypeVariableMarker, TypeVariableMarker> = emptyMap()
+        override val builtFunctionalTypesForPostponedArgumentsByExpectedTypeVariables: Map<TypeConstructorMarker, KotlinTypeMarker> = emptyMap()
     }
 }
 

compiler/resolution.common/src/org/jetbrains/kotlin/resolve/calls/inference/model/MutableConstraintStorage.kt

@@ -202,6 +202,8 @@ internal class MutableConstraintStorage : ConstraintStorage {
     override val hasContradiction: Boolean get() = errors.any { !it.applicability.isSuccess }
     override val fixedTypeVariables: MutableMap<TypeConstructorMarker, KotlinTypeMarker> = LinkedHashMap()
     override val postponedTypeVariables: MutableList<TypeVariableMarker> = SmartList()
-    override val revisedVariablesForParameters: MutableMap<Pair<TypeVariableMarker, Int>, TypeVariableMarker> = LinkedHashMap()
+    override val builtFunctionalTypesForPostponedArgumentsByTopLevelTypeVariables: MutableMap<Pair<TypeConstructorMarker, List<Pair<TypeConstructorMarker, Int>>>, KotlinTypeMarker> =
-    override val revisedReturnTypes: MutableMap<TypeVariableMarker, TypeVariableMarker> = LinkedHashMap()
+        LinkedHashMap()
+    override val builtFunctionalTypesForPostponedArgumentsByExpectedTypeVariables: MutableMap<TypeConstructorMarker, KotlinTypeMarker> =
+        LinkedHashMap()
 }

compiler/resolution.common/src/org/jetbrains/kotlin/resolve/calls/inference/model/NewConstraintSystemImpl.kt

@@ -115,21 +115,28 @@ class NewConstraintSystemImpl(
         storage.postponedTypeVariables.clear()
     }
 
-    override fun getRevisedVariableForParameter(expectedType: TypeVariableMarker, index: Int): TypeVariableMarker? {
+    override fun putBuiltFunctionalExpectedTypeForPostponedArgument(
-        return storage.revisedVariablesForParameters[expectedType to index]
+        topLevelVariable: TypeConstructorMarker,
+        pathToExpectedType: List<Pair<TypeConstructorMarker, Int>>,
+        builtFunctionalType: KotlinTypeMarker
+    ) {
+        storage.builtFunctionalTypesForPostponedArgumentsByTopLevelTypeVariables[topLevelVariable to pathToExpectedType] = builtFunctionalType
     }
 
-    override fun getRevisedVariableForReturnType(expectedType: TypeVariableMarker): TypeVariableMarker? {
+    override fun putBuiltFunctionalExpectedTypeForPostponedArgument(
-        return storage.revisedReturnTypes[expectedType]
+        expectedTypeVariable: TypeConstructorMarker,
+        builtFunctionalType: KotlinTypeMarker
+    ) {
+        storage.builtFunctionalTypesForPostponedArgumentsByExpectedTypeVariables[expectedTypeVariable] = builtFunctionalType
     }
 
-    override fun putRevisedVariableForParameter(expectedType: TypeVariableMarker, index: Int, newVariable: TypeVariableMarker) {
+    override fun getBuiltFunctionalExpectedTypeForPostponedArgument(
-        storage.revisedVariablesForParameters[expectedType to index] = newVariable
+        topLevelVariable: TypeConstructorMarker,
-    }
+        pathToExpectedType: List<Pair<TypeConstructorMarker, Int>>
+    ) = storage.builtFunctionalTypesForPostponedArgumentsByTopLevelTypeVariables[topLevelVariable to pathToExpectedType]
 
-    override fun putRevisedVariableForReturnType(expectedType: TypeVariableMarker, newVariable: TypeVariableMarker) {
+    override fun getBuiltFunctionalExpectedTypeForPostponedArgument(expectedTypeVariable: TypeConstructorMarker) =
-        storage.revisedReturnTypes[expectedType] = newVariable
+        storage.builtFunctionalTypesForPostponedArgumentsByExpectedTypeVariables[expectedTypeVariable]
-    }
 
     override fun addSubtypeConstraint(lowerType: KotlinTypeMarker, upperType: KotlinTypeMarker, position: ConstraintPosition) =
         constraintInjector.addInitialSubtypeConstraint(

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

@@ -66,6 +66,8 @@ class KotlinConstraintSystemCompleter(
         collectVariablesFromContext: Boolean,
         analyze: (PostponedResolvedAtom) -> Unit
     ) {
+        val topLevelTypeVariables = topLevelType.extractTypeVariables()
+
         completion@ while (true) {
             // TODO
             val postponedArguments = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
@@ -89,7 +91,11 @@ class KotlinConstraintSystemCompleter(
 
             // Stage 2: collect parameter types for postponed arguments
             val wasBuiltNewExpectedTypeForSomeArgument = postponedArgumentInputTypesResolver.collectParameterTypesAndBuildNewExpectedTypes(
-                asConstraintSystemCompletionContext(), postponedArgumentsWithRevisableType, completionMode, dependencyProvider
+                asConstraintSystemCompletionContext(),
+                postponedArgumentsWithRevisableType,
+                completionMode,
+                dependencyProvider,
+                topLevelTypeVariables
             )
 
             if (wasBuiltNewExpectedTypeForSomeArgument)

compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/basic.fir.kt

@@ -143,7 +143,7 @@ fun main() {
 
     // Resolution of extension/non-extension functions combination
     val x19: String.() -> Unit = select(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.String>")!>id { <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String"), DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>this<!> }<!>, <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.Unit>")!>id(fun(x: String) {})<!>)
-    val x20: String.() -> Unit = select(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.Unit>")!>{ <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String"), DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>this<!> }<!>, (fun(x: String) {}))
+    val x20: String.() -> Unit = select(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.String>")!>{ <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String"), DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>this<!> }<!>, (fun(x: String) {}))
     val x21: String.() -> Unit = select(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.Unit>")!>id(fun(x: String) {})<!>, <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.Unit>")!>id(fun(x: String) {})<!>)
     select(id<String.() -> Unit>(fun(x: String) {}), <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.Unit>")!>id(fun(x: String) {})<!>)
     select(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function2<kotlin.String, kotlin.String, kotlin.Unit>")!>id(fun String.(x: String) {})<!>, <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function2<kotlin.String, kotlin.String, kotlin.Unit>")!>id(fun(x: String, y: String) {})<!>)

compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForIdLambdas.kt

@@ -0,0 +1,40 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+// !DIAGNOSTICS: -UNUSED_ANONYMOUS_PARAMETER
+
+fun <T> id(x: T) = x
+
+class A {
+    val x: Map<String, (String, String, String, String) -> Unit> =
+        mapOf(
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> },
+            "" to id { a, b, c, d -> }
+        )
+}

compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForLambdas.kt

@@ -0,0 +1,38 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+// !DIAGNOSTICS: -UNUSED_ANONYMOUS_PARAMETER
+
+class A {
+    val x: Map<String, (String, String, String, String) -> Unit> =
+        mapOf(
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> },
+            "" to { a, b, c, d -> }
+        )
+}

compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairOfLambdas.kt

@@ -0,0 +1,42 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_ANONYMOUS_PARAMETER -CAST_NEVER_SUCCEEDS
+
+interface Foo<out K, out V> {}
+
+fun <K, V> fooOf(vararg pairs: Pair<K, V>) = null as Foo<K, V>
+
+fun <A, B> to(th: A, that: B): Pair<A, B> = Pair(th, that)
+
+class A {
+    val x: Foo<(Int, String, String, String) -> Unit, (String, String, String, String) -> Unit> =
+        fooOf(
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+            to({ a, b, c, d -> }, { a, b, c, d -> }),
+        )
+}

compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfDeepLambdas.kt

@@ -0,0 +1,46 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_ANONYMOUS_PARAMETER -CAST_NEVER_SUCCEEDS
+
+fun <T> id(x: T) = x
+
+interface Foo<out K, out V> {}
+
+fun <K, V> fooOf(vararg pairs: Pair<K, V>) = null as Foo<K, V>
+
+fun <A, B> to(th: A, that: B): Pair<A, B> = Pair(th, that)
+
+class Inv<K>(x: K)
+
+class A {
+    val x: Foo<Inv<(Int) -> Unit>, Inv<(String, String, String, String) -> Unit>> =
+        fooOf(
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+            to(Inv { a -> }, Inv { a, b, c, d -> }),
+        )
+}

compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfDeepMixedLambdas.kt

@@ -0,0 +1,48 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+// !DIAGNOSTICS: -UNUSED_ANONYMOUS_PARAMETER -UNUSED_PARAMETER -CAST_NEVER_SUCCEEDS
+
+fun <T> id(x: T) = x
+
+interface Foo<out K, out V> {}
+
+fun <K, V> fooOf(vararg pairs: Pair<K, V>) = null as Foo<K, V>
+
+fun <A, B> to(th: A, that: B): Pair<A, B> = Pair(th, that)
+
+class Inv<K>(x: K)
+
+fun <A: (Int, Int, Int, Int) -> Unit> take(x: A, y: A) = null as (Int, String, String, String) -> Unit
+
+class A {
+    val x: Foo<Inv<(Int, String, String, String) -> Unit>, Inv<(Int, String, String, String) -> Unit>> =
+        fooOf(
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv(take({ a, b, c, d -> }, { a, b, c, d -> }))),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv(take({ a, b, c, d -> }, { a, b, c, d -> }))),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv(take({ a, b, c, d -> }, { a, b, c, d -> }))),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv(take({ a, b, c, d -> }, { a, b, c, d -> }))),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv(take({ a, b, c, d -> }, { a, b, c, d -> }))),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+            to(Inv(take({ a, b, c, d -> }, { a, b, c, d -> })), Inv { a, b, c, d -> }),
+        )
+}

compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfIdLambdas.kt

@@ -0,0 +1,44 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_ANONYMOUS_PARAMETER -CAST_NEVER_SUCCEEDS
+
+fun <T> id(x: T) = x
+
+interface Foo<out K, out V> {}
+
+fun <K, V> fooOf(vararg pairs: Pair<K, V>) = null as Foo<K, V>
+
+fun <A, B> to(th: A, that: B): Pair<A, B> = Pair(th, that)
+
+class A {
+    val x: Foo<(Int, String, String, String) -> Unit, (String, String, String, String) -> Unit> =
+        fooOf(
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+            to(id { a, b, c, d -> }, id { a, b, c, d -> }),
+        )
+}

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

@@ -3384,6 +3384,49 @@ public class DiagnosticsTestWithStdLibGenerated extends AbstractDiagnosticsTestW
                 public void testSuspendFunctions() throws Exception {
                     runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/suspendFunctions.kt");
                 }
+
+                @TestMetadata("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance")
+                @TestDataPath("$PROJECT_ROOT")
+                @RunWith(JUnit3RunnerWithInners.class)
+                public static class Performance extends AbstractDiagnosticsTestWithStdLib {
+                    private void runTest(String testDataFilePath) throws Exception {
+                        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+                    }
+
+                    public void testAllFilesPresentInPerformance() throws Exception {
+                        KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForIdLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForIdLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForIdLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairOfLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairOfLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairOfLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairsOfDeepLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairsOfDeepLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfDeepLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairsOfDeepMixedLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairsOfDeepMixedLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfDeepMixedLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairsOfIdLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairsOfIdLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfIdLambdas.kt");
+                    }
+                }
             }
         }
 

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

@@ -3384,6 +3384,49 @@ public class DiagnosticsTestWithStdLibUsingJavacGenerated extends AbstractDiagno
                 public void testSuspendFunctions() throws Exception {
                     runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/suspendFunctions.kt");
                 }
+
+                @TestMetadata("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance")
+                @TestDataPath("$PROJECT_ROOT")
+                @RunWith(JUnit3RunnerWithInners.class)
+                public static class Performance extends AbstractDiagnosticsTestWithStdLibUsingJavac {
+                    private void runTest(String testDataFilePath) throws Exception {
+                        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+                    }
+
+                    public void testAllFilesPresentInPerformance() throws Exception {
+                        KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForIdLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForIdLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForIdLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairOfLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairOfLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairOfLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairsOfDeepLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairsOfDeepLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfDeepLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairsOfDeepMixedLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairsOfDeepMixedLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfDeepMixedLambdas.kt");
+                    }
+
+                    @TestMetadata("reuseBuiltFunctionalTypesForPairsOfIdLambdas.kt")
+                    public void testReuseBuiltFunctionalTypesForPairsOfIdLambdas() throws Exception {
+                        runTest("compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/performance/reuseBuiltFunctionalTypesForPairsOfIdLambdas.kt");
+                    }
+                }
             }
         }
 

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

@@ -7,6 +7,8 @@ package org.jetbrains.kotlin.types.model
 
 import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
 import org.jetbrains.kotlin.types.Variance
+import java.util.*
+import kotlin.collections.ArrayList
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.contract
 
@@ -207,6 +209,25 @@ interface TypeSystemInferenceExtensionContext : TypeSystemContext, TypeSystemBui
     fun getFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker
 
     fun getKFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker
+
+    private fun KotlinTypeMarker.extractTypeVariables(to: MutableSet<TypeVariableTypeConstructorMarker>) {
+        for (i in 0 until argumentsCount()) {
+            val argument = getArgument(i)
+
+            if (argument.isStarProjection()) continue
+
+            val argumentType = argument.getType()
+            val argumentTypeConstructor = argumentType.typeConstructor()
+            if (argumentTypeConstructor is TypeVariableTypeConstructorMarker) {
+                to.add(argumentTypeConstructor)
+            } else if (argumentType.argumentsCount() != 0) {
+                argumentType.extractTypeVariables(to)
+            }
+        }
+    }
+
+    @OptIn(ExperimentalStdlibApi::class)
+    fun KotlinTypeMarker.extractTypeVariables() = buildSet { extractTypeVariables(this) }
 }