Partial body resolve: more precise smart casts analysis for if-statement

idea/ide-common/src/org/jetbrains/jet/lang/resolve/lazy/PartialBodyResolveFilter.kt

@@ -83,15 +83,30 @@ class PartialBodyResolveFilter(elementToResolve: JetElement, private val body: J
     private fun potentialSmartCastPlaces(expression: JetExpression, filter: (String) -> Boolean = { true }): Map<String, List<JetExpression>> {
         val map = HashMap<String, ArrayList<JetExpression>>(0)
 
-        fun addIfCanBeSmartCasted(expression: JetExpression) {
-            val name = expression.smartCastedExpressionName() ?: return
-            if (!filter(name)) return
+        fun addPlace(name: String, place: JetExpression) {
             var list = map[name]
             if (list == null) {
                 list = ArrayList(1)
                 map[name] = list
             }
-            list!!.add(expression)
+            list!!.add(place)
+        }
+
+        fun addPlaces(name: String, places: Collection<JetExpression>) {
+            assert(!places.isEmpty())
+            var list = map[name]
+            if (list == null) {
+                list = ArrayList(places.size)
+                map[name] = list
+            }
+            list!!.addAll(places)
+        }
+
+        fun addIfCanBeSmartCasted(expression: JetExpression) {
+            val name = expression.smartCastedExpressionName() ?: return
+            if (filter(name)) {
+                addPlace(name, expression)
+            }
         }
 
         expression.accept(object : ControlFlowVisitor(){
@@ -121,21 +136,29 @@ class PartialBodyResolveFilter(elementToResolve: JetElement, private val body: J
                     val exits = collectAlwaysExitPoints(thenBranch) + collectAlwaysExitPoints(elseBranch)
                     if (exits.isNotEmpty()) {
                         for (name in smartCastedNames) {
-                            var list = map[name]
-                            if (list == null) {
-                                list = ArrayList(exits.size)
-                                map[name] = list
-                            }
-                            list!!.addAll(exits)
+                            addPlaces(name, exits)
                         }
                     }
                 }
 
                 condition.acceptChildren(this)
+
                 if (thenBranch != null && elseBranch != null) {
-                    //TODO: merge casts!
-                    thenBranch.acceptChildren(this)
-                    elseBranch.acceptChildren(this)
+                    val thenCasts = potentialSmartCastPlaces(thenBranch, filter)
+                    if (!thenCasts.isEmpty()) {
+                        val elseCasts = potentialSmartCastPlaces(elseBranch) { filter(it) && thenCasts.containsKey(it) }
+                        if (!elseCasts.isEmpty()) {
+                            for ((name, places) in thenCasts) {
+                                if (elseCasts.containsKey(name)) { // need filtering by cast names in else-branch
+                                    addPlaces(name, places)
+                                }
+                            }
+
+                            for ((name, places) in elseCasts) { // already filtered by cast names in then-branch
+                                addPlaces(name, places)
+                            }
+                        }
+                    }
                 }
             }
 

idea/testData/resolve/partialBodyResolve/IfBranchesAutoCasts.dump

@@ -0,0 +1,3 @@
+Resolve target: null
+Skipped statements:
+if (x()) { y(p!!) } else { z(p1!!) }

idea/testData/resolve/partialBodyResolve/IfBranchesAutoCasts.kt

@@ -0,0 +1,10 @@
+fun foo(p: Any?, p1, Any?) {
+    if (x()) {
+        y(p!!)
+    }
+    else {
+        z(p1!!)
+    }
+
+    <caret>xxx
+}

idea/testData/resolve/partialBodyResolve/IfBranchesAutoCasts2.dump

@@ -0,0 +1,4 @@
+Resolve target: value-parameter val p1: kotlin.Any? smart-casted to kotlin.Any
+Skipped statements:
+print(p!!)
+print(p2!!)

idea/testData/resolve/partialBodyResolve/IfBranchesAutoCasts2.kt

@@ -0,0 +1,12 @@
+fun foo(p: Any?, p1: Any?, p2: Any?) {
+    if (x()) {
+        print(p!!)
+        print(p1!!)
+    }
+    else {
+        print(p1 as String)
+        print(p2!!)
+    }
+
+    <caret>p1.hashCode()
+}

idea/tests/org/jetbrains/jet/resolve/PartialBodyResolveTestGenerated.java

@@ -19,7 +19,6 @@ package org.jetbrains.jet.resolve;
 import com.intellij.testFramework.TestDataPath;
 import org.jetbrains.jet.JUnit3RunnerWithInners;
 import org.jetbrains.jet.JetTestUtils;
-import org.jetbrains.jet.test.InnerTestClasses;
 import org.jetbrains.jet.test.TestMetadata;
 import org.junit.runner.RunWith;
 
@@ -54,6 +53,18 @@ public class PartialBodyResolveTestGenerated extends AbstractPartialBodyResolveT
         doTest(fileName);
     }
 
+    @TestMetadata("IfBranchesAutoCasts.kt")
+    public void testIfBranchesAutoCasts() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/resolve/partialBodyResolve/IfBranchesAutoCasts.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("IfBranchesAutoCasts2.kt")
+    public void testIfBranchesAutoCasts2() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/resolve/partialBodyResolve/IfBranchesAutoCasts2.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("IfBranchesSmartCast.kt")
     public void testIfBranchesSmartCast() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/testData/resolve/partialBodyResolve/IfBranchesSmartCast.kt");