[NI] Fix data flow for arguments in special calls

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinToResolvedCallTransformer.kt

@@ -409,13 +409,7 @@ sealed class NewAbstractResolvedCall<D : CallableDescriptor>(): ResolvedCall<D>
             if (externalPsiCallArgument?.valueArgument == valueArgument) {
                 return externalPsiCallArgument.dataFlowInfoAfterThisArgument
             }
-            kotlinCall.argumentsInParenthesis.find { it.psiCallArgument.valueArgument == valueArgument }?.let {
-                return it.psiCallArgument.dataFlowInfoAfterThisArgument
-            }
-
-            // valueArgument is not found
-            // may be we should return initial DataFlowInfo but I think that it isn't important
-            return kotlinCall.psiKotlinCall.resultDataFlowInfo
+            return kotlinCall.psiKotlinCall.dataFlowInfoForArguments.getInfo(valueArgument)
         }
     }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/PSICallResolver.kt

@@ -60,7 +60,6 @@ import org.jetbrains.kotlin.resolve.scopes.SyntheticScopes
 import org.jetbrains.kotlin.resolve.scopes.receivers.*
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.expressions.*
-import org.jetbrains.kotlin.types.expressions.ControlStructureTypingUtils.ControlStructureDataFlowInfo
 import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 import java.util.*
 
@@ -385,9 +384,6 @@ class PSICallResolver(
         val resolvedExplicitReceiver = resolveExplicitReceiver(context, forcedExplicitReceiver?: oldCall.explicitReceiver, oldCall.isSafeCall())
         val resolvedTypeArguments = resolveTypeArguments(context, oldCall.typeArguments)
 
-        // this is hack for special calls. Note that special call has only arguments in parenthesis.
-        val givenDataFlowInfo: ControlStructureDataFlowInfo? = context.dataFlowInfoForArguments as? ControlStructureDataFlowInfo
-
         val argumentsInParenthesis = if (oldCall.callType != Call.CallType.ARRAY_SET_METHOD && oldCall.functionLiteralArguments.isEmpty()) {
             oldCall.valueArguments
         }
@@ -395,10 +391,9 @@ class PSICallResolver(
             oldCall.valueArguments.dropLast(1)
         }
 
-        val (resolvedArgumentsInParenthesis, dataFlowInfoAfterArgumentsInParenthesis) = resolveArgumentsInParenthesis(
-                context, context.dataFlowInfoForArguments.resultInfo, argumentsInParenthesis, givenDataFlowInfo)
-
         val externalLambdaArguments = oldCall.functionLiteralArguments
+        val resolvedArgumentsInParenthesis = resolveArgumentsInParenthesis(context, argumentsInParenthesis)
+
         val externalArgument = if (oldCall.callType == Call.CallType.ARRAY_SET_METHOD) {
             assert(externalLambdaArguments.isEmpty()) {
                 "Unexpected lambda parameters for call $oldCall"
@@ -415,11 +410,17 @@ class PSICallResolver(
             externalLambdaArguments.firstOrNull()
         }
 
+        val dataFlowInfoAfterArgumentsInParenthesis =
+                if (externalArgument != null && resolvedArgumentsInParenthesis.isNotEmpty())
+                    resolvedArgumentsInParenthesis.last().psiCallArgument.dataFlowInfoAfterThisArgument
+                else
+                    context.dataFlowInfoForArguments.resultInfo
+
         val astExternalArgument = externalArgument?.let { resolveValueArgument(context, dataFlowInfoAfterArgumentsInParenthesis, it) }
         val resultDataFlowInfo = astExternalArgument?.dataFlowInfoAfterThisArgument ?: dataFlowInfoAfterArgumentsInParenthesis
 
         return PSIKotlinCallImpl(kotlinCallKind, oldCall, tracingStrategy, resolvedExplicitReceiver, name, resolvedTypeArguments, resolvedArgumentsInParenthesis,
-                                 astExternalArgument, context.dataFlowInfo, resultDataFlowInfo)
+                                 astExternalArgument, context.dataFlowInfo, resultDataFlowInfo, context.dataFlowInfoForArguments)
     }
 
     private fun resolveExplicitReceiver(context: BasicCallResolutionContext, oldReceiver: Receiver?, isSafeCall: Boolean): ReceiverKotlinCallArgument? =
@@ -472,26 +473,14 @@ class PSICallResolver(
 
     private fun resolveArgumentsInParenthesis(
             context: BasicCallResolutionContext,
-            dataFlowInfoForArguments: DataFlowInfo,
-            arguments: List<ValueArgument>,
-            givenDataFlowInfo: ControlStructureDataFlowInfo?
-    ): Pair<List<KotlinCallArgument>, DataFlowInfo> {
-        if (givenDataFlowInfo != null) {
-            val resolvedArguments = arguments.map {
-                resolveValueArgument(context, givenDataFlowInfo.getInfo(it), it)
+            arguments: List<ValueArgument>
+    ): List<KotlinCallArgument> {
+        val dataFlowInfoForArguments = context.dataFlowInfoForArguments
+        return arguments.map { argument ->
+            resolveValueArgument(context, dataFlowInfoForArguments.getInfo(argument), argument).also { resolvedArgument ->
+                dataFlowInfoForArguments.updateInfo(argument, resolvedArgument.dataFlowInfoAfterThisArgument)
             }
-            return resolvedArguments to givenDataFlowInfo.resultInfo
         }
-
-        var dataFlowInfo = dataFlowInfoForArguments
-
-        val resolvedArguments = arguments.map {
-            val argument = resolveValueArgument(context, dataFlowInfo, it)
-            dataFlowInfo = argument.dataFlowInfoAfterThisArgument
-            argument
-        }
-
-        return resolvedArguments to dataFlowInfo
     }
 
     private fun resolveValueArgument(

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/PSIKotlinCalls.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2016 JetBrains s.r.o.
+ * Copyright 2010-2017 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -37,6 +37,7 @@ abstract class PSIKotlinCall : KotlinCall {
     abstract val psiCall: Call
     abstract val startingDataFlowInfo: DataFlowInfo
     abstract val resultDataFlowInfo: DataFlowInfo
+    abstract val dataFlowInfoForArguments: DataFlowInfoForArguments
     abstract val tracingStrategy: TracingStrategy
 
     override fun toString() = "$psiCall"
@@ -52,7 +53,8 @@ class PSIKotlinCallImpl(
         override val argumentsInParenthesis: List<KotlinCallArgument>,
         override val externalArgument: KotlinCallArgument?,
         override val startingDataFlowInfo: DataFlowInfo,
-        override val resultDataFlowInfo: DataFlowInfo
+        override val resultDataFlowInfo: DataFlowInfo,
+        override val dataFlowInfoForArguments: DataFlowInfoForArguments
 ) : PSIKotlinCall()
 
 class PSIKotlinCallForVariable(
@@ -67,6 +69,7 @@ class PSIKotlinCallForVariable(
 
     override val startingDataFlowInfo: DataFlowInfo get() = baseCall.startingDataFlowInfo
     override val resultDataFlowInfo: DataFlowInfo get() = baseCall.startingDataFlowInfo
+    override val dataFlowInfoForArguments: DataFlowInfoForArguments get() = baseCall.dataFlowInfoForArguments
 
     override val tracingStrategy: TracingStrategy get() = baseCall.tracingStrategy
     override val psiCall: Call = CallTransformer.stripCallArguments(baseCall.psiCall).let {
@@ -87,6 +90,7 @@ class PSIKotlinCallForInvoke(
 
     override val startingDataFlowInfo: DataFlowInfo get() = baseCall.startingDataFlowInfo
     override val resultDataFlowInfo: DataFlowInfo get() = baseCall.resultDataFlowInfo
+    override val dataFlowInfoForArguments: DataFlowInfoForArguments get() = baseCall.dataFlowInfoForArguments
     override val psiCall: Call
     override val tracingStrategy: TracingStrategy