[NI] Add substitution of lambda return type in diagnostics

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt

@@ -11,9 +11,7 @@ import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.*
-import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
-import org.jetbrains.kotlin.resolve.calls.CallTransformer
-import org.jetbrains.kotlin.resolve.calls.DiagnosticReporterByTrackingStrategy
+import org.jetbrains.kotlin.resolve.calls.*
 import org.jetbrains.kotlin.resolve.calls.callResolverUtil.getEffectiveExpectedType
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.callUtil.isFakeElement
@@ -135,7 +133,7 @@ class KotlinToResolvedCallTransformer(
                 }
 
                 @Suppress("UNCHECKED_CAST") val resolvedCall =
-                    ktPrimitiveCompleter.completeResolvedCall(candidate, baseResolvedCall.diagnostics) as ResolvedCall<D>
+                    ktPrimitiveCompleter.completeResolvedCall(candidate, baseResolvedCall.completedDiagnostic) as ResolvedCall<D>
                 forwardCallToInferenceSession(baseResolvedCall, context, resolvedCall, tracingStrategy)
 
                 resolvedCall

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

@@ -8,12 +8,16 @@ package org.jetbrains.kotlin.resolve.calls.model
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.resolve.calls.components.*
+import org.jetbrains.kotlin.resolve.calls.inference.buildResultingSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.FreshVariableNewTypeSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
+import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintError
 import org.jetbrains.kotlin.resolve.calls.inference.model.TypeVariableForLambdaReturnType
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
+import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.UnwrappedType
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 /**
  * Call, Callable reference, lambda & function expression, collection literal.
@@ -180,6 +184,17 @@ sealed class CallResolutionResult(
         super.setAnalyzedResults(subResolvedAtoms)
     }
 
+    val completedDiagnostic: List<KotlinCallDiagnostic>
+        get() {
+            val substitutor = constraintSystem.buildResultingSubstitutor()
+            return diagnostics.map {
+                if (it !is NewConstraintError) return@map it
+                val lowerType = it.lowerType.safeAs<KotlinType>()?.unwrap() ?: return@map it
+                val newLowerType = substitutor.safeSubstitute(lowerType)
+                NewConstraintError(newLowerType, it.upperType, it.position)
+            }
+        }
+
     override val atom: ResolutionAtom? get() = null
 
     override fun toString() = "diagnostics: (${diagnostics.joinToString()})"