[FIR] Prepare commonization of PostponedArgumentInputTypesResolver

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

@@ -7,7 +7,12 @@ package org.jetbrains.kotlin.fir.resolve.inference
 
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemUtilContext
+import org.jetbrains.kotlin.resolve.calls.inference.model.ArgumentConstraintPosition
+import org.jetbrains.kotlin.resolve.calls.inference.model.FixVariableConstraintPosition
+import org.jetbrains.kotlin.resolve.calls.model.LambdaWithTypeVariableAsExpectedTypeMarker
+import org.jetbrains.kotlin.resolve.calls.model.PostponedAtomWithRevisableExpectedType
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
 object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
@@ -35,4 +40,68 @@ object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
     override fun KotlinTypeMarker.refineType(): KotlinTypeMarker {
         return this
     }
+
+    override fun KotlinTypeMarker.isFunctionOrKFunctionTypeWithAnySuspendability(): Boolean {
+        TODO()
+    }
+
+    override fun KotlinTypeMarker.isSuspendFunctionTypeOrSubtype(): Boolean {
+        TODO()
+    }
+
+    override fun extractFunctionalTypeFromSupertypes(type: KotlinTypeMarker): KotlinTypeMarker {
+        TODO()
+    }
+
+    override fun KotlinTypeMarker.extractArgumentsForFunctionalTypeOrSubtype(): List<KotlinTypeMarker> {
+        TODO()
+    }
+
+    override fun <T> createArgumentConstraintPosition(argument: T): ArgumentConstraintPosition<T> {
+        TODO()
+    }
+
+    override fun <T> createFixVariableConstraintPosition(variable: TypeVariableMarker, atom: T): FixVariableConstraintPosition<T> {
+        TODO()
+    }
+
+    override fun extractParameterTypesFromDeclaration(declaration: PostponedAtomWithRevisableExpectedType): List<ConeKotlinType?>? {
+        TODO()
+    }
+
+    override fun KotlinTypeMarker.isExtensionFunctionType(): Boolean {
+        TODO()
+    }
+
+    override fun getFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker {
+        TODO()
+    }
+
+    override fun getKFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker {
+        TODO()
+    }
+
+    override fun isAnonymousFunction(argument: PostponedAtomWithRevisableExpectedType): Boolean {
+        TODO()
+    }
+
+    override fun PostponedAtomWithRevisableExpectedType.isFunctionExpressionWithReceiver(): Boolean {
+        TODO()
+    }
+
+    override fun createTypeVariableForLambdaReturnType(): TypeVariableMarker {
+        TODO()
+    }
+
+    override fun createTypeVariableForLambdaParameterType(argument: PostponedAtomWithRevisableExpectedType, index: Int): TypeVariableMarker {
+        TODO()
+    }
+
+    override fun createTypeVariableForCallableReferenceParameterType(argument: PostponedAtomWithRevisableExpectedType, index: Int): TypeVariableMarker {
+        TODO()
+    }
+
+    override fun createTypeVariableForCallableReferenceReturnType(): TypeVariableMarker {
+        TODO()
+    }
 }

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

@@ -17,8 +17,10 @@ import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeTypeVariable
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.resolve.calls.model.PostponedAtomWithRevisableExpectedType
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 
 //  --------------------------- Variables ---------------------------
 
@@ -30,7 +32,7 @@ sealed class PostponedResolvedAtom : PostponedResolvedAtomMarker {
     abstract override val inputTypes: Collection<ConeKotlinType>
     abstract override val outputType: ConeKotlinType?
     override var analyzed: Boolean = false
-    abstract val expectedType: ConeKotlinType?
+    abstract override val expectedType: ConeKotlinType?
 }
 
 //  ------------- Lambdas -------------
@@ -62,13 +64,20 @@ class LambdaWithTypeVariableAsExpectedTypeAtom(
     override val expectedType: ConeKotlinType,
     val expectedTypeRef: FirTypeRef,
     val candidateOfOuterCall: Candidate
-) : PostponedResolvedAtom() {
+) : PostponedResolvedAtom(), PostponedAtomWithRevisableExpectedType {
     init {
         candidateOfOuterCall.postponedAtoms += this
     }
 
     override val inputTypes: Collection<ConeKotlinType> get() = listOf(expectedType)
     override val outputType: ConeKotlinType? get() = null
+    override var revisedExpectedType: ConeKotlinType? = null
+        private set
+
+    override fun reviseExpectedType(expectedType: KotlinTypeMarker) {
+        require(expectedType is ConeKotlinType)
+        revisedExpectedType = expectedType
+    }
 }
 
 //  ------------- References -------------

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

@@ -160,6 +160,12 @@ 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,6 +94,12 @@ 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

@@ -5,7 +5,11 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference.components
 
+import org.jetbrains.kotlin.resolve.calls.inference.model.ArgumentConstraintPosition
+import org.jetbrains.kotlin.resolve.calls.inference.model.FixVariableConstraintPosition
+import org.jetbrains.kotlin.resolve.calls.model.PostponedAtomWithRevisableExpectedType
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 
 /*
@@ -19,4 +23,25 @@ interface ConstraintSystemUtilContext {
     fun KotlinTypeMarker.unCapture(): KotlinTypeMarker
     fun TypeVariableMarker.isReified(): Boolean
     fun KotlinTypeMarker.refineType(): KotlinTypeMarker
+
+    // PostponedArgumentInputTypesResolver
+    fun extractFunctionalTypeFromSupertypes(type: KotlinTypeMarker): KotlinTypeMarker
+    fun KotlinTypeMarker.extractArgumentsForFunctionalTypeOrSubtype(): List<KotlinTypeMarker>
+    fun KotlinTypeMarker.isFunctionOrKFunctionTypeWithAnySuspendability(): Boolean
+    fun KotlinTypeMarker.isSuspendFunctionTypeOrSubtype(): Boolean
+    fun <T> createArgumentConstraintPosition(argument: T): ArgumentConstraintPosition<T>
+    fun <T> createFixVariableConstraintPosition(variable: TypeVariableMarker, atom: T): FixVariableConstraintPosition<T>
+    fun extractParameterTypesFromDeclaration(declaration: PostponedAtomWithRevisableExpectedType): List<KotlinTypeMarker?>?
+    fun KotlinTypeMarker.isExtensionFunctionType(): Boolean
+    fun getFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker
+    fun getKFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker
+    fun isAnonymousFunction(argument: PostponedAtomWithRevisableExpectedType): Boolean
+    fun PostponedAtomWithRevisableExpectedType.isFunctionExpressionWithReceiver(): Boolean
+    fun createTypeVariableForLambdaReturnType(): TypeVariableMarker
+    fun createTypeVariableForLambdaParameterType(argument: PostponedAtomWithRevisableExpectedType, index: Int): TypeVariableMarker
+    fun createTypeVariableForCallableReferenceParameterType(
+        argument: PostponedAtomWithRevisableExpectedType,
+        index: Int
+    ): TypeVariableMarker
+    fun createTypeVariableForCallableReferenceReturnType(): TypeVariableMarker
 }

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

@@ -10,5 +10,20 @@ import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 interface PostponedResolvedAtomMarker {
     val inputTypes: Collection<KotlinTypeMarker>
     val outputType: KotlinTypeMarker?
+    val expectedType: KotlinTypeMarker?
     val analyzed: Boolean
 }
+
+interface PostponedAtomWithRevisableExpectedType : PostponedResolvedAtomMarker {
+    val revisedExpectedType: KotlinTypeMarker?
+
+    fun reviseExpectedType(expectedType: KotlinTypeMarker)
+}
+
+interface PostponedCallableReferenceMarker : PostponedAtomWithRevisableExpectedType
+
+interface LambdaWithTypeVariableAsExpectedTypeMarker : PostponedAtomWithRevisableExpectedType {
+    val parameterTypesFromDeclaration: List<KotlinTypeMarker?>?
+
+    fun updateParameterTypesFromDeclaration(types: List<KotlinTypeMarker?>?)
+}

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

@@ -5,17 +5,27 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference.components
 
+import org.jetbrains.kotlin.builtins.*
+import org.jetbrains.kotlin.builtins.getPureArgumentsForFunctionalTypeOrSubtype
 import org.jetbrains.kotlin.resolve.calls.components.CreateFreshVariablesSubstitutor.shouldBeFlexible
-import org.jetbrains.kotlin.resolve.calls.inference.model.NewTypeVariable
+import org.jetbrains.kotlin.resolve.calls.inference.components.PostponedArgumentInputTypesResolver.Companion.TYPE_VARIABLE_NAME_FOR_CR_RETURN_TYPE
-import org.jetbrains.kotlin.resolve.calls.inference.model.TypeVariableFromCallableDescriptor
+import org.jetbrains.kotlin.resolve.calls.inference.components.PostponedArgumentInputTypesResolver.Companion.TYPE_VARIABLE_NAME_FOR_LAMBDA_RETURN_TYPE
+import org.jetbrains.kotlin.resolve.calls.inference.components.PostponedArgumentInputTypesResolver.Companion.TYPE_VARIABLE_NAME_PREFIX_FOR_CR_PARAMETER_TYPE
+import org.jetbrains.kotlin.resolve.calls.inference.components.PostponedArgumentInputTypesResolver.Companion.TYPE_VARIABLE_NAME_PREFIX_FOR_LAMBDA_PARAMETER_TYPE
+import org.jetbrains.kotlin.resolve.calls.inference.model.*
+import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.checker.KotlinTypeRefiner
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 import org.jetbrains.kotlin.types.refinement.TypeRefinement
 import org.jetbrains.kotlin.types.typeUtil.unCapture as unCaptureKotlinType
 
-class ClassicConstraintSystemUtilContext(val kotlinTypeRefiner: KotlinTypeRefiner) : ConstraintSystemUtilContext {
+class ClassicConstraintSystemUtilContext(
+    val kotlinTypeRefiner: KotlinTypeRefiner,
+    val builtIns: KotlinBuiltIns,
+) : ConstraintSystemUtilContext {
     override fun TypeVariableMarker.shouldBeFlexible(): Boolean {
         return this is TypeVariableFromCallableDescriptor && this.originalTypeParameter.shouldBeFlexible()
     }
@@ -40,4 +50,110 @@ class ClassicConstraintSystemUtilContext(val kotlinTypeRefiner: KotlinTypeRefine
         require(this is KotlinType)
         return kotlinTypeRefiner.refineType(this)
     }
+
+    override fun extractFunctionalTypeFromSupertypes(type: KotlinTypeMarker): KotlinTypeMarker {
+        require(type is KotlinType)
+        return type.extractFunctionalTypeFromSupertypes()
+    }
+
+    override fun KotlinTypeMarker.extractArgumentsForFunctionalTypeOrSubtype(): List<KotlinTypeMarker> {
+        require(this is KotlinType)
+        return this.getPureArgumentsForFunctionalTypeOrSubtype()
+    }
+
+    override fun KotlinTypeMarker.isFunctionOrKFunctionTypeWithAnySuspendability(): Boolean {
+        require(this is KotlinType)
+        return this.isFunctionOrKFunctionTypeWithAnySuspendability
+    }
+
+    override fun KotlinTypeMarker.isSuspendFunctionTypeOrSubtype(): Boolean {
+        require(this is KotlinType)
+        return this.isSuspendFunctionTypeOrSubtype
+    }
+
+    override fun <T> createArgumentConstraintPosition(argument: T): ArgumentConstraintPosition<T> {
+        require(argument is ResolvedAtom)
+        @Suppress("UNCHECKED_CAST")
+        return ArgumentConstraintPositionImpl(argument.atom as KotlinCallArgument) as ArgumentConstraintPosition<T>
+    }
+
+    override fun <T> createFixVariableConstraintPosition(variable: TypeVariableMarker, atom: T): FixVariableConstraintPosition<T> {
+        require(atom is ResolvedAtom)
+        @Suppress("UNCHECKED_CAST")
+        return FixVariableConstraintPositionImpl(variable, atom) as FixVariableConstraintPosition<T>
+    }
+
+    override fun extractParameterTypesFromDeclaration(declaration: PostponedAtomWithRevisableExpectedType): List<KotlinTypeMarker?>? {
+        require(declaration is ResolvedAtom)
+        return when (val atom = declaration.atom) {
+            is FunctionExpression -> {
+                val receiverType = atom.receiverType
+                if (receiverType != null) listOf(receiverType) + atom.parametersTypes else atom.parametersTypes.toList()
+            }
+            is LambdaKotlinCallArgument -> atom.parametersTypes?.toList()
+            else -> null
+        }
+    }
+
+    override fun KotlinTypeMarker.isExtensionFunctionType(): Boolean {
+        require(this is KotlinType)
+        return this.isExtensionFunctionType
+    }
+
+    override fun getFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker {
+        return getFunctionDescriptor(builtIns, parametersNumber, isSuspend).typeConstructor
+    }
+
+    override fun getKFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker {
+        return getKFunctionDescriptor(builtIns, parametersNumber, isSuspend).typeConstructor
+    }
+
+    override fun isAnonymousFunction(argument: PostponedAtomWithRevisableExpectedType): Boolean {
+        require(argument is ResolvedAtom)
+        return argument.atom is FunctionExpression
+    }
+
+    override fun PostponedAtomWithRevisableExpectedType.isFunctionExpressionWithReceiver(): Boolean {
+        require(this is ResolvedAtom)
+        val atom = this.atom
+        return atom is FunctionExpression && atom.receiverType != null
+    }
+
+    override fun createTypeVariableForLambdaReturnType(): TypeVariableMarker {
+        return TypeVariableForLambdaReturnType(
+            builtIns,
+            TYPE_VARIABLE_NAME_FOR_LAMBDA_RETURN_TYPE
+        )
+    }
+
+    override fun createTypeVariableForLambdaParameterType(
+        argument: PostponedAtomWithRevisableExpectedType,
+        index: Int
+    ): TypeVariableMarker {
+        require(argument is ResolvedAtom)
+        val atom = argument.atom as PostponableKotlinCallArgument
+        return TypeVariableForLambdaParameterType(
+            atom,
+            index,
+            builtIns,
+            TYPE_VARIABLE_NAME_PREFIX_FOR_LAMBDA_PARAMETER_TYPE + (index + 1)
+        )
+    }
+
+    override fun createTypeVariableForCallableReferenceParameterType(
+        argument: PostponedAtomWithRevisableExpectedType,
+        index: Int
+    ): TypeVariableMarker {
+        return TypeVariableForCallableReferenceParameterType(
+            builtIns,
+            TYPE_VARIABLE_NAME_PREFIX_FOR_CR_PARAMETER_TYPE + (index + 1)
+        )
+    }
+
+    override fun createTypeVariableForCallableReferenceReturnType(): TypeVariableMarker {
+        return TypeVariableForCallableReferenceReturnType(
+            builtIns,
+            TYPE_VARIABLE_NAME_FOR_CR_RETURN_TYPE
+        )
+    }
 }

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

@@ -14,13 +14,15 @@ import org.jetbrains.kotlin.types.model.TypeVariableMarker
 import org.jetbrains.kotlin.types.model.safeSubstitute
 import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.utils.addIfNotNull
+import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class KotlinConstraintSystemCompleter(
     private val resultTypeResolver: ResultTypeResolver,
     val variableFixationFinder: VariableFixationFinder,
+    private val ctx: ConstraintSystemUtilContext,
 ) {
-    private val postponedArgumentInputTypesResolver = PostponedArgumentInputTypesResolver(resultTypeResolver, variableFixationFinder)
+    private val postponedArgumentInputTypesResolver = PostponedArgumentInputTypesResolver(resultTypeResolver, variableFixationFinder, ctx)
 
     fun runCompletion(
         c: ConstraintSystemCompletionContext,
@@ -103,9 +105,10 @@ class KotlinConstraintSystemCompleter(
                         argument,
                         postponedArguments,
                         topLevelType,
-                        topLevelAtoms,
                         dependencyProvider,
-                    )
+                    ) {
+                        findResolvedAtomBy(it, topLevelAtoms) ?: topLevelAtoms.firstOrNull()
+                    }
 
                     if (wasFixedSomeVariable)
                         continue@completion
@@ -113,7 +116,7 @@ class KotlinConstraintSystemCompleter(
 
                 // Stage 4: create atoms with revised expected types if needed
                 for (argument in postponedArgumentsWithRevisableType) {
-                    val wasTransformedSomeArgument = postponedArgumentInputTypesResolver.transformToAtomWithNewFunctionalExpectedType(
+                    val wasTransformedSomeArgument = transformToAtomWithNewFunctionalExpectedType(
                         asConstraintSystemCompletionContext(), argument, diagnosticsHolder
                     )
 
@@ -143,6 +146,27 @@ class KotlinConstraintSystemCompleter(
         }
     }
 
+    private fun transformToAtomWithNewFunctionalExpectedType(
+        c: ConstraintSystemCompletionContext,
+        argument: PostponedAtomWithRevisableExpectedType,
+        diagnosticsHolder: KotlinDiagnosticsHolder
+    ): Boolean = with(ctx) {
+        val revisedExpectedType: UnwrappedType =
+            argument.revisedExpectedType?.takeIf { it.isFunctionOrKFunctionTypeWithAnySuspendability() }?.cast() ?: return false
+
+        when (argument) {
+            is PostponedCallableReferenceAtom ->
+                CallableReferenceWithRevisedExpectedTypeAtom(argument.atom, revisedExpectedType).also {
+                    argument.setAnalyzedResults(null, listOf(it))
+                }
+            is LambdaWithTypeVariableAsExpectedTypeAtom ->
+                argument.transformToResolvedLambda(c.getBuilder(), diagnosticsHolder, revisedExpectedType)
+            else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
+        }
+
+        return true
+    }
+
     private fun ConstraintSystemCompletionContext.analyzeArgumentWithFixedParameterTypes(
         postponedArguments: List<PostponedResolvedAtom>,
         analyze: (PostponedResolvedAtom) -> Unit

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

@@ -5,34 +5,31 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference.components
 
-import org.jetbrains.kotlin.builtins.*
-import org.jetbrains.kotlin.descriptors.annotations.Annotations
-import org.jetbrains.kotlin.resolve.calls.components.transformToResolvedLambda
 import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.model.*
-import org.jetbrains.kotlin.types.*
+import org.jetbrains.kotlin.types.model.*
-import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
-import org.jetbrains.kotlin.types.typeUtil.builtIns
 import org.jetbrains.kotlin.utils.SmartSet
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 private typealias Context = ConstraintSystemCompletionContext
+private typealias ResolvedAtomProvider = (TypeVariableMarker) -> Any?
 
 class PostponedArgumentInputTypesResolver(
     private val resultTypeResolver: ResultTypeResolver,
-    private val variableFixationFinder: VariableFixationFinder
+    private val variableFixationFinder: VariableFixationFinder,
+    private val resolutionTypeSystemContext: ConstraintSystemUtilContext,
 ) {
     private class ParameterTypesInfo(
-        val parametersFromDeclaration: List<UnwrappedType?>?,
+        val parametersFromDeclaration: List<KotlinTypeMarker?>?,
-        val parametersFromDeclarationOfRelatedLambdas: Set<List<UnwrappedType?>>?,
+        val parametersFromDeclarationOfRelatedLambdas: Set<List<KotlinTypeMarker?>>?,
         val parametersFromConstraints: Set<List<TypeWithKind>>?,
-        val annotations: Annotations,
+        val isExtensionFunction: Boolean,
         val isSuspend: Boolean,
         val isNullable: Boolean
     )
 
     data class TypeWithKind(
-        val type: KotlinType,
+        val type: KotlinTypeMarker,
         val direction: ConstraintKind = ConstraintKind.UPPER
     )
 
@@ -41,38 +38,27 @@ class PostponedArgumentInputTypesResolver(
         variableDependencyProvider: TypeVariableDependencyInformationProvider
     ): List<TypeWithKind>? {
         fun List<Constraint>.extractFunctionalTypes() = mapNotNull { constraint ->
-            val type = constraint.type as? KotlinType ?: return@mapNotNull null
+            TypeWithKind(resolutionTypeSystemContext.extractFunctionalTypeFromSupertypes(constraint.type), constraint.kind)
-            TypeWithKind(type.extractFunctionalTypeFromSupertypes(), constraint.kind)
         }
 
-        val typeVariableTypeConstructor = variable.typeVariable.freshTypeConstructor() as? TypeVariableTypeConstructor ?: return null
+        val typeVariableTypeConstructor = variable.typeVariable.freshTypeConstructor()
         val dependentVariables =
             variableDependencyProvider.getShallowlyDependentVariables(typeVariableTypeConstructor).orEmpty() + typeVariableTypeConstructor
 
         return dependentVariables.flatMap { type ->
             val constraints = notFixedTypeVariables[type]?.constraints ?: return@flatMap emptyList()
-            val constraintsWithFunctionalType = constraints.filter { (it.type as? KotlinType)?.isBuiltinFunctionalTypeOrSubtype == true }
+            val constraintsWithFunctionalType = constraints.filter { it.type.isBuiltinFunctionalTypeOrSubtype() }
             constraintsWithFunctionalType.extractFunctionalTypes()
         }
     }
 
-    private fun extractParameterTypesFromDeclaration(atom: ResolutionAtom) =
-        when (atom) {
-            is FunctionExpression -> {
-                val receiverType = atom.receiverType
-                if (receiverType != null) listOf(receiverType) + atom.parametersTypes else atom.parametersTypes.toList()
-            }
-            is LambdaKotlinCallArgument -> atom.parametersTypes?.toList()
-            else -> null
-        }
-
     private fun Context.extractParameterTypesInfo(
         argument: PostponedAtomWithRevisableExpectedType,
         postponedArguments: List<PostponedAtomWithRevisableExpectedType>,
         variableDependencyProvider: TypeVariableDependencyInformationProvider
-    ): ParameterTypesInfo? {
+    ): ParameterTypesInfo? = with(resolutionTypeSystemContext) {
         val expectedType = argument.expectedType ?: return null
-        val variableWithConstraints = notFixedTypeVariables[expectedType.constructor] ?: return null
+        val variableWithConstraints = notFixedTypeVariables[expectedType.typeConstructor()] ?: return null
         val functionalTypesFromConstraints = findFunctionalTypesInConstraints(variableWithConstraints, variableDependencyProvider)
 
         // Don't create functional expected type for further error reporting about a different number of arguments
@@ -80,10 +66,10 @@ class PostponedArgumentInputTypesResolver(
             return null
 
         val parameterTypesFromDeclaration =
-            if (argument is LambdaWithTypeVariableAsExpectedTypeAtom) argument.parameterTypesFromDeclaration else null
+            if (argument is LambdaWithTypeVariableAsExpectedTypeMarker) argument.parameterTypesFromDeclaration else null
 
         val parameterTypesFromConstraints = functionalTypesFromConstraints?.mapTo(SmartSet.create()) { typeWithKind ->
-            typeWithKind.type.getPureArgumentsForFunctionalTypeOrSubtype().map {
+            typeWithKind.type.extractArgumentsForFunctionalTypeOrSubtype().map {
                 // We should use opposite kind as lambda's parameters are contravariant
                 TypeWithKind(it, typeWithKind.direction.opposite())
             }
@@ -93,21 +79,15 @@ class PostponedArgumentInputTypesResolver(
         val (parameterTypesFromDeclarationOfRelatedLambdas, isThereExtensionFunctionAmongRelatedLambdas) =
             getDeclaredParametersFromRelatedLambdas(argument, postponedArguments, variableDependencyProvider)
 
-        val annotationsFromConstraints = functionalTypesFromConstraints?.run {
+        val extensionFunctionTypePresentInConstraints = functionalTypesFromConstraints?.any { it.type.isExtensionFunctionType() } == true
-            Annotations.create(flatMapTo(SmartSet.create()) { it.type.annotations }.toList())
-        } ?: Annotations.EMPTY
-
-        val annotations = if (isThereExtensionFunctionAmongRelatedLambdas) {
-            annotationsFromConstraints.withExtensionFunctionAnnotation(expectedType.builtIns)
-        } else annotationsFromConstraints
 
         var isSuspend = false
         var isNullable = false
         if (!functionalTypesFromConstraints.isNullOrEmpty()) {
             isNullable = true
             for (funType in functionalTypesFromConstraints) {
-                if (!isSuspend && funType.type.isSuspendFunctionTypeOrSubtype) isSuspend = true
+                if (!isSuspend && funType.type.isSuspendFunctionTypeOrSubtype()) isSuspend = true
-                if (isNullable && !funType.type.isMarkedNullable) isNullable = false
+                if (isNullable && !funType.type.isMarkedNullable()) isNullable = false
                 if (isSuspend && !isNullable) break
             }
         }
@@ -115,7 +95,7 @@ class PostponedArgumentInputTypesResolver(
             parameterTypesFromDeclaration,
             parameterTypesFromDeclarationOfRelatedLambdas,
             parameterTypesFromConstraints,
-            annotations,
+            isExtensionFunction = isThereExtensionFunctionAmongRelatedLambdas || extensionFunctionTypePresentInConstraints,
             isSuspend = isSuspend,
             isNullable = isNullable
         )
@@ -125,20 +105,20 @@ class PostponedArgumentInputTypesResolver(
         argument: PostponedAtomWithRevisableExpectedType,
         postponedArguments: List<PostponedAtomWithRevisableExpectedType>,
         dependencyProvider: TypeVariableDependencyInformationProvider
-    ): Pair<Set<List<UnwrappedType?>>?, Boolean> {
+    ): Pair<Set<List<KotlinTypeMarker?>>?, Boolean> = with(resolutionTypeSystemContext) {
         val parameterTypesFromDeclarationOfRelatedLambdas = postponedArguments
             .mapNotNull { anotherArgument ->
                 when {
-                    anotherArgument !is LambdaWithTypeVariableAsExpectedTypeAtom -> null
+                    anotherArgument !is LambdaWithTypeVariableAsExpectedTypeMarker -> null
                     anotherArgument.parameterTypesFromDeclaration == null || anotherArgument == argument -> null
                     else -> {
                         val argumentExpectedTypeConstructor = argument.expectedType?.typeConstructor() ?: return@mapNotNull null
-                        val anotherArgumentExpectedTypeConstructor = anotherArgument.expectedType.typeConstructor()
+                        val anotherArgumentExpectedTypeConstructor =
+                            anotherArgument.expectedType?.typeConstructor() ?: return@mapNotNull null
                         val areTypeVariablesRelated = dependencyProvider.areVariablesDependentShallowly(
                             argumentExpectedTypeConstructor, anotherArgumentExpectedTypeConstructor
                         )
-                        val anotherAtom = anotherArgument.atom
+                        val isAnonymousExtensionFunction = anotherArgument.isFunctionExpressionWithReceiver()
-                        val isAnonymousExtensionFunction = anotherAtom is FunctionExpression && anotherAtom.receiverType != null
                         val parameterTypesFromDeclarationOfRelatedLambda = anotherArgument.parameterTypesFromDeclaration
 
                         if (areTypeVariablesRelated && parameterTypesFromDeclarationOfRelatedLambda != null) {
@@ -151,57 +131,44 @@ class PostponedArgumentInputTypesResolver(
         return parameterTypesFromDeclarationOfRelatedLambdas.run { mapTo(SmartSet.create()) { it.first } to any { it.second } }
     }
 
-    private fun Context.createTypeVariableForReturnType(argument: PostponedAtomWithRevisableExpectedType): NewTypeVariable {
+    private fun Context.createTypeVariableForReturnType(argument: PostponedAtomWithRevisableExpectedType): TypeVariableMarker =
-        val expectedType = argument.expectedType
+        with(resolutionTypeSystemContext) {
-            ?: throw IllegalStateException("Postponed argument's expected type must not be null")
+            val expectedType = argument.expectedType
+                ?: throw IllegalStateException("Postponed argument's expected type must not be null")
 
-        val variable = getBuilder().currentStorage().allTypeVariables[expectedType.constructor]
+            val variable = getBuilder().currentStorage().allTypeVariables[expectedType.typeConstructor()]
-        if (variable != null) {
+            if (variable != null) {
-            val revisedVariableForReturnType = getBuilder().getRevisedVariableForReturnType(variable)
+                val revisedVariableForReturnType = getBuilder().getRevisedVariableForReturnType(variable)
-            if (revisedVariableForReturnType != null) return revisedVariableForReturnType as NewTypeVariable
+                if (revisedVariableForReturnType != null) return revisedVariableForReturnType
+            }
+
+            return when (argument) {
+                is LambdaWithTypeVariableAsExpectedTypeMarker -> createTypeVariableForLambdaReturnType()
+                is PostponedCallableReferenceMarker -> createTypeVariableForCallableReferenceReturnType()
+                else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
+            }.also {
+                if (variable != null) getBuilder().putRevisedVariableForReturnType(variable, it)
+
+                getBuilder().registerVariable(it)
+            }
         }
 
-        return when (argument) {
-            is LambdaWithTypeVariableAsExpectedTypeAtom -> TypeVariableForLambdaReturnType(
-                expectedType.builtIns,
-                TYPE_VARIABLE_NAME_FOR_LAMBDA_RETURN_TYPE
-            )
-            is PostponedCallableReferenceAtom -> TypeVariableForCallableReferenceReturnType(
-                expectedType.builtIns,
-                TYPE_VARIABLE_NAME_FOR_CR_RETURN_TYPE
-            )
-            else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
-        }.also {
-            if (variable != null) getBuilder().putRevisedVariableForReturnType(variable, it)
-
-            getBuilder().registerVariable(it)
-        }
-    }
-
     private fun Context.createTypeVariableForParameterType(
         argument: PostponedAtomWithRevisableExpectedType,
         index: Int
-    ): NewTypeVariable {
+    ): TypeVariableMarker = with(resolutionTypeSystemContext) {
         val expectedType = argument.expectedType
             ?: throw IllegalStateException("Postponed argument's expected type must not be null")
 
-        val variable = getBuilder().currentStorage().allTypeVariables[expectedType.constructor]
+        val variable = getBuilder().currentStorage().allTypeVariables[expectedType.typeConstructor()]
         if (variable != null) {
             val revisedVariableForParameter = getBuilder().getRevisedVariableForParameter(variable, index)
-            if (revisedVariableForParameter != null) return revisedVariableForParameter as NewTypeVariable
+            if (revisedVariableForParameter != null) return revisedVariableForParameter
         }
 
         return when (argument) {
-            is LambdaWithTypeVariableAsExpectedTypeAtom -> TypeVariableForLambdaParameterType(
+            is LambdaWithTypeVariableAsExpectedTypeMarker -> createTypeVariableForLambdaParameterType(argument, index)
-                argument.atom,
+            is PostponedCallableReferenceMarker -> createTypeVariableForCallableReferenceParameterType(argument, index)
-                index,
-                expectedType.builtIns,
-                TYPE_VARIABLE_NAME_PREFIX_FOR_LAMBDA_PARAMETER_TYPE + (index + 1)
-            )
-            is PostponedCallableReferenceAtom -> TypeVariableForCallableReferenceParameterType(
-                expectedType.builtIns,
-                TYPE_VARIABLE_NAME_PREFIX_FOR_CR_PARAMETER_TYPE + (index + 1)
-            )
             else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
         }.also {
             if (variable != null) getBuilder().putRevisedVariableForParameter(variable, index, it)
@@ -213,14 +180,13 @@ class PostponedArgumentInputTypesResolver(
     private fun Context.createTypeVariablesForParameters(
         argument: PostponedAtomWithRevisableExpectedType,
         parameterTypes: List<List<TypeWithKind?>>
-    ): List<TypeProjection> {
+    ): List<TypeArgumentMarker> = with(resolutionTypeSystemContext) {
-        val atom = argument.atom
         val csBuilder = getBuilder()
         val allGroupedParameterTypes = parameterTypes.first().indices.map { i -> parameterTypes.map { it.getOrNull(i) } }
 
         return allGroupedParameterTypes.mapIndexed { index, types ->
             val parameterTypeVariable = createTypeVariableForParameterType(argument, index)
-            val typeVariableConstructor = parameterTypeVariable.freshTypeConstructor
+            val typeVariableConstructor = parameterTypeVariable.freshTypeConstructor()
 
             for (typeWithKind in types) {
                 if (typeVariableConstructor in fixedTypeVariables) break
@@ -228,24 +194,20 @@ class PostponedArgumentInputTypesResolver(
 
                 when (typeWithKind.direction) {
                     ConstraintKind.EQUALITY -> csBuilder.addEqualityConstraint(
-                        parameterTypeVariable.defaultType, typeWithKind.type, ArgumentConstraintPositionImpl(atom)
+                        parameterTypeVariable.defaultType(), typeWithKind.type, createArgumentConstraintPosition(argument)
                     )
                     ConstraintKind.UPPER -> csBuilder.addSubtypeConstraint(
-                        parameterTypeVariable.defaultType, typeWithKind.type, ArgumentConstraintPositionImpl(atom)
+                        parameterTypeVariable.defaultType(), typeWithKind.type, createArgumentConstraintPosition(argument)
                     )
                     ConstraintKind.LOWER -> csBuilder.addSubtypeConstraint(
-                        typeWithKind.type, parameterTypeVariable.defaultType, ArgumentConstraintPositionImpl(atom)
+                        typeWithKind.type, parameterTypeVariable.defaultType(), createArgumentConstraintPosition(argument)
                     )
                 }
             }
 
-            val resultType = if (typeVariableConstructor in fixedTypeVariables) {
+            val resultType = fixedTypeVariables[typeVariableConstructor] ?: parameterTypeVariable.defaultType()
-                fixedTypeVariables[typeVariableConstructor] as KotlinType
-            } else {
-                parameterTypeVariable.defaultType
-            }
 
-            resultType.asTypeProjection()
+            resultType.asTypeArgument()
         }
     }
 
@@ -254,27 +216,27 @@ class PostponedArgumentInputTypesResolver(
         parametersNumber: Int,
         isSuspend: Boolean,
         resultTypeResolver: ResultTypeResolver
-    ): TypeConstructor {
+    ): TypeConstructorMarker = with(resolutionTypeSystemContext) {
         val expectedType = argument.expectedType
             ?: throw IllegalStateException("Postponed argument's expected type must not be null")
 
-        val expectedTypeConstructor = expectedType.constructor
+        val expectedTypeConstructor = expectedType.typeConstructor()
 
         return when (argument) {
-            is LambdaWithTypeVariableAsExpectedTypeAtom ->
+            is LambdaWithTypeVariableAsExpectedTypeMarker ->
-                getFunctionDescriptor(expectedTypeConstructor.builtIns, parametersNumber, isSuspend).typeConstructor
+                getFunctionTypeConstructor(parametersNumber, isSuspend)
-            is PostponedCallableReferenceAtom -> {
+            is PostponedCallableReferenceMarker -> {
                 val computedResultType = resultTypeResolver.findResultType(
-                    this,
+                    this@computeResultingFunctionalConstructor,
                     notFixedTypeVariables.getValue(expectedTypeConstructor),
                     TypeVariableDirectionCalculator.ResolveDirection.TO_SUPERTYPE
                 )
 
                 // Avoid KFunction<...>/Function<...> types
                 if (computedResultType.isBuiltinFunctionalTypeOrSubtype() && computedResultType.argumentsCount() > 1) {
-                    computedResultType.typeConstructor() as TypeConstructor
+                    computedResultType.typeConstructor()
                 } else {
-                    getKFunctionDescriptor(expectedTypeConstructor.builtIns, parametersNumber, isSuspend).typeConstructor
+                    getKFunctionTypeConstructor(parametersNumber, isSuspend)
                 }
             }
             else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
@@ -284,17 +246,16 @@ class PostponedArgumentInputTypesResolver(
     private fun Context.buildNewFunctionalExpectedType(
         argument: PostponedAtomWithRevisableExpectedType,
         parameterTypesInfo: ParameterTypesInfo
-    ): UnwrappedType? {
+    ): KotlinTypeMarker? = with(resolutionTypeSystemContext) {
         val expectedType = argument.expectedType
 
-        if (expectedType == null || expectedType.constructor !in notFixedTypeVariables)
+        if (expectedType == null || expectedType.typeConstructor() !in notFixedTypeVariables)
             return null
 
-        val atom = argument.atom
         val parametersFromConstraints = parameterTypesInfo.parametersFromConstraints
         val parametersFromDeclaration = getDeclaredParametersConsideringExtensionFunctionsPresence(parameterTypesInfo)
         val areAllParameterTypesSpecified = !parametersFromDeclaration.isNullOrEmpty() && parametersFromDeclaration.all { it != null }
-        val isExtensionFunction = parameterTypesInfo.annotations.hasExtensionFunctionAnnotation()
+        val isExtensionFunction = parameterTypesInfo.isExtensionFunction
         val parametersFromDeclarations = parameterTypesInfo.parametersFromDeclarationOfRelatedLambdas.orEmpty() + parametersFromDeclaration
 
         /*
@@ -306,7 +267,9 @@ class PostponedArgumentInputTypesResolver(
             return null
 
         val allParameterTypes =
-            (parametersFromConstraints.orEmpty() + parametersFromDeclarations.map { parameters -> parameters?.map { it.wrapToTypeWithKind() } }).filterNotNull()
+            (parametersFromConstraints.orEmpty() + parametersFromDeclarations.map { parameters ->
+                parameters?.map { it.wrapToTypeWithKind() }
+            }).filterNotNull()
 
         if (allParameterTypes.isEmpty())
             return null
@@ -320,7 +283,7 @@ class PostponedArgumentInputTypesResolver(
             resultTypeResolver
         )
 
-        val isExtensionFunctionType = parameterTypesInfo.annotations.hasExtensionFunctionAnnotation()
+        val isExtensionFunctionType = parameterTypesInfo.isExtensionFunction
         val areParametersNumberInDeclarationAndConstraintsEqual =
             !parametersFromDeclaration.isNullOrEmpty() && !parametersFromConstraints.isNullOrEmpty()
                     && parametersFromDeclaration.size == parametersFromConstraints.first().size
@@ -339,27 +302,22 @@ class PostponedArgumentInputTypesResolver(
          *
          * Example: `val x = id(fun String.() = this)`
          */
-        val shouldAddExtensionFunctionAnnotation = atom is FunctionExpression && atom.receiverType != null
 
-        val annotations = when {
+        val newExpectedType = createSimpleType(
-            shouldDiscriminateExtensionFunctionAnnotation ->
-                parameterTypesInfo.annotations.withoutExtensionFunctionAnnotation()
-            shouldAddExtensionFunctionAnnotation ->
-                parameterTypesInfo.annotations.withExtensionFunctionAnnotation(expectedType.builtIns)
-            else -> parameterTypesInfo.annotations
-        }
-
-        val newExpectedType = KotlinTypeFactory.simpleType(
-            annotations,
             functionalConstructor,
-            variablesForParameterTypes + variableForReturnType.defaultType.asTypeProjection(),
+            variablesForParameterTypes + variableForReturnType.defaultType().asTypeArgument(),
-            parameterTypesInfo.isNullable
+            parameterTypesInfo.isNullable,
+            isExtensionFunction = when {
+                shouldDiscriminateExtensionFunctionAnnotation -> false
+                argument.isFunctionExpressionWithReceiver() -> true
+                else -> parameterTypesInfo.isExtensionFunction
+            }
         )
 
         getBuilder().addSubtypeConstraint(
             newExpectedType,
             expectedType,
-            ArgumentConstraintPositionImpl(argument.atom)
+            createArgumentConstraintPosition(argument)
         )
 
         return newExpectedType
@@ -370,12 +328,12 @@ class PostponedArgumentInputTypesResolver(
         postponedArguments: List<PostponedAtomWithRevisableExpectedType>,
         completionMode: ConstraintSystemCompletionMode,
         dependencyProvider: TypeVariableDependencyInformationProvider
-    ): Boolean {
+    ): Boolean = with(resolutionTypeSystemContext) {
         // We can collect parameter types from declaration in any mode, they can't change during completion.
         for (argument in postponedArguments) {
-            if (argument !is LambdaWithTypeVariableAsExpectedTypeAtom) continue
+            if (argument !is LambdaWithTypeVariableAsExpectedTypeMarker) continue
             if (argument.parameterTypesFromDeclaration != null) continue
-            argument.parameterTypesFromDeclaration = extractParameterTypesFromDeclaration(argument.atom)
+            argument.updateParameterTypesFromDeclaration(extractParameterTypesFromDeclaration(argument))
         }
 
         return postponedArguments.any { argument ->
@@ -394,54 +352,34 @@ class PostponedArgumentInputTypesResolver(
             val newExpectedType =
                 c.buildNewFunctionalExpectedType(argument, parameterTypesInfo) ?: return@any false
 
-            argument.revisedExpectedType = newExpectedType
+            argument.reviseExpectedType(newExpectedType)
 
             true
         }
     }
 
-    fun transformToAtomWithNewFunctionalExpectedType(
+    private fun Context.getAllDeeplyRelatedTypeVariables(
-        c: Context,
+        type: KotlinTypeMarker,
-        argument: PostponedAtomWithRevisableExpectedType,
-        diagnosticsHolder: KotlinDiagnosticsHolder
-    ): Boolean {
-        val revisedExpectedType = argument.revisedExpectedType?.takeIf { it.isFunctionOrKFunctionTypeWithAnySuspendability } ?: return false
-
-        when (argument) {
-            is PostponedCallableReferenceAtom ->
-                CallableReferenceWithRevisedExpectedTypeAtom(argument.atom, revisedExpectedType).also {
-                    argument.setAnalyzedResults(null, listOf(it))
-                }
-            is LambdaWithTypeVariableAsExpectedTypeAtom ->
-                argument.transformToResolvedLambda(c.getBuilder(), diagnosticsHolder, revisedExpectedType)
-            else -> throw IllegalStateException("Unsupported postponed argument type of $argument")
-        }
-
-        return true
-    }
-
-    private fun getAllDeeplyRelatedTypeVariables(
-        type: KotlinType,
         variableDependencyProvider: TypeVariableDependencyInformationProvider
-    ): List<TypeVariableTypeConstructor> {
+    ): List<TypeVariableTypeConstructorMarker> {
-        val typeConstructor = type.constructor
+        val typeConstructor = type.typeConstructor()
 
         return when {
-            typeConstructor is TypeVariableTypeConstructor -> {
+            typeConstructor is TypeVariableTypeConstructorMarker -> {
                 val relatedVariables = variableDependencyProvider.getDeeplyDependentVariables(typeConstructor).orEmpty()
-                listOf(typeConstructor) + relatedVariables.filterIsInstance<TypeVariableTypeConstructor>()
+                listOf(typeConstructor) + relatedVariables.filterIsInstance<TypeVariableTypeConstructorMarker>()
             }
-            type.arguments.isNotEmpty() -> {
+            type.argumentsCount() > 0 -> {
-                type.arguments.flatMap {
+                type.getArguments().flatMap {
-                    if (it.isStarProjection) emptyList() else getAllDeeplyRelatedTypeVariables(it.type, variableDependencyProvider)
+                    if (it.isStarProjection()) emptyList() else getAllDeeplyRelatedTypeVariables(it.getType(), variableDependencyProvider)
                 }
             }
             else -> emptyList()
         }
     }
 
-    private fun getDeclaredParametersConsideringExtensionFunctionsPresence(parameterTypesInfo: ParameterTypesInfo): List<UnwrappedType?>? =
+    private fun getDeclaredParametersConsideringExtensionFunctionsPresence(parameterTypesInfo: ParameterTypesInfo): List<KotlinTypeMarker?>? =
-        with (parameterTypesInfo) {
+        with(parameterTypesInfo) {
 
             if (parametersFromConstraints.isNullOrEmpty() || parametersFromDeclaration.isNullOrEmpty())
                 parametersFromDeclaration
@@ -449,7 +387,7 @@ class PostponedArgumentInputTypesResolver(
                 val oneLessParameterInDeclarationThanInConstraints =
                     parametersFromConstraints.first().size == parametersFromDeclaration.size + 1
 
-                if (oneLessParameterInDeclarationThanInConstraints && annotations.hasExtensionFunctionAnnotation()) {
+                if (oneLessParameterInDeclarationThanInConstraints && isExtensionFunction) {
                     listOf(null) + parametersFromDeclaration
                 } else {
                     parametersFromDeclaration
@@ -459,23 +397,22 @@ class PostponedArgumentInputTypesResolver(
 
     fun fixNextReadyVariableForParameterTypeIfNeeded(
         c: Context,
-        argument: PostponedResolvedAtom,
+        argument: PostponedResolvedAtomMarker,
-        postponedArguments: List<PostponedResolvedAtom>,
+        postponedArguments: List<PostponedResolvedAtomMarker>,
-        topLevelType: UnwrappedType,
+        topLevelType: KotlinTypeMarker,
-        topLevelAtoms: List<ResolvedAtom>,
         dependencyProvider: TypeVariableDependencyInformationProvider,
-    ): Boolean {
+        resolvedAtomProvider: ResolvedAtomProvider
+    ): Boolean = with(resolutionTypeSystemContext) {
         val expectedType = argument.run { safeAs<PostponedAtomWithRevisableExpectedType>()?.revisedExpectedType ?: expectedType }
 
-        if (expectedType != null && expectedType.isFunctionOrKFunctionTypeWithAnySuspendability) {
+        if (expectedType != null && expectedType.isFunctionOrKFunctionTypeWithAnySuspendability()) {
-            val wasFixedSomeVariable =
+            val wasFixedSomeVariable = c.fixNextReadyVariableForParameterType(
-                c.fixNextReadyVariableForParameterType(
+                expectedType,
-                    expectedType,
+                postponedArguments,
-                    postponedArguments,
+                topLevelType,
-                    topLevelType,
+                dependencyProvider,
-                    topLevelAtoms,
+                resolvedAtomProvider
-                    dependencyProvider,
+            )
-                )
 
             if (wasFixedSomeVariable)
                 return true
@@ -485,16 +422,16 @@ class PostponedArgumentInputTypesResolver(
     }
 
     private fun Context.fixNextReadyVariableForParameterType(
-        type: KotlinType,
+        type: KotlinTypeMarker,
-        postponedArguments: List<PostponedResolvedAtom>,
+        postponedArguments: List<PostponedResolvedAtomMarker>,
-        topLevelType: UnwrappedType,
+        topLevelType: KotlinTypeMarker,
-        topLevelAtoms: List<ResolvedAtom>,
         dependencyProvider: TypeVariableDependencyInformationProvider,
-    ): Boolean {
+        resolvedAtomByTypeVariableProvider: ResolvedAtomProvider,
-        val relatedVariables = type.getPureArgumentsForFunctionalTypeOrSubtype()
+    ): Boolean = with(resolutionTypeSystemContext) {
+        val relatedVariables = type.extractArgumentsForFunctionalTypeOrSubtype()
             .flatMap { getAllDeeplyRelatedTypeVariables(it, dependencyProvider) }
         val variableForFixation = variableFixationFinder.findFirstVariableForFixation(
-            this, relatedVariables, postponedArguments, ConstraintSystemCompletionMode.FULL, topLevelType
+            this@fixNextReadyVariableForParameterType, relatedVariables, postponedArguments, ConstraintSystemCompletionMode.FULL, topLevelType
         )
 
         if (variableForFixation == null || !variableForFixation.hasProperConstraint)
@@ -502,24 +439,28 @@ class PostponedArgumentInputTypesResolver(
 
         val variableWithConstraints = notFixedTypeVariables.getValue(variableForFixation.variable)
         val resultType =
-            resultTypeResolver.findResultType(this, variableWithConstraints, TypeVariableDirectionCalculator.ResolveDirection.UNKNOWN)
+            resultTypeResolver.findResultType(
+                this@fixNextReadyVariableForParameterType,
+                variableWithConstraints,
+                TypeVariableDirectionCalculator.ResolveDirection.UNKNOWN
+            )
         val variable = variableWithConstraints.typeVariable
-        val resolvedAtom = KotlinConstraintSystemCompleter.findResolvedAtomBy(variable, topLevelAtoms)
-            ?: topLevelAtoms.firstOrNull()
 
-        fixVariable(variable, resultType, FixVariableConstraintPositionImpl(variable, resolvedAtom))
+        fixVariable(
+            variable,
+            resultType,
+            createFixVariableConstraintPosition(variable, resolvedAtomByTypeVariableProvider(variable))
+        )
 
         return true
     }
 
-    private fun KotlinType?.wrapToTypeWithKind() = this?.let { TypeWithKind(it) }
+    private fun KotlinTypeMarker?.wrapToTypeWithKind() = this?.let { TypeWithKind(it) }
-
-    private fun isAnonymousFunction(argument: PostponedAtomWithRevisableExpectedType) = argument.atom is FunctionExpression
 
     companion object {
-        private const val TYPE_VARIABLE_NAME_PREFIX_FOR_LAMBDA_PARAMETER_TYPE = "_RP"
+        const val TYPE_VARIABLE_NAME_PREFIX_FOR_LAMBDA_PARAMETER_TYPE = "_RP"
-        private const val TYPE_VARIABLE_NAME_FOR_LAMBDA_RETURN_TYPE = "_R"
+        const val TYPE_VARIABLE_NAME_FOR_LAMBDA_RETURN_TYPE = "_R"
-        private const val TYPE_VARIABLE_NAME_PREFIX_FOR_CR_PARAMETER_TYPE = "_QP"
+        const val TYPE_VARIABLE_NAME_PREFIX_FOR_CR_PARAMETER_TYPE = "_QP"
-        private const val TYPE_VARIABLE_NAME_FOR_CR_RETURN_TYPE = "_Q"
+        const val TYPE_VARIABLE_NAME_FOR_CR_RETURN_TYPE = "_Q"
     }
-}
+}

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/ResolutionAtoms.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeConstructor
 import org.jetbrains.kotlin.types.TypeSubstitutor
 import org.jetbrains.kotlin.types.UnwrappedType
+import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.typeUtil.unCapture
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
@@ -92,28 +93,36 @@ class ResolvedSubCallArgument(override val atom: SubKotlinCallArgument, resolveI
     }
 }
 
-interface PostponedAtomWithRevisableExpectedType {
-    var revisedExpectedType: UnwrappedType?
-    val expectedType: UnwrappedType?
-    val atom: PostponableKotlinCallArgument
-}
 
 sealed class PostponedResolvedAtom : ResolvedAtom(), PostponedResolvedAtomMarker {
     abstract override val inputTypes: Collection<UnwrappedType>
     abstract override val outputType: UnwrappedType?
-    abstract val expectedType: UnwrappedType?
+    abstract override val expectedType: UnwrappedType?
 }
 
 class LambdaWithTypeVariableAsExpectedTypeAtom(
     override val atom: LambdaKotlinCallArgument,
     override val expectedType: UnwrappedType
-) : PostponedResolvedAtom(), PostponedAtomWithRevisableExpectedType {
+) : PostponedResolvedAtom(), LambdaWithTypeVariableAsExpectedTypeMarker {
     override val inputTypes: Collection<UnwrappedType> get() = listOf(expectedType)
     override val outputType: UnwrappedType? get() = null
 
     override var revisedExpectedType: UnwrappedType? = null
+        private set
 
-    var parameterTypesFromDeclaration: List<UnwrappedType?>? = null
+    override var parameterTypesFromDeclaration: List<UnwrappedType?>? = null
+        private set
+
+    override fun updateParameterTypesFromDeclaration(types: List<KotlinTypeMarker?>?) {
+        @Suppress("UNCHECKED_CAST")
+        types as List<UnwrappedType?>?
+        parameterTypesFromDeclaration = types
+    }
+
+    override fun reviseExpectedType(expectedType: KotlinTypeMarker) {
+        require(expectedType is UnwrappedType)
+        revisedExpectedType = expectedType
+    }
 
     fun setAnalyzed(resolvedLambdaAtom: ResolvedLambdaAtom) {
         setAnalyzedResults(listOf(resolvedLambdaAtom))
@@ -202,8 +211,17 @@ class CallableReferenceWithRevisedExpectedTypeAtom(
 
 class PostponedCallableReferenceAtom(
     eagerCallableReferenceAtom: EagerCallableReferenceAtom
-) : AbstractPostponedCallableReferenceAtom(eagerCallableReferenceAtom.atom, eagerCallableReferenceAtom.expectedType), PostponedAtomWithRevisableExpectedType {
+) : AbstractPostponedCallableReferenceAtom(eagerCallableReferenceAtom.atom, eagerCallableReferenceAtom.expectedType),
+    PostponedCallableReferenceMarker,
+    PostponedAtomWithRevisableExpectedType
+{
     override var revisedExpectedType: UnwrappedType? = null
+        private set
+
+    override fun reviseExpectedType(expectedType: KotlinTypeMarker) {
+        require(expectedType is UnwrappedType)
+        revisedExpectedType = expectedType
+    }
 }
 
 class ResolvedCollectionLiteralAtom(

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

@@ -232,6 +232,7 @@ 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)

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

@@ -161,6 +161,11 @@ 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