KT-8560 Smart completion for lambda value and more in smart completion

#KT-8560 Fixed

.idea/dictionaries/valentin.xml

@@ -1,3 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
 <component name="ProjectDictionaryState">
   <dictionary name="valentin">
     <words>
@@ -9,6 +10,7 @@
       <w>negatable</w>
       <w>pparent</w>
       <w>processings</w>
+      <w>rbrace</w>
       <w>rbracket</w>
       <w>renderers</w>
       <w>rparenth</w>

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/ExpectedInfos.kt

@@ -23,10 +23,10 @@ import org.jetbrains.kotlin.frontend.di.createContainerForMacros
 import org.jetbrains.kotlin.idea.caches.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.completion.smart.toList
 import org.jetbrains.kotlin.idea.core.mapArgumentsToParameters
-import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
 import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DelegatingBindingTrace
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
@@ -47,12 +47,14 @@ import org.jetbrains.kotlin.types.JetType
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingUtils
 import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
+import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
 import java.util.HashSet
 
 enum class Tail {
     COMMA,
     RPARENTH,
-    ELSE
+    ELSE,
+    RBRACE
 }
 
 data class ItemOptions(val starPrefix: Boolean) {
@@ -130,6 +132,15 @@ class ExpectedInfos(
     }
 
     public fun calculateForArgument(call: Call, argument: ValueArgument): Collection<ExpectedInfo>? {
+        val callExpression = (call.callElement as? JetExpression)?.getQualifiedExpressionForSelectorOrThis()
+        var expectedTypes = callExpression?.let { calculate(it) }?.map { it.type }
+        if (expectedTypes == null || expectedTypes.isEmpty()) {
+            expectedTypes = listOf(TypeUtils.NO_EXPECTED_TYPE)
+        }
+        return expectedTypes.flatMap { calculateForArgument(call, it, argument) ?: emptyList() }.toSet()
+    }
+
+    private fun calculateForArgument(call: Call, callExpectedType: JetType, argument: ValueArgument): Collection<ExpectedInfo>? {
         val argumentIndex = call.getValueArguments().indexOf(argument)
         assert(argumentIndex >= 0) {
             "Could not find argument '$argument' among arguments of call: $call"
@@ -138,8 +149,8 @@ class ExpectedInfos(
         val isFunctionLiteralArgument = argument is FunctionLiteralArgument
         val argumentPosition = ArgumentPosition(argumentIndex, argumentName, isFunctionLiteralArgument)
 
-        val callElement = call.getCallElement()
-        val calleeExpression = call.getCalleeExpression()
+        val project = call.callElement.getProject()
+        val moduleDescriptor = resolutionFacade.findModuleDescriptor(call.callElement)
 
         // leave only arguments before the current one
         val truncatedCall = object : DelegatingCall(call) {
@@ -149,19 +160,15 @@ class ExpectedInfos(
             override fun getFunctionLiteralArguments() = emptyList<FunctionLiteralArgument>()
             override fun getValueArgumentList() = null
         }
-        val resolutionScope = bindingContext[BindingContext.RESOLUTION_SCOPE, calleeExpression] ?: return null //TODO: discuss it
+        val resolutionScope = bindingContext[BindingContext.RESOLUTION_SCOPE, call.calleeExpression] ?: return null //TODO: discuss it
 
-        val expectedType = (callElement as? JetExpression)?.let { bindingContext[BindingContext.EXPECTED_EXPRESSION_TYPE, it] } ?: TypeUtils.NO_EXPECTED_TYPE
-        val dataFlowInfo = bindingContext.getDataFlowInfo(calleeExpression)
+        val dataFlowInfo = bindingContext.getDataFlowInfo(call.calleeExpression)
         val bindingTrace = DelegatingBindingTrace(bindingContext, "Temporary trace for completion")
-        val context = BasicCallResolutionContext.create(bindingTrace, resolutionScope, truncatedCall, expectedType, dataFlowInfo,
+        val context = BasicCallResolutionContext.create(bindingTrace, resolutionScope, truncatedCall, callExpectedType, dataFlowInfo,
                                                         ContextDependency.INDEPENDENT, CheckArgumentTypesMode.CHECK_VALUE_ARGUMENTS,
                                                         CompositeChecker(listOf()), SymbolUsageValidator.Empty, AdditionalTypeChecker.Composite(listOf()), false)
         val callResolutionContext = context.replaceCollectAllCandidates(true)
-        val callResolver = createContainerForMacros(
-                callElement.getProject(),
-                resolutionFacade.findModuleDescriptor(callElement)
-        ).callResolver
+        val callResolver = createContainerForMacros(project, moduleDescriptor).callResolver
         val results: OverloadResolutionResults<FunctionDescriptor> = callResolver.resolveFunctionCall(callResolutionContext)
 
         val expectedInfos = HashSet<ExpectedInfo>()
@@ -226,7 +233,7 @@ class ExpectedInfos(
                 if (alreadyHasStar) continue
 
                 val parameterType = if (useHeuristicSignatures)
-                    HeuristicSignatures.correctedParameterType(descriptor, parameter, moduleDescriptor, callElement.getProject()) ?: parameter.getType()
+                    HeuristicSignatures.correctedParameterType(descriptor, parameter, moduleDescriptor, project) ?: parameter.getType()
                 else
                     parameter.getType()
 
@@ -315,9 +322,21 @@ class ExpectedInfos(
     }
 
     private fun calculateForBlockExpression(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
-        val block = expressionWithType.getParent() as? JetBlockExpression ?: return null
-        if (expressionWithType != block.getStatements().last()) return null
-        return calculate(block)?.map { ExpectedInfo(it.type, it.expectedName, null) }
+        val block = expressionWithType.parent as? JetBlockExpression ?: return null
+        if (expressionWithType != block.statements.last()) return null
+
+        val functionLiteral = block.parent as? JetFunctionLiteral
+        if (functionLiteral != null) {
+            val literalExpression = functionLiteral.parent as JetFunctionLiteralExpression
+            return calculate(literalExpression)
+                    ?.map { it.type }
+                    ?.filter { KotlinBuiltIns.isExactFunctionOrExtensionFunctionType(it) }
+                    ?.map { KotlinBuiltIns.getReturnTypeFromFunctionType(it) }
+                    ?.map { ExpectedInfo(it, null, Tail.RBRACE) }
+        }
+        else {
+            return calculate(block)?.map { ExpectedInfo(it.type, it.expectedName, null) }
+        }
     }
 
     private fun calculateForWhenEntryValue(expressionWithType: JetExpression): Collection<ExpectedInfo>? {

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/handlers/WithTailInsertHandler.kt

@@ -52,26 +52,29 @@ class WithTailInsertHandler(val tailText: String,
 
         val moveCaret = context.getEditor().getCaretModel().getOffset() == tailOffset
 
+        //TODO: analyze parenthesis balance to decide whether to replace or not
+        var insert = true
         if (overwriteText) {
             fun isCharAt(offset: Int, c: Char) = offset < document.getTextLength() && document.getCharsSequence().charAt(offset) == c
             fun isTextAt(offset: Int, text: String) = offset + text.length() <= document.getTextLength() && document.getText(TextRange(offset, offset + text.length())) == text
 
-            if (spaceBefore && isCharAt(tailOffset, ' ')) {
-                document.deleteString(tailOffset, tailOffset + 1)
-            }
+            var offset = document.charsSequence.skipSpacesAndLineBreaks(tailOffset)
+            if (isTextAt(offset, tailText)) {
+                insert = false
+                offset += tailText.length()
+                tailOffset = offset
 
-            if (isTextAt(tailOffset, tailText)) {
-                document.deleteString(tailOffset, tailOffset + tailText.length())
-
-                if (spaceAfter && isCharAt(tailOffset, ' ')) {
-                    document.deleteString(tailOffset, tailOffset + 1)
+                if (spaceAfter && isCharAt(offset, ' ')) {
+                    document.deleteString(offset, offset + 1)
                 }
             }
         }
 
-
-        var textToInsert = tailText
-        if (spaceBefore) textToInsert = " " + textToInsert
+        var textToInsert = ""
+        if (insert) {
+            textToInsert = tailText
+            if (spaceBefore) textToInsert = " " + textToInsert
+        }
         if (spaceAfter) textToInsert += " "
 
         document.insertString(tailOffset, textToInsert)
@@ -88,6 +91,7 @@ class WithTailInsertHandler(val tailText: String,
     companion object {
         fun commaTail() = WithTailInsertHandler(",", spaceBefore = false, spaceAfter = true /*TODO: use code style option*/)
         fun rparenthTail() = WithTailInsertHandler(")", spaceBefore = false, spaceAfter = false)
+        fun rbraceTail() = WithTailInsertHandler("}", spaceBefore = true, spaceAfter = false)
         fun elseTail() = WithTailInsertHandler("else", spaceBefore = true, spaceAfter = true)
         fun eqTail() = WithTailInsertHandler("=", spaceBefore = true, spaceAfter = true) /*TODO: use code style options*/
         fun spaceTail() = WithTailInsertHandler(" ", spaceBefore = false, spaceAfter = false, overwriteText = false)

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/handlers/handlerUtils.kt

@@ -28,3 +28,6 @@ fun CharSequence.indexOfSkippingSpace(c: Char, startIndex: Int): Int? {
 fun CharSequence.skipSpaces(index: Int): Int
         = (index..length() - 1).firstOrNull { val c = this[it]; c != ' ' && c != '\t' } ?: this.length()
 
+fun CharSequence.skipSpacesAndLineBreaks(index: Int): Int
+        = (index..length() - 1).firstOrNull { val c = this[it]; c != ' ' && c != '\t' && c != '\n' && c != '\r' } ?: this.length()
+

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/smart/Utils.kt

@@ -80,6 +80,12 @@ fun LookupElement.addTail(tail: Tail?): LookupElement {
                 WithTailInsertHandler.elseTail().handleInsert(context, getDelegate())
             }
         }
+
+        Tail.RBRACE -> object: LookupElementDecorator<LookupElement>(this) {
+            override fun handleInsert(context: InsertionContext) {
+                WithTailInsertHandler.rbraceTail().handleInsert(context, getDelegate())
+            }
+        }
     }
 }
 

idea/idea-completion/testData/handlers/smart/LambdaValue1.kt

@@ -0,0 +1,5 @@
+fun foo(list: List<String>) {
+    list.filter { it.<caret> }
+}
+
+// ELEMENT: isEmpty

idea/idea-completion/testData/handlers/smart/LambdaValue1.kt.after

@@ -0,0 +1,5 @@
+fun foo(list: List<String>) {
+    list.filter { it.isEmpty() }<caret>
+}
+
+// ELEMENT: isEmpty

idea/idea-completion/testData/handlers/smart/LambdaValue2.kt

@@ -0,0 +1,7 @@
+fun foo(list: List<String>) {
+    list.filter {
+        it.<caret>
+    }
+}
+
+// ELEMENT: isEmpty

idea/idea-completion/testData/handlers/smart/LambdaValue2.kt.after

@@ -0,0 +1,7 @@
+fun foo(list: List<String>) {
+    list.filter {
+        it.isEmpty()
+    }<caret>
+}
+
+// ELEMENT: isEmpty

idea/idea-completion/testData/smart/LambdaValue1.kt

@@ -0,0 +1,7 @@
+fun foo(list: List<String>) {
+    list.filter { it.<caret> }
+}
+
+// EXIST: isEmpty
+// EXIST: isBlank
+// ABSENT: substring

idea/idea-completion/testData/smart/LambdaValue2.kt

@@ -0,0 +1,6 @@
+fun foo(list: List<String>): Collection<Int> {
+    return list.map { it.<caret> }
+}
+
+// EXIST: length
+// ABSENT: isEmpty

idea/idea-completion/testData/smart/LambdaValue3.kt

@@ -0,0 +1,13 @@
+fun foo(list: List<String>): Collection<Int> {
+    bar(list.map { it.<caret> })
+}
+
+fun bar(p: Collection<Int>) {
+}
+
+fun bar(p: Collection<String>, b: Boolean) {
+}
+
+// EXIST: length
+// EXIST: substring
+// ABSENT: isEmpty

idea/idea-completion/testData/smart/MapTo.kt

@@ -0,0 +1,7 @@
+fun foo(list: List<String>, intList: MutableList<Int>, stringList: MutableList<String>): Collection<Int> {
+    return list.mapTo(<caret>)
+}
+
+// EXIST: intList
+// EXIST: arrayListOf
+// ABSENT: stringList

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JvmSmartCompletionTestGenerated.java

@@ -257,6 +257,30 @@ public class JvmSmartCompletionTestGenerated extends AbstractJvmSmartCompletionT
         doTest(fileName);
     }
 
+    @TestMetadata("LambdaValue1.kt")
+    public void testLambdaValue1() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/smart/LambdaValue1.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("LambdaValue2.kt")
+    public void testLambdaValue2() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/smart/LambdaValue2.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("LambdaValue3.kt")
+    public void testLambdaValue3() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/smart/LambdaValue3.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("MapTo.kt")
+    public void testMapTo() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/smart/MapTo.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("MethodCallArgument.kt")
     public void testMethodCallArgument() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/smart/MethodCallArgument.kt");

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/handlers/SmartCompletionHandlerTestGenerated.java

@@ -485,6 +485,18 @@ public class SmartCompletionHandlerTestGenerated extends AbstractSmartCompletion
         doTest(fileName);
     }
 
+    @TestMetadata("LambdaValue1.kt")
+    public void testLambdaValue1() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/handlers/smart/LambdaValue1.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("LambdaValue2.kt")
+    public void testLambdaValue2() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/handlers/smart/LambdaValue2.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("LastNonOptionalParamIsFunction.kt")
     public void testLastNonOptionalParamIsFunction() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/handlers/smart/LastNonOptionalParamIsFunction.kt");