Smart completion for then and else branches of if

#KT-4910 Fixed

idea/src/org/jetbrains/jet/plugin/completion/ExpectedInfos.kt

@@ -42,19 +42,26 @@ import org.jetbrains.jet.lang.descriptors.ModuleDescriptor
 import org.jetbrains.jet.lang.types.JetType
 import org.jetbrains.jet.lang.psi.JetBinaryExpression
 import org.jetbrains.jet.lexer.JetTokens
+import org.jetbrains.jet.lang.psi.JetIfExpression
+import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns
+import org.jetbrains.jet.lang.psi.JetContainerNode
+import org.jetbrains.jet.plugin.completion.smart.isSubtypeOf
 
 enum class Tail {
     COMMA
     PARENTHESIS
+    ELSE
 }
 
 data class ExpectedInfo(val `type`: JetType, val tail: Tail?)
 
 class ExpectedInfos(val bindingContext: BindingContext, val moduleDescriptor: ModuleDescriptor) {
     public fun calculate(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
-        return calculateForArgument(expressionWithType)
+        val forArgument = calculateForArgument(expressionWithType)
+        return forArgument
             ?: calculateForFunctionLiteralArgument(expressionWithType)
             ?: calculateForEq(expressionWithType)
+            ?: calculateForIf(expressionWithType)
             ?: getFromBindingContext(expressionWithType)
     }
 
@@ -142,6 +149,26 @@ class ExpectedInfos(val bindingContext: BindingContext, val moduleDescriptor: Mo
         return null
     }
 
+    private fun calculateForIf(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
+        val ifExpression = (expressionWithType.getParent() as? JetContainerNode)?.getParent() as? JetIfExpression ?: return null
+        return when (expressionWithType) {
+            ifExpression.getCondition() -> listOf(ExpectedInfo(KotlinBuiltIns.getInstance().getBooleanType(), Tail.PARENTHESIS))
+
+            ifExpression.getThen() -> calculate(ifExpression)?.map { ExpectedInfo(it.`type`, Tail.ELSE) }
+
+            ifExpression.getElse() -> {
+                val ifExpectedInfo = calculate(ifExpression)
+                val thenType = bindingContext[BindingContext.EXPRESSION_TYPE, ifExpression.getThen()]
+                if (thenType != null)
+                    ifExpectedInfo?.filter { it.`type`.isSubtypeOf(thenType) }
+                else
+                    ifExpectedInfo
+            }
+
+            else -> return null
+        }
+    }
+
     private fun getFromBindingContext(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
         val expectedType = bindingContext[BindingContext.EXPECTED_EXPRESSION_TYPE, expressionWithType] ?: return null
         return listOf(ExpectedInfo(expectedType, null))

idea/src/org/jetbrains/jet/plugin/completion/handlers/WithTailInsertHandler.kt

@@ -23,10 +23,12 @@ import com.intellij.codeInsight.AutoPopupController
 import com.intellij.codeInsight.lookup.Lookup
 import org.jetbrains.jet.plugin.completion.smart.SmartCompletion
 import org.jetbrains.jet.plugin.completion.smart.KEEP_OLD_ARGUMENT_LIST_ON_TAB_KEY
+import com.intellij.openapi.util.TextRange
 
-abstract class WithTailInsertHandlerBase : InsertHandler<LookupElement> {
+class WithTailInsertHandler(val tailText: String,
-    protected abstract fun insertTail(context: InsertionContext, offset: Int, moveCaret: Boolean)
+                            val spaceBefore: Boolean,
-
+                            val spaceAfter: Boolean,
+                            val overwriteText: Boolean = true) : InsertHandler<LookupElement> {
     override fun handleInsert(context: InsertionContext, item: LookupElement) {
         val document = context.getDocument()
 
@@ -38,40 +40,39 @@ abstract class WithTailInsertHandlerBase : InsertHandler<LookupElement> {
             val offset = context.getOffsetMap().getOffset(SmartCompletion.OLD_ARGUMENTS_REPLACEMENT_OFFSET)
             if (offset != -1) tailOffset = offset
         }
-        insertTail(context, tailOffset, context.getEditor().getCaretModel().getOffset() == tailOffset)
-    }
-}
 
-class WithTailCharInsertHandler(val tailChar: Char, val spaceAfter: Boolean) : WithTailInsertHandlerBase() {
+        val moveCaret = context.getEditor().getCaretModel().getOffset() == tailOffset
-    override fun insertTail(context: InsertionContext, offset: Int, moveCaret: Boolean) {
+
-        val document = context.getDocument()
         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 (isCharAt(offset, tailChar)) {
+        if (overwriteText) {
-            document.deleteString(offset, offset + 1)
+            if (spaceBefore && isCharAt(tailOffset, ' ')) {
+                document.deleteString(tailOffset, tailOffset + 1)
+            }
 
-            if (spaceAfter && isCharAt(offset, ' ')) {
+            if (isTextAt(tailOffset, tailText)) {
-                document.deleteString(offset, offset + 1)
+                document.deleteString(tailOffset, tailOffset + tailText.length)
+
+                if (spaceAfter && isCharAt(tailOffset, ' ')) {
+                    document.deleteString(tailOffset, tailOffset + 1)
+                }
             }
         }
 
-        val textToInsert = if (spaceAfter) tailChar + " " else tailChar.toString()
-        document.insertString(offset, textToInsert)
-        if (moveCaret) {
-            context.getEditor().getCaretModel().moveToOffset(offset + textToInsert.length)
 
-            if (tailChar == ',') {
+        var textToInsert = tailText
+        if (spaceBefore) textToInsert = " " + textToInsert
+        if (spaceAfter) textToInsert += " "
+
+        document.insertString(tailOffset, textToInsert)
+
+        if (moveCaret) {
+            context.getEditor().getCaretModel().moveToOffset(tailOffset + textToInsert.length)
+
+            if (tailText == ",") {
                 AutoPopupController.getInstance(context.getProject())?.autoPopupParameterInfo(context.getEditor(), null)
             }
         }
     }
 }
-
-class WithTailStringInsertHandler(val tail: String) : WithTailInsertHandlerBase() {
-    override fun insertTail(context: InsertionContext, offset: Int, moveCaret: Boolean) {
-        context.getDocument().insertString(offset, tail)
-        if (moveCaret) {
-            context.getEditor().getCaretModel().moveToOffset(offset + tail.length)
-        }
-    }
-}

idea/src/org/jetbrains/jet/plugin/completion/smart/Utils.kt

@@ -24,7 +24,6 @@ import org.jetbrains.jet.lang.psi.JetFile
 import org.jetbrains.jet.plugin.codeInsight.ShortenReferences
 import java.util.HashSet
 import com.intellij.codeInsight.lookup.LookupElementDecorator
-import org.jetbrains.jet.plugin.completion.handlers.WithTailCharInsertHandler
 import com.intellij.codeInsight.lookup.AutoCompletionPolicy
 import org.jetbrains.jet.lang.descriptors.FunctionDescriptor
 import org.jetbrains.jet.lang.types.JetType
@@ -37,7 +36,7 @@ import com.intellij.openapi.util.Key
 import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor
 import org.jetbrains.jet.plugin.project.ResolveSessionForBodies
 import org.jetbrains.jet.lang.resolve.BindingContext
-import org.jetbrains.jet.plugin.completion.handlers.WithTailStringInsertHandler
+import org.jetbrains.jet.plugin.completion.handlers.WithTailInsertHandler
 
 class ArtificialElementInsertHandler(
         val textBeforeCaret: String, val textAfterCaret: String, val shortenRefs: Boolean) : InsertHandler<LookupElement>{
@@ -70,13 +69,19 @@ fun LookupElement.addTail(tail: Tail?): LookupElement {
 
         Tail.COMMA -> object: LookupElementDecorator<LookupElement>(this) {
             override fun handleInsert(context: InsertionContext) {
-                WithTailCharInsertHandler(',', true /*TODO: use code style option*/).handleInsert(context, getDelegate())
+                WithTailInsertHandler(",", spaceBefore = false, spaceAfter = true /*TODO: use code style option*/).handleInsert(context, getDelegate())
             }
         }
 
         Tail.PARENTHESIS -> object: LookupElementDecorator<LookupElement>(this) {
             override fun handleInsert(context: InsertionContext) {
-                WithTailCharInsertHandler(')', false).handleInsert(context, getDelegate())
+                handlers.WithTailInsertHandler(")", spaceBefore = false, spaceAfter = false).handleInsert(context, getDelegate())
+            }
+        }
+
+        Tail.ELSE -> object: LookupElementDecorator<LookupElement>(this) {
+            override fun handleInsert(context: InsertionContext) {
+                handlers.WithTailInsertHandler("else", spaceBefore = true, spaceAfter = true).handleInsert(context, getDelegate())
             }
         }
     }
@@ -132,7 +137,7 @@ fun MutableCollection<LookupElement>.addLookupElementsForNullable(factory: () ->
                 presentation.setItemText("!! " + presentation.getItemText())
             }
             override fun handleInsert(context: InsertionContext) {
-                WithTailStringInsertHandler("!!").handleInsert(context, getDelegate())
+                WithTailInsertHandler("!!", spaceBefore = false, spaceAfter = false).handleInsert(context, getDelegate())
             }
         }
         lookupElement = lookupElement!!.suppressAutoInsertion()
@@ -147,7 +152,7 @@ fun MutableCollection<LookupElement>.addLookupElementsForNullable(factory: () ->
                 presentation.setItemText("?: " + presentation.getItemText())
             }
             override fun handleInsert(context: InsertionContext) {
-                WithTailStringInsertHandler(" ?: ").handleInsert(context, getDelegate()) //TODO: code style
+                handlers.WithTailInsertHandler("?:", spaceBefore = true, spaceAfter = true).handleInsert(context, getDelegate()) //TODO: code style
             }
         }
         lookupElement = lookupElement!!.suppressAutoInsertion()

idea/testData/completion/handlers/smart/IfCondition.kt

@@ -0,0 +1,5 @@
+fun bar(b: Boolean){
+    if (<caret>
+}
+
+// ELEMENT: b

idea/testData/completion/handlers/smart/IfCondition.kt.after

@@ -0,0 +1,5 @@
+fun bar(b: Boolean){
+    if (b)<caret>
+}
+
+// ELEMENT: b

idea/testData/completion/handlers/smart/IfValue1.kt

@@ -0,0 +1,8 @@
+fun foo(s: String, i: Int){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String){
+    foo(if (b) <caret>)
+}
+
+// ELEMENT: s

idea/testData/completion/handlers/smart/IfValue1.kt.after

@@ -0,0 +1,8 @@
+fun foo(s: String, i: Int){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String){
+    foo(if (b) s else <caret>)
+}
+
+// ELEMENT: s

idea/testData/completion/handlers/smart/IfValue2.kt

@@ -0,0 +1,8 @@
+fun foo(s: String, i: Int){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String){
+    foo(if (b) "abc" else <caret>)
+}
+
+// ELEMENT: s

idea/testData/completion/handlers/smart/IfValue2.kt.after

@@ -0,0 +1,8 @@
+fun foo(s: String, i: Int){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String){
+    foo(if (b) "abc" else s, <caret>)
+}
+
+// ELEMENT: s

idea/testData/completion/handlers/smart/IfValue3.kt

@@ -0,0 +1,9 @@
+fun foo(s: String, i: Int){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String){
+    foo(if (b) <caret>xxx else "")
+}
+
+// ELEMENT: s
+// CHAR: \t

idea/testData/completion/handlers/smart/IfValue3.kt.after

@@ -0,0 +1,9 @@
+fun foo(s: String, i: Int){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String){
+    foo(if (b) s else <caret>"")
+}
+
+// ELEMENT: s
+// CHAR: \t

idea/testData/completion/smart/IfCondition.kt

@@ -0,0 +1,6 @@
+fun bar(b: Boolean, c: Char){
+    if (<caret>
+}
+
+// EXIST: b
+// ABSENT: c

idea/testData/completion/smart/IfValue1.kt

@@ -0,0 +1,10 @@
+fun foo(s: String){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String, c: Char){
+    foo(if (b) <caret>)
+}
+
+// EXIST: s
+// EXIST: c
+// ABSENT: b

idea/testData/completion/smart/IfValue2.kt

@@ -0,0 +1,10 @@
+fun foo(s: String){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String, c: Char){
+    foo(if (b) "abc" else <caret>)
+}
+
+// EXIST: s
+// ABSENT: c
+// ABSENT: b

idea/testData/completion/smart/IfValue3.kt

@@ -0,0 +1,10 @@
+fun foo(s: String){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String, c: Char){
+    foo(if (b) xxx else <caret>)
+}
+
+// EXIST: s
+// EXIST: c
+// ABSENT: b

idea/tests/org/jetbrains/jet/completion/JvmSmartCompletionTestGenerated.java

@@ -191,6 +191,26 @@ public class JvmSmartCompletionTestGenerated extends AbstractJvmSmartCompletionT
         doTest("idea/testData/completion/smart/FunctionReference9.kt");
     }
     
+    @TestMetadata("IfCondition.kt")
+    public void testIfCondition() throws Exception {
+        doTest("idea/testData/completion/smart/IfCondition.kt");
+    }
+    
+    @TestMetadata("IfValue1.kt")
+    public void testIfValue1() throws Exception {
+        doTest("idea/testData/completion/smart/IfValue1.kt");
+    }
+    
+    @TestMetadata("IfValue2.kt")
+    public void testIfValue2() throws Exception {
+        doTest("idea/testData/completion/smart/IfValue2.kt");
+    }
+    
+    @TestMetadata("IfValue3.kt")
+    public void testIfValue3() throws Exception {
+        doTest("idea/testData/completion/smart/IfValue3.kt");
+    }
+    
     @TestMetadata("InaccessibleConstructor.kt")
     public void testInaccessibleConstructor() throws Exception {
         doTest("idea/testData/completion/smart/InaccessibleConstructor.kt");

idea/tests/org/jetbrains/jet/completion/handlers/SmartCompletionHandlerTestGenerated.java

@@ -206,6 +206,26 @@ public class SmartCompletionHandlerTestGenerated extends AbstractSmartCompletion
         doTest("idea/testData/completion/handlers/smart/GenericFunction.kt");
     }
     
+    @TestMetadata("IfCondition.kt")
+    public void testIfCondition() throws Exception {
+        doTest("idea/testData/completion/handlers/smart/IfCondition.kt");
+    }
+    
+    @TestMetadata("IfValue1.kt")
+    public void testIfValue1() throws Exception {
+        doTest("idea/testData/completion/handlers/smart/IfValue1.kt");
+    }
+    
+    @TestMetadata("IfValue2.kt")
+    public void testIfValue2() throws Exception {
+        doTest("idea/testData/completion/handlers/smart/IfValue2.kt");
+    }
+    
+    @TestMetadata("IfValue3.kt")
+    public void testIfValue3() throws Exception {
+        doTest("idea/testData/completion/handlers/smart/IfValue3.kt");
+    }
+    
     @TestMetadata("JavaEnumMemberInsertsImport.kt")
     public void testJavaEnumMemberInsertsImport() throws Exception {
         doTest("idea/testData/completion/handlers/smart/JavaEnumMemberInsertsImport.kt");