diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArgumentsAnalyzer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArgumentsAnalyzer.kt
index e0f1d86aed9..4d07fae70ae 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArgumentsAnalyzer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArgumentsAnalyzer.kt
@@ -26,10 +26,7 @@ import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzerC
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
 import org.jetbrains.kotlin.resolve.calls.inference.model.BuilderInferencePosition
-import org.jetbrains.kotlin.types.model.StubTypeMarker
-import org.jetbrains.kotlin.types.model.TypeVariableMarker
-import org.jetbrains.kotlin.types.model.freshTypeConstructor
-import org.jetbrains.kotlin.types.model.safeSubstitute
+import org.jetbrains.kotlin.types.model.*
 
 data class ReturnArgumentsAnalysisResult(
     val returnArguments: Collection<FirStatement>,
@@ -165,26 +162,8 @@ class PostponedArgumentsAnalyzer(
         substitute: (ConeKotlinType) -> ConeKotlinType = c.createSubstituteFunctorForLambdaAnalysis()
     ) {
         val (returnArguments, inferenceSession) = results
-        if (inferenceSession != null) {
-            val storageSnapshot = c.getBuilder().currentStorage()
-
-            val postponedVariables = inferenceSession.inferPostponedVariables(lambda, storageSnapshot, completionMode)
-
-            if (postponedVariables == null) {
-                c.getBuilder().removePostponedVariables()
-            } else {
-                for ((constructor, resultType) in postponedVariables) {
-                    val variableWithConstraints = storageSnapshot.notFixedTypeVariables[constructor] ?: continue
-                    val variable = variableWithConstraints.typeVariable as ConeTypeVariable
-
-                    c.getBuilder().unmarkPostponedVariable(variable)
-                    c.getBuilder().addEqualityConstraint(variable.defaultType, resultType, BuilderInferencePosition)
-                }
-            }
-        }
 
         returnArguments.forEach { c.addSubsystemFromExpression(it) }
-
         val checkerSink: CheckerSink = CheckerSinkImpl(candidate)
 
         val lastExpression = lambda.atom.body?.statements?.lastOrNull() as? FirExpression
@@ -223,6 +202,27 @@ class PostponedArgumentsAnalyzer(
 
         lambda.analyzed = true
         lambda.returnStatements = returnArguments
+
+        if (inferenceSession != null) {
+            val storageSnapshot = c.getBuilder().currentStorage()
+
+            val postponedVariables = inferenceSession.inferPostponedVariables(lambda, storageSnapshot, completionMode)
+
+            if (postponedVariables == null) {
+                c.getBuilder().removePostponedVariables()
+                return
+            }
+
+            for ((constructor, resultType) in postponedVariables) {
+                val variableWithConstraints = storageSnapshot.notFixedTypeVariables[constructor] ?: continue
+                val variable = variableWithConstraints.typeVariable as ConeTypeVariable
+
+                c.getBuilder().unmarkPostponedVariable(variable)
+                c.getBuilder().addSubtypeConstraint(resultType, variable.defaultType(c), BuilderInferencePosition)
+            }
+
+            c.removePostponedTypeVariablesFromConstraints(postponedVariables.keys)
+        }
     }
 
     fun PostponedArgumentsAnalyzerContext.createSubstituteFunctorForLambdaAnalysis(): (ConeKotlinType) -> ConeKotlinType {