KT-4977 Smart completion should work for arguments in brackets

#KT-4977 Fixed

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

@@ -61,6 +61,7 @@ import java.util.*
 enum class Tail {
     COMMA,
     RPARENTH,
+    RBRACKET,
     ELSE,
     RBRACE
 }
@@ -158,6 +159,7 @@ class ExpectedInfos(
     public fun calculate(expressionWithType: JetExpression): Collection<ExpectedInfo> {
         val expectedInfos = calculateForArgument(expressionWithType)
                             ?: calculateForFunctionLiteralArgument(expressionWithType)
+                            ?: calculateForIndexingArgument(expressionWithType)
                             ?: calculateForEqAndAssignment(expressionWithType)
                             ?: calculateForIf(expressionWithType)
                             ?: calculateForElvis(expressionWithType)
@@ -189,6 +191,15 @@ class ExpectedInfos(
         return calculateForArgument(callExpression, literalArgument)
     }
 
+    private fun calculateForIndexingArgument(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
+        val containerNode = expressionWithType.parent as? JetContainerNode ?: return null
+        val arrayAccessExpression = containerNode.parent as? JetArrayAccessExpression ?: return null
+        if (containerNode != arrayAccessExpression.indicesNode) return null
+        val call = arrayAccessExpression.getCall(bindingContext) ?: return null
+        val argument = call.valueArguments.firstOrNull { it.getArgumentExpression() == expressionWithType } ?: return null
+        return calculateForArgument(call, argument)
+    }
+
     private fun calculateForArgument(callElement: JetCallElement, argument: ValueArgument): Collection<ExpectedInfo>? {
         val call = callElement.getCall(bindingContext) ?: return null
         return calculateForArgument(call, argument)
@@ -312,6 +323,18 @@ class ExpectedInfos(
             ArgumentPositionData.Positional(descriptor, argumentIndex, isFunctionLiteralArgument, namedArgumentCandidates)
         }
 
+        val callType = call.callType
+        val isArrayAccess = callType == Call.CallType.ARRAY_GET_METHOD || callType == Call.CallType.ARRAY_SET_METHOD
+        val rparenthTail = if (isArrayAccess) Tail.RBRACKET else Tail.RPARENTH
+
+        var parameters = descriptor.valueParameters
+        if (callType == Call.CallType.ARRAY_SET_METHOD) { // last parameter in set is used for value assigned
+            if (parameter == parameters.last()) {
+                parameter = null
+            }
+            parameters = parameters.dropLast(1)
+        }
+
         if (parameter == null) {
             if (argumentPositionData is ArgumentPositionData.Positional && argumentPositionData.namedArgumentCandidates.isNotEmpty()) {
                 add(ExpectedInfo.createForNamedArgumentExpected(argumentPositionData))
@@ -321,19 +344,17 @@ class ExpectedInfos(
 
         val expectedName = if (descriptor.hasSynthesizedParameterNames()) null else parameter.getName().asString()
 
-        val parameters = descriptor.valueParameters
-
         fun needCommaForParameter(parameter: ValueParameterDescriptor): Boolean {
             if (parameter.hasDefaultValue()) return false // parameter is optional
             if (parameter.getVarargElementType() != null) return false // vararg arguments list can be empty
             // last parameter of functional type can be placed outside parenthesis:
-            if (parameter == parameters.last() && KotlinBuiltIns.isExactFunctionOrExtensionFunctionType(parameter.getType())) return false
+            if (!isArrayAccess && parameter == parameters.last() && KotlinBuiltIns.isExactFunctionOrExtensionFunctionType(parameter.getType())) return false
             return true
         }
 
         val tail = if (argumentName == null) {
             if (parameter == parameters.last())
-                Tail.RPARENTH //TODO: support square brackets
+                rparenthTail
             else if (parameters.dropWhile { it != parameter }.drop(1).any(::needCommaForParameter))
                 Tail.COMMA
             else
@@ -349,7 +370,7 @@ class ExpectedInfos(
         if (varargElementType != null) {
             if (isFunctionLiteralArgument) return
 
-            val varargTail = if (argumentName == null && tail == Tail.RPARENTH)
+            val varargTail = if (argumentName == null && tail == rparenthTail)
                 null /* even if it's the last parameter, there can be more arguments for the same parameter */
             else
                 tail
@@ -391,7 +412,7 @@ class ExpectedInfos(
         val usedParameterNames = (argumentToParameter.values().map { it.getName() } + listOf(argumentName)).toSet()
         val notUsedParameters = descriptor.getValueParameters().filter { it.getName() !in usedParameterNames }
         return if (notUsedParameters.isEmpty())
-            Tail.RPARENTH
+            Tail.RPARENTH // named arguments no supported for []
         else if (notUsedParameters.all { it.hasDefaultValue() })
             null
         else

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

@@ -88,6 +88,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 rbracketTail() = 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*/

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

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

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

@@ -0,0 +1,5 @@
+fun foo(map: Map<String, Int>, p: String) {
+    map[<caret>]
+}
+
+// ELEMENT: p

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

@@ -0,0 +1,5 @@
+fun foo(map: Map<String, Int>, p: String) {
+    map[p]<caret>
+}
+
+// ELEMENT: p

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

@@ -0,0 +1,5 @@
+fun foo(map: Map<String, Int>, p: String) {
+    map[<caret>] = 1
+}
+
+// ELEMENT: p

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

@@ -0,0 +1,5 @@
+fun foo(map: Map<String, Int>, p: String) {
+    map[p]<caret> = 1
+}
+
+// ELEMENT: p

idea/idea-completion/testData/smart/heuristicSignatures/GetWithBrackets.kt

@@ -0,0 +1,6 @@
+fun foo(map: Map<String, Int>, p1: Any, p2: String) {
+    map[<caret>]
+}
+
+// ABSENT: p1
+// EXIST: p2

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

@@ -1012,6 +1012,12 @@ public class JvmSmartCompletionTestGenerated extends AbstractJvmSmartCompletionT
             doTest(fileName);
         }
 
+        @TestMetadata("GetWithBrackets.kt")
+        public void testGetWithBrackets() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/smart/heuristicSignatures/GetWithBrackets.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("InOperator.kt")
         public void testInOperator() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/smart/heuristicSignatures/InOperator.kt");

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

@@ -389,6 +389,12 @@ public class SmartCompletionHandlerTestGenerated extends AbstractSmartCompletion
         doTest(fileName);
     }
 
+    @TestMetadata("GetWithBrackets.kt")
+    public void testGetWithBrackets() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/handlers/smart/GetWithBrackets.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("IfCondition.kt")
     public void testIfCondition() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/handlers/smart/IfCondition.kt");
@@ -695,6 +701,12 @@ public class SmartCompletionHandlerTestGenerated extends AbstractSmartCompletion
         doTest(fileName);
     }
 
+    @TestMetadata("SetWithBrackets.kt")
+    public void testSetWithBrackets() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/handlers/smart/SetWithBrackets.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("TabReplaceComma1.kt")
     public void testTabReplaceComma1() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/handlers/smart/TabReplaceComma1.kt");