From 1b034c5fc3effe0b57b9e933d782871850d3a7b6 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Wed, 30 Apr 2014 20:45:26 +0400
Subject: [PATCH] Smart completion works for block's result

---
 .../jet/plugin/completion/ExpectedInfos.kt        |  8 ++++++++
 .../completion/handlers/smart/IfValueInBlock.kt   | 12 ++++++++++++
 .../handlers/smart/IfValueInBlock.kt.after        | 12 ++++++++++++
 idea/testData/completion/smart/IfValueInBlock1.kt | 13 +++++++++++++
 idea/testData/completion/smart/IfValueInBlock2.kt | 15 +++++++++++++++
 .../JvmSmartCompletionTestGenerated.java          | 10 ++++++++++
 .../SmartCompletionHandlerTestGenerated.java      |  5 +++++
 7 files changed, 75 insertions(+)
 create mode 100644 idea/testData/completion/handlers/smart/IfValueInBlock.kt
 create mode 100644 idea/testData/completion/handlers/smart/IfValueInBlock.kt.after
 create mode 100644 idea/testData/completion/smart/IfValueInBlock1.kt
 create mode 100644 idea/testData/completion/smart/IfValueInBlock2.kt

diff --git a/idea/src/org/jetbrains/jet/plugin/completion/ExpectedInfos.kt b/idea/src/org/jetbrains/jet/plugin/completion/ExpectedInfos.kt
index b975537776a..09dc57e2726 100644
--- a/idea/src/org/jetbrains/jet/plugin/completion/ExpectedInfos.kt
+++ b/idea/src/org/jetbrains/jet/plugin/completion/ExpectedInfos.kt
@@ -50,6 +50,7 @@ import org.jetbrains.jet.lang.resolve.calls.util.DelegatingCall
 import org.jetbrains.jet.lang.resolve.calls.util.noErrorsInValueArguments
 import org.jetbrains.jet.lang.resolve.calls.util.hasUnmappedParameters
 import org.jetbrains.jet.lang.descriptors.Visibilities
+import org.jetbrains.jet.lang.psi.JetBlockExpression
 
 enum class Tail {
     COMMA
@@ -66,6 +67,7 @@ class ExpectedInfos(val bindingContext: BindingContext, val moduleDescriptor: Mo
             ?: calculateForEq(expressionWithType)
             ?: calculateForIf(expressionWithType)
             ?: calculateForElvis(expressionWithType)
+            ?: calculateForBlockExpression(expressionWithType)
             ?: getFromBindingContext(expressionWithType)
     }
 
@@ -208,6 +210,12 @@ class ExpectedInfos(val bindingContext: BindingContext, val moduleDescriptor: Mo
         return null
     }
 
+    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`, null) }
+    }
+
     private fun getFromBindingContext(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
         val expectedType = bindingContext[BindingContext.EXPECTED_EXPRESSION_TYPE, expressionWithType] ?: return null
         return listOf(ExpectedInfo(expectedType, null))
diff --git a/idea/testData/completion/handlers/smart/IfValueInBlock.kt b/idea/testData/completion/handlers/smart/IfValueInBlock.kt
new file mode 100644
index 00000000000..e7e48f81c78
--- /dev/null
+++ b/idea/testData/completion/handlers/smart/IfValueInBlock.kt
@@ -0,0 +1,12 @@
+fun foo(s: String, i: Int){}
+
+fun bar(b: Boolean, s: String){
+    foo(if (b)
+            "abc"
+        else {
+            println()
+            <caret>
+    })
+}
+
+// ELEMENT: s
diff --git a/idea/testData/completion/handlers/smart/IfValueInBlock.kt.after b/idea/testData/completion/handlers/smart/IfValueInBlock.kt.after
new file mode 100644
index 00000000000..2cf8708726f
--- /dev/null
+++ b/idea/testData/completion/handlers/smart/IfValueInBlock.kt.after
@@ -0,0 +1,12 @@
+fun foo(s: String, i: Int){}
+
+fun bar(b: Boolean, s: String){
+    foo(if (b)
+            "abc"
+        else {
+            println()
+            s<caret>
+    })
+}
+
+// ELEMENT: s
diff --git a/idea/testData/completion/smart/IfValueInBlock1.kt b/idea/testData/completion/smart/IfValueInBlock1.kt
new file mode 100644
index 00000000000..825b9b90470
--- /dev/null
+++ b/idea/testData/completion/smart/IfValueInBlock1.kt
@@ -0,0 +1,13 @@
+fun foo(s: String){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String, c: Char){
+    foo(if (b) {
+        println()
+        <caret>
+    })
+}
+
+// EXIST: s
+// EXIST: c
+// ABSENT: b
diff --git a/idea/testData/completion/smart/IfValueInBlock2.kt b/idea/testData/completion/smart/IfValueInBlock2.kt
new file mode 100644
index 00000000000..6cff230967c
--- /dev/null
+++ b/idea/testData/completion/smart/IfValueInBlock2.kt
@@ -0,0 +1,15 @@
+fun foo(s: String){}
+fun foo(c: Char){}
+
+fun bar(b: Boolean, s: String, c: Char){
+    foo(if (b)
+            "abc"
+        else {
+            println()
+            <caret>
+    })
+}
+
+// EXIST: s
+// ABSENT: c
+// ABSENT: b
diff --git a/idea/tests/org/jetbrains/jet/completion/JvmSmartCompletionTestGenerated.java b/idea/tests/org/jetbrains/jet/completion/JvmSmartCompletionTestGenerated.java
index f56e9a104a5..3b6198b418c 100644
--- a/idea/tests/org/jetbrains/jet/completion/JvmSmartCompletionTestGenerated.java
+++ b/idea/tests/org/jetbrains/jet/completion/JvmSmartCompletionTestGenerated.java
@@ -216,6 +216,16 @@ public class JvmSmartCompletionTestGenerated extends AbstractJvmSmartCompletionT
         doTest("idea/testData/completion/smart/IfValue3.kt");
     }
     
+    @TestMetadata("IfValueInBlock1.kt")
+    public void testIfValueInBlock1() throws Exception {
+        doTest("idea/testData/completion/smart/IfValueInBlock1.kt");
+    }
+    
+    @TestMetadata("IfValueInBlock2.kt")
+    public void testIfValueInBlock2() throws Exception {
+        doTest("idea/testData/completion/smart/IfValueInBlock2.kt");
+    }
+    
     @TestMetadata("InElvisOperator1.kt")
     public void testInElvisOperator1() throws Exception {
         doTest("idea/testData/completion/smart/InElvisOperator1.kt");
diff --git a/idea/tests/org/jetbrains/jet/completion/handlers/SmartCompletionHandlerTestGenerated.java b/idea/tests/org/jetbrains/jet/completion/handlers/SmartCompletionHandlerTestGenerated.java
index cd25eac0aa2..db384196ac8 100644
--- a/idea/tests/org/jetbrains/jet/completion/handlers/SmartCompletionHandlerTestGenerated.java
+++ b/idea/tests/org/jetbrains/jet/completion/handlers/SmartCompletionHandlerTestGenerated.java
@@ -226,6 +226,11 @@ public class SmartCompletionHandlerTestGenerated extends AbstractSmartCompletion
         doTest("idea/testData/completion/handlers/smart/IfValue3.kt");
     }
     
+    @TestMetadata("IfValueInBlock.kt")
+    public void testIfValueInBlock() throws Exception {
+        doTest("idea/testData/completion/handlers/smart/IfValueInBlock.kt");
+    }
+    
     @TestMetadata("InElvisOperator.kt")
     public void testInElvisOperator() throws Exception {
         doTest("idea/testData/completion/handlers/smart/InElvisOperator.kt");