Add Parameter Fix: Support smart casts

#KT-24207 Fixed

idea/src/org/jetbrains/kotlin/idea/quickfix/AddFunctionParametersFix.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.core.CollectingNameValidator
+import org.jetbrains.kotlin.idea.quickfix.createFromUsage.callableBuilder.getDataFlowAwareTypes
 import org.jetbrains.kotlin.idea.refactoring.changeSignature.*
 import org.jetbrains.kotlin.psi.KtCallElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -141,7 +142,7 @@ class AddFunctionParametersFix(
     ): KotlinParameterInfo {
         val name = getNewArgumentName(argument, validator)
         val expression = argument.getArgumentExpression()
-        val type = expression?.let { it.analyze().getType(it) } ?: functionDescriptor.builtIns.nullableAnyType
+        val type = expression?.let { getDataFlowAwareTypes(it).firstOrNull() } ?: functionDescriptor.builtIns.nullableAnyType
         return KotlinParameterInfo(functionDescriptor, -1, name, KotlinTypeInfo(false, null))
                 .apply { currentTypeInfo = KotlinTypeInfo(false, type) }
     }

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/typeUtils.kt

@@ -20,6 +20,7 @@ import com.intellij.refactoring.psi.SearchUtils
 import org.jetbrains.kotlin.builtins.isFunctionType
 import org.jetbrains.kotlin.cfg.pseudocode.*
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.project.builtIns
 import org.jetbrains.kotlin.idea.project.languageVersionSettings
@@ -44,6 +45,7 @@ import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
 import org.jetbrains.kotlin.types.typeUtil.supertypes
+import org.jetbrains.kotlin.utils.ifEmpty
 import java.util.*
 
 internal operator fun KotlinType.contains(inner: KotlinType): Boolean {
@@ -162,12 +164,7 @@ fun KtExpression.guessTypes(
         // if we know the actual type of the expression
         val theType1 = context.getType(this)
         if (theType1 != null && isAcceptable(theType1)) {
-            val dataFlowInfo = context.getDataFlowInfoAfter(this)
-            val dataFlowValueFactory = this.getResolutionFacade().frontendService<DataFlowValueFactory>()
-            val dataFlowValue = dataFlowValueFactory.createDataFlowValue(this, theType1, context, module)
-
-            val possibleTypes = dataFlowInfo.getCollectedTypes(dataFlowValue, languageVersionSettings)
-            return if (possibleTypes.isNotEmpty()) possibleTypes.toTypedArray() else arrayOf(theType1)
+            return getDataFlowAwareTypes(this, context).toTypedArray()
         }
     }
 
@@ -343,3 +340,15 @@ private fun TypePredicate.getRepresentativeTypes(): Set<KotlinType> {
         else -> throw AssertionError("Invalid type predicate: ${this}")
     }
 }
+
+fun getDataFlowAwareTypes(expression: KtExpression, bindingContext: BindingContext = expression.analyze()): List<KotlinType> {
+    val originalType = bindingContext.getType(expression) ?: return emptyList()
+    val dataFlowInfo = bindingContext.getDataFlowInfoAfter(expression)
+    val dataFlowValueFactory = expression.getResolutionFacade().frontendService<DataFlowValueFactory>()
+    val dataFlowValue = dataFlowValueFactory.createDataFlowValue(
+            expression,
+            bindingContext.getType(expression)!!,
+            bindingContext, expression.getResolutionFacade().moduleDescriptor
+    )
+    return dataFlowInfo.getCollectedTypes(dataFlowValue, expression.languageVersionSettings).ifEmpty { listOf(originalType) }
+}

idea/testData/quickfix/changeSignature/addFunctionParameterWithSmartCast.kt

@@ -0,0 +1,13 @@
+// "Add parameter to function 'fooFun'" "true"
+sealed class Foo {
+    class SubFoo : Foo()
+    class Sub2Foo : Foo()
+}
+
+fun fooFun(): Int = 1
+
+val subFoo: Foo = Foo.SubFoo()
+val bar = when(subFoo){
+    is Foo.SubFoo -> fooFun(<caret>subFoo)
+    else -> 0
+}

idea/testData/quickfix/changeSignature/addFunctionParameterWithSmartCast.kt.after

@@ -0,0 +1,13 @@
+// "Add parameter to function 'fooFun'" "true"
+sealed class Foo {
+    class SubFoo : Foo()
+    class Sub2Foo : Foo()
+}
+
+fun fooFun(subFoo: Foo.SubFoo): Int = 1
+
+val subFoo: Foo = Foo.SubFoo()
+val bar = when(subFoo){
+    is Foo.SubFoo -> fooFun(<caret>subFoo)
+    else -> 0
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -1509,6 +1509,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/changeSignature/addFunctionParameterLongNameRuntime.kt");
         }
 
+        @TestMetadata("addFunctionParameterWithSmartCast.kt")
+        public void testAddFunctionParameterWithSmartCast() throws Exception {
+            runTest("idea/testData/quickfix/changeSignature/addFunctionParameterWithSmartCast.kt");
+        }
+
         @TestMetadata("addNothingReturnType.kt")
         public void testAddNothingReturnType() throws Exception {
             runTest("idea/testData/quickfix/changeSignature/addNothingReturnType.kt");
@@ -3086,8 +3091,7 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
 
                 @TestMetadata("nullableReceiver.kt")
                 public void testNullableReceiver() throws Exception {
-                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createFunction/call/nullableReceiver.kt");
-                    doTest(fileName);
+                    runTest("idea/testData/quickfix/createFromUsage/createFunction/call/nullableReceiver.kt");
                 }
 
                 @TestMetadata("objectMemberFunNoReceiver.kt")
@@ -4607,8 +4611,7 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
 
                 @TestMetadata("nullableReceiver.kt")
                 public void testNullableReceiver() throws Exception {
-                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createVariable/property/nullableReceiver.kt");
-                    doTest(fileName);
+                    runTest("idea/testData/quickfix/createFromUsage/createVariable/property/nullableReceiver.kt");
                 }
 
                 @TestMetadata("objectMemberValNoReceiver.kt")