From 3ec28f12427c3c6a820b9447a87e09e13e7df4bc Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Sun, 19 Mar 2017 18:22:54 +0300
Subject: [PATCH] Attempt to preserve partial substitution when detecting
 expected type for smart completion

---
 .../testData/smart/lambdaSignature/Map.kt     |  6 ++++
 .../test/JvmSmartCompletionTestGenerated.java |  6 ++++
 .../kotlin/idea/core/ExpectedInfos.kt         | 30 ++++++++++++-------
 3 files changed, 31 insertions(+), 11 deletions(-)
 create mode 100644 idea/idea-completion/testData/smart/lambdaSignature/Map.kt

diff --git a/idea/idea-completion/testData/smart/lambdaSignature/Map.kt b/idea/idea-completion/testData/smart/lambdaSignature/Map.kt
new file mode 100644
index 00000000000..0c43eecbf0c
--- /dev/null
+++ b/idea/idea-completion/testData/smart/lambdaSignature/Map.kt
@@ -0,0 +1,6 @@
+fun foo() {
+    listOf(1, 2).map { <caret> }
+}
+
+// EXIST: "i ->"
+// EXIST: "i: Int ->"
diff --git a/idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JvmSmartCompletionTestGenerated.java b/idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JvmSmartCompletionTestGenerated.java
index ee95d3f969a..2c2cc71ab7f 100644
--- a/idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JvmSmartCompletionTestGenerated.java
+++ b/idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JvmSmartCompletionTestGenerated.java
@@ -1472,6 +1472,12 @@ public class JvmSmartCompletionTestGenerated extends AbstractJvmSmartCompletionT
             doTest(fileName);
         }
 
+        @TestMetadata("Map.kt")
+        public void testMap() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/idea-completion/testData/smart/lambdaSignature/Map.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("MultipleParameters.kt")
         public void testMultipleParameters() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/idea-completion/testData/smart/lambdaSignature/MultipleParameters.kt");
diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ExpectedInfos.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ExpectedInfos.kt
index 9956766bb98..ef9eb4656a0 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ExpectedInfos.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ExpectedInfos.kt
@@ -35,10 +35,7 @@ import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.util.DelegatingCall
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.resolve.descriptorUtil.hasDefaultValue
-import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.types.TypeSubstitutor
-import org.jetbrains.kotlin.types.TypeUtils
-import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.typeUtil.*
 import java.util.*
 
@@ -281,14 +278,15 @@ class ExpectedInfos(
         var descriptor = candidate.resultingDescriptor
         if (descriptor.valueParameters.isEmpty()) return
 
-        var argumentToParameter = call.mapArgumentsToParameters(descriptor)
+        val argumentToParameter = call.mapArgumentsToParameters(descriptor)
         var parameter = argumentToParameter[argument]
+        var parameterType = parameter?.type
 
-        //TODO: we can loose partially inferred substitution here but what to do?
-        if (parameter != null && parameter.type.containsError()) {
+        if (parameterType != null && parameterType.containsError()) {
+            val originalParameter = descriptor.original.valueParameters[parameter!!.index]
+            parameter = originalParameter
             descriptor = descriptor.original
-            parameter = descriptor.valueParameters[parameter.index]!!
-            argumentToParameter = call.mapArgumentsToParameters(descriptor)
+            parameterType = fixSubstitutedType(parameterType, originalParameter.type)
         }
 
         val argumentName = argument.getArgumentName()?.asName
@@ -316,6 +314,7 @@ class ExpectedInfos(
         if (callType == Call.CallType.ARRAY_SET_METHOD) { // last parameter in set is used for value assigned
             if (parameter == parameters.last()) {
                 parameter = null
+                parameterType = null
             }
             parameters = parameters.dropLast(1)
         }
@@ -326,6 +325,7 @@ class ExpectedInfos(
             }
             return
         }
+        parameterType!!
 
         val expectedName = if (descriptor.hasSynthesizedParameterNames()) null else parameter.name.asString()
 
@@ -365,12 +365,11 @@ class ExpectedInfos(
             }
 
             val starOptions = if (!alreadyHasStar) ItemOptions.STAR_PREFIX else ItemOptions.DEFAULT
-            add(ExpectedInfo.createForArgument(parameter.type, expectedName, varargTail, argumentPositionData, starOptions))
+            add(ExpectedInfo.createForArgument(parameterType, expectedName, varargTail, argumentPositionData, starOptions))
         }
         else {
             if (alreadyHasStar) return
 
-            val parameterType = parameter.type
             if (isFunctionLiteralArgument) {
                 if (parameterType.isFunctionType) {
                     add(ExpectedInfo.createForArgument(parameterType, expectedName, null, argumentPositionData))
@@ -382,6 +381,15 @@ class ExpectedInfos(
         }
     }
 
+    private fun fixSubstitutedType(substitutedType: KotlinType, originalType: KotlinType): KotlinType {
+        if (substitutedType.isError) return originalType
+        if (substitutedType.arguments.size != originalType.arguments.size) return originalType
+        val newTypeArguments = substitutedType.arguments.zip(originalType.arguments).map { (argument, originalArgument) ->
+            if (argument.type.containsError()) originalArgument else argument
+        }
+        return substitutedType.replace(newTypeArguments)
+    }
+
     private fun <D : CallableDescriptor> ResolvedCall<D>.allArgumentsMatched()
             = call.valueArguments.none { argument -> getArgumentMapping(argument).isError() && !argument.hasError() /* ignore arguments that has error type */ }