From ad1b8a6fb1e0eb5cf166fab1315ae62f137b2c1b Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Fri, 26 Apr 2019 13:53:30 +0300
Subject: [PATCH] [NI] Add substitution of lambda return type in diagnostics

---
 .../tower/KotlinToResolvedCallTransformer.kt      |  6 ++----
 .../kotlin/resolve/calls/model/ResolutionAtoms.kt | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt
index d65e527ab6a..2b8a15838e9 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt
+++ b/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
diff --git a/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/ResolutionAtoms.kt b/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/ResolutionAtoms.kt
index 64b3b3dc824..9176e65e0bd 100644
--- a/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/ResolutionAtoms.kt
+++ b/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()})"