Remove dependencies on data flow aware expression type in BindingContext

`BindingContext.getType`'s behaviour on argument expression has been different for OI and NI for some time.
New inference used to rewrite argument's type after smartcast, which led to missing subsequent smartcasts in some cases.
This commit makes retrieval of smartcasted argument type explicit.

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

@@ -28,7 +28,7 @@ import org.jetbrains.kotlin.diagnostics.DiagnosticFactory
 import org.jetbrains.kotlin.diagnostics.Errors
 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.quickfix.quickfixUtil.getDataFlowAwareTypes
 import org.jetbrains.kotlin.idea.refactoring.changeSignature.*
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType

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

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.idea.core.KotlinNameSuggester
 import org.jetbrains.kotlin.idea.core.mapArgumentsToParameters
+import org.jetbrains.kotlin.idea.quickfix.quickfixUtil.getDataFlowAwareTypes
 import org.jetbrains.kotlin.idea.refactoring.canRefactor
 import org.jetbrains.kotlin.idea.refactoring.changeSignature.KotlinChangeSignatureConfiguration
 import org.jetbrains.kotlin.idea.refactoring.changeSignature.KotlinMethodDescriptor
@@ -102,8 +103,10 @@ abstract class ChangeFunctionSignatureFix(
                     val call = originalElement.getCall(bindingContext) ?: return null
                     val argumentToParameter = call.mapArgumentsToParameters(functionDescriptor)
                     val hasTypeMismatches = argumentToParameter.any { (argument, parameter) ->
-                        val argumentType = argument.getArgumentExpression()?.let { bindingContext.getType(it) }
+                        val argumentTypes = argument.getArgumentExpression()?.let { getDataFlowAwareTypes(it, bindingContext) }
-                        argumentType == null || !KotlinTypeChecker.DEFAULT.isSubtypeOf(argumentType, parameter.type)
+                        argumentTypes?.none { dataFlowAwareType ->
+                            KotlinTypeChecker.DEFAULT.isSubtypeOf(dataFlowAwareType, parameter.type)
+                        } ?: true
                     }
                     return AddFunctionParametersFix(originalElement, functionDescriptor, hasTypeMismatches)
                 }

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

@@ -9,12 +9,9 @@ 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
+import org.jetbrains.kotlin.idea.quickfix.quickfixUtil.getDataFlowAwareTypes
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
-import org.jetbrains.kotlin.idea.resolve.frontendService
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.name.FqName
@@ -24,9 +21,7 @@ import org.jetbrains.kotlin.psi.psiUtil.getAssignmentByLHS
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfoAfter
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsStatement
-import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.constants.IntegerLiteralTypeConstructor
 import org.jetbrains.kotlin.resolve.descriptorUtil.resolveTopLevelClass
 import org.jetbrains.kotlin.resolve.scopes.HierarchicalScope
@@ -35,7 +30,6 @@ 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 {
@@ -335,18 +329,3 @@ private fun TypePredicate.getRepresentativeTypes(): Set<KotlinType> {
         else -> throw AssertionError("Invalid type predicate: $this")
     }
 }
-
-fun getDataFlowAwareTypes(
-    expression: KtExpression,
-    bindingContext: BindingContext = expression.analyze(),
-    originalType: KotlinType? = bindingContext.getType(expression)
-): Collection<KotlinType> {
-    if (originalType == null) return emptyList()
-    val dataFlowInfo = bindingContext.getDataFlowInfoAfter(expression)
-    val dataFlowValueFactory = expression.getResolutionFacade().frontendService<DataFlowValueFactory>()
-    val expressionType = bindingContext.getType(expression) ?: return listOf(originalType)
-    val dataFlowValue = dataFlowValueFactory.createDataFlowValue(
-        expression, expressionType, bindingContext, expression.getResolutionFacade().moduleDescriptor
-    )
-    return dataFlowInfo.getCollectedTypes(dataFlowValue, expression.languageVersionSettings).ifEmpty { listOf(originalType) }
-}

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createCallable/CreateCallableFromCallActionFactory.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.analyzeAndGetResult
 import org.jetbrains.kotlin.idea.caches.resolve.unsafeResolveToDescriptor
 import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.idea.quickfix.createFromUsage.callableBuilder.*
+import org.jetbrains.kotlin.idea.quickfix.quickfixUtil.getDataFlowAwareTypes
 import org.jetbrains.kotlin.idea.refactoring.canRefactor
 import org.jetbrains.kotlin.idea.refactoring.getExtractionContainers
 import org.jetbrains.kotlin.idea.refactoring.isInterfaceClass
@@ -119,7 +120,11 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
             is ReceiverValue -> {
                 val originalType = receiver.type
                 val finalType = if (receiver is ExpressionReceiver) {
-                    getDataFlowAwareTypes(receiver.expression, context, originalType).firstOrNull() ?: originalType
+                    getDataFlowAwareTypes(
+                        receiver.expression,
+                        context,
+                        originalType
+                    ).firstOrNull() ?: originalType
                 } else originalType
                 TypeInfo(finalType, Variance.IN_VARIANCE)
             }

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

@@ -19,11 +19,21 @@ package org.jetbrains.kotlin.idea.quickfix.quickfixUtil
 import com.intellij.codeInsight.intention.IntentionAction
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.idea.project.languageVersionSettings
 import org.jetbrains.kotlin.idea.quickfix.KotlinQuickFixAction
 import org.jetbrains.kotlin.idea.quickfix.KotlinSingleIntentionActionFactory
+import org.jetbrains.kotlin.idea.resolve.frontendService
+import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtPrimaryConstructor
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfoAfter
+import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.utils.ifEmpty
 
 inline fun <reified T : PsiElement> Diagnostic.createIntentionForFirstParentOfType(
     factory: (T) -> KotlinQuickFixAction<T>?
@@ -40,3 +50,18 @@ fun KtPrimaryConstructor.addConstructorKeyword(): PsiElement? {
     val keyword = KtPsiFactory(this).createConstructorKeyword()
     return addAfter(keyword, modifierList ?: return null)
 }
+
+fun getDataFlowAwareTypes(
+    expression: KtExpression,
+    bindingContext: BindingContext = expression.analyze(),
+    originalType: KotlinType? = bindingContext.getType(expression)
+): Collection<KotlinType> {
+    if (originalType == null) return emptyList()
+    val dataFlowInfo = bindingContext.getDataFlowInfoAfter(expression)
+    val dataFlowValueFactory = expression.getResolutionFacade().frontendService<DataFlowValueFactory>()
+    val expressionType = bindingContext.getType(expression) ?: return listOf(originalType)
+    val dataFlowValue = dataFlowValueFactory.createDataFlowValue(
+        expression, expressionType, bindingContext, expression.getResolutionFacade().moduleDescriptor
+    )
+    return dataFlowInfo.getCollectedTypes(dataFlowValue, expression.languageVersionSettings).ifEmpty { listOf(originalType) }
+}

idea/tests/org/jetbrains/kotlin/idea/resolve/AbstractPartialBodyResolveTest.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.idea.quickfix.quickfixUtil.getDataFlowAwareTypes
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.configureCompilerOptions
@@ -148,7 +149,9 @@ abstract class AbstractPartialBodyResolveTest : KotlinLightCodeInsightFixtureTes
         } else {
             expression
         }
-        val type = bindingContext.getType(expressionWithType)
+
+        val nonDataFlowAwareType = bindingContext.getType(expressionWithType)
+        val type = getDataFlowAwareTypes(expressionWithType, bindingContext, nonDataFlowAwareType).singleOrNull() ?: nonDataFlowAwareType
 
         return ResolveData(target, type, processedStatements)
     }

nj2k/testData/newJ2k/postProcessing/IfToElvis.kt

@@ -1,5 +1,5 @@
 internal class C {
-    fun foo(s: String?): String? {
+    fun foo(s: String?): String {
         return s ?: ""
     }
 }