[NI] Add common supertype for PSI lambda arguments

Also move initial dataFlowInfo to lambda arguments.
As result, we can not store outer call for postpone call arguments

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

@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.psi.KtPsiUtil
 import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
@@ -69,31 +68,23 @@ class KotlinResolutionCallbacksImpl(
     }
 
     override fun analyzeAndGetLambdaResultArguments(
-            outerCall: KotlinCall,
             lambdaArgument: LambdaKotlinCallArgument,
             isSuspend: Boolean,
             receiverType: UnwrappedType?,
             parameters: List<UnwrappedType>,
             expectedReturnType: UnwrappedType?
     ): List<SimpleKotlinCallArgument> {
-        val psiCallArgument = lambdaArgument.psiCallArgument
+        val psiCallArgument = lambdaArgument.psiCallArgument as PSIFunctionKotlinCallArgument
-        val outerCallContext = (psiCallArgument as? LambdaKotlinCallArgumentImpl)?.outerCallContext ?:
+        val outerCallContext = psiCallArgument.outerCallContext
-                               (psiCallArgument as FunctionExpressionImpl).outerCallContext
 
         fun createCallArgument(ktExpression: KtExpression, typeInfo: KotlinTypeInfo) =
                 createSimplePSICallArgument(trace.bindingContext, outerCallContext.statementFilter, outerCallContext.scope.ownerDescriptor,
                                             CallMaker.makeExternalValueArgument(ktExpression), DataFlowInfo.EMPTY, typeInfo)
 
-        val expression: KtExpression = (psiCallArgument as? LambdaKotlinCallArgumentImpl)?.ktLambdaExpression ?:
-                               (psiCallArgument as FunctionExpressionImpl).ktFunction
-
-        val ktFunction: KtFunction = (psiCallArgument as? LambdaKotlinCallArgumentImpl)?.ktLambdaExpression?.functionLiteral ?:
-                         (psiCallArgument as FunctionExpressionImpl).ktFunction
-
         val lambdaInfo = LambdaInfo(expectedReturnType ?: TypeUtils.NO_EXPECTED_TYPE,
                                     if (expectedReturnType == null) ContextDependency.DEPENDENT else ContextDependency.INDEPENDENT)
 
-        trace.record(BindingContext.NEW_INFERENCE_LAMBDA_INFO, ktFunction, lambdaInfo)
+        trace.record(BindingContext.NEW_INFERENCE_LAMBDA_INFO, psiCallArgument.ktFunction, lambdaInfo)
 
         val builtIns = outerCallContext.scope.ownerDescriptor.builtIns
         val expectedType = createFunctionType(builtIns, Annotations.EMPTY, receiverType, parameters, null,
@@ -105,10 +96,10 @@ class KotlinResolutionCallbacksImpl(
                 .replaceBindingTrace(trace)
                 .replaceContextDependency(lambdaInfo.contextDependency)
                 .replaceExpectedType(approximatesExpectedType)
-                .replaceDataFlowInfo(outerCall.psiKotlinCall.resultDataFlowInfo)
+                .replaceDataFlowInfo(psiCallArgument.lambdaInitialDataFlowInfo)
 
-        val functionTypeInfo = expressionTypingServices.getTypeInfo(expression, actualContext)
+        val functionTypeInfo = expressionTypingServices.getTypeInfo(psiCallArgument.expression, actualContext)
-        trace.record(BindingContext.NEW_INFERENCE_LAMBDA_INFO, ktFunction, LambdaInfo.STUB_EMPTY)
+        trace.record(BindingContext.NEW_INFERENCE_LAMBDA_INFO, psiCallArgument.ktFunction, LambdaInfo.STUB_EMPTY)
 
         var hasReturnWithoutExpression = false
         val returnArguments = lambdaInfo.returnStatements.mapNotNullTo(ArrayList()) { (expression, typeInfo) ->

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

@@ -81,30 +81,43 @@ class ParseErrorKotlinCallArgument(
         get() = dataFlowInfoAfterThisArgument
 }
 
-class LambdaKotlinCallArgumentImpl(
+abstract class PSIFunctionKotlinCallArgument(
         val outerCallContext: BasicCallResolutionContext,
         override val valueArgument: ValueArgument,
         override val dataFlowInfoBeforeThisArgument: DataFlowInfo,
-        val ktLambdaExpression: KtLambdaExpression,
+        override val argumentName: Name?
-        override val argumentName: Name?,
-        override val parametersTypes: Array<UnwrappedType?>?
 ) : LambdaKotlinCallArgument, PSIKotlinCallArgument() {
-    override val dataFlowInfoAfterThisArgument: DataFlowInfo
+    override val dataFlowInfoAfterThisArgument: DataFlowInfo // todo drop this and use only lambdaInitialDataFlowInfo
         get() = dataFlowInfoBeforeThisArgument
+
+    abstract val ktFunction: KtFunction
+    abstract val expression: KtExpression
+    lateinit var lambdaInitialDataFlowInfo: DataFlowInfo
+}
+
+class LambdaKotlinCallArgumentImpl(
+        outerCallContext: BasicCallResolutionContext,
+        valueArgument: ValueArgument,
+        dataFlowInfoBeforeThisArgument: DataFlowInfo,
+        argumentName: Name?,
+        val ktLambdaExpression: KtLambdaExpression,
+        override val parametersTypes: Array<UnwrappedType?>?
+) : PSIFunctionKotlinCallArgument(outerCallContext, valueArgument, dataFlowInfoBeforeThisArgument, argumentName) {
+    override val ktFunction get() = ktLambdaExpression.functionLiteral
+    override val expression get() = ktLambdaExpression
 }
 
 class FunctionExpressionImpl(
-        val outerCallContext: BasicCallResolutionContext,
+        outerCallContext: BasicCallResolutionContext,
-        override val valueArgument: ValueArgument,
+        valueArgument: ValueArgument,
-        override val dataFlowInfoBeforeThisArgument: DataFlowInfo,
+        dataFlowInfoBeforeThisArgument: DataFlowInfo,
-        val ktFunction: KtNamedFunction,
+        argumentName: Name?,
-        override val argumentName: Name?,
+        override val ktFunction: KtNamedFunction,
         override val receiverType: UnwrappedType?,
         override val parametersTypes: Array<UnwrappedType?>,
         override val returnType: UnwrappedType?
-) : FunctionExpression, PSIKotlinCallArgument() {
+) : FunctionExpression, PSIFunctionKotlinCallArgument(outerCallContext, valueArgument, dataFlowInfoBeforeThisArgument, argumentName) {
-    override val dataFlowInfoAfterThisArgument: DataFlowInfo
+    override val expression get() = ktFunction
-        get() = dataFlowInfoBeforeThisArgument
 }
 
 class CallableReferenceKotlinCallArgumentImpl(
@@ -171,6 +184,13 @@ class EmptyLabeledReturn(
     override val isSafeCall: Boolean get() = false
 }
 
+internal fun KotlinCallArgument.setResultDataFlowInfoIfRelevant(resultDataFlowInfo: DataFlowInfo) {
+    if (this is PSIFunctionKotlinCallArgument) {
+        lambdaInitialDataFlowInfo = resultDataFlowInfo
+    }
+}
+
+
 // context here is context for value argument analysis
 internal fun createSimplePSICallArgument(
         contextForArgument: BasicCallResolutionContext,

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

@@ -423,6 +423,9 @@ class PSICallResolver(
         val astExternalArgument = externalArgument?.let { resolveValueArgument(context, dataFlowInfoAfterArgumentsInParenthesis, it) }
         val resultDataFlowInfo = astExternalArgument?.dataFlowInfoAfterThisArgument ?: dataFlowInfoAfterArgumentsInParenthesis
 
+        resolvedArgumentsInParenthesis.forEach { it.setResultDataFlowInfoIfRelevant(resultDataFlowInfo) }
+        astExternalArgument?.setResultDataFlowInfoIfRelevant(resultDataFlowInfo)
+
         return PSIKotlinCallImpl(kotlinCallKind, oldCall, tracingStrategy, resolvedExplicitReceiver, name, resolvedTypeArguments, resolvedArgumentsInParenthesis,
                                  astExternalArgument, context.dataFlowInfo, resultDataFlowInfo, context.dataFlowInfoForArguments)
     }
@@ -501,14 +504,14 @@ class PSICallResolver(
 
         val lambdaArgument: PSIKotlinCallArgument? = when (ktExpression) {
             is KtLambdaExpression ->
-                LambdaKotlinCallArgumentImpl(outerCallContext, valueArgument, startDataFlowInfo, ktExpression, argumentName,
+                LambdaKotlinCallArgumentImpl(outerCallContext, valueArgument, startDataFlowInfo, argumentName, ktExpression,
                                              resolveParametersTypes(outerCallContext, ktExpression.functionLiteral))
             is KtNamedFunction -> {
                 val receiverType = resolveType(outerCallContext, ktExpression.receiverTypeReference)
                 val parametersTypes = resolveParametersTypes(outerCallContext, ktExpression) ?: emptyArray()
                 val returnType = resolveType(outerCallContext, ktExpression.typeReference) ?:
                                  if (ktExpression.hasBlockBody()) builtIns.unitType else null
-                FunctionExpressionImpl(outerCallContext, valueArgument, startDataFlowInfo, ktExpression, argumentName, receiverType, parametersTypes, returnType)
+                FunctionExpressionImpl(outerCallContext, valueArgument, startDataFlowInfo, argumentName, ktExpression, receiverType, parametersTypes, returnType)
             }
             else -> null
         }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/CallableReferenceResolution.kt

@@ -156,7 +156,6 @@ class CallableReferencesCandidateFactory(
         val compatibilityChecker: ((ConstraintSystemOperation) -> Unit) -> Unit,
         val expectedType: UnwrappedType?
 ) : CandidateFactory<CallableReferenceCandidate> {
-    private val position = ArgumentConstraintPosition(argument)
 
     fun createCallableProcessor(explicitReceiver: DetailedReceiver?) =
             createCallableReferenceProcessor(outerCallContext.scopeTower, argument.rhsName, this, explicitReceiver)

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ExternalComponents.kt

@@ -34,7 +34,6 @@ interface KotlinResolutionExternalPredicates {
 // This components hold state (trace). Work with this carefully.
 interface KotlinResolutionCallbacks {
     fun analyzeAndGetLambdaResultArguments(
-            outerCall: KotlinCall,
             lambdaArgument: LambdaKotlinCallArgument,
             isSuspend: Boolean,
             receiverType: UnwrappedType?,

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/KotlinCallCompleter.kt

@@ -199,7 +199,7 @@ class KotlinCallCompleter(
         val parameters = lambda.parameters.map(::substitute)
         val expectedType = lambda.returnType.takeIf { c.canBeProper(it) }?.let(::substitute)
         lambda.analyzed = true
-        lambda.resultArguments = resolutionCallbacks.analyzeAndGetLambdaResultArguments(lambda.outerCall, lambda.argument, lambda.isSuspend, receiver, parameters, expectedType)
+        lambda.resultArguments = resolutionCallbacks.analyzeAndGetLambdaResultArguments(lambda.argument, lambda.isSuspend, receiver, parameters, expectedType)
 
         for (resultLambdaArgument in lambda.resultArguments) {
             checkSimpleArgument(c.getBuilder(), resultLambdaArgument, lambda.returnType.let(::substitute))

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponeArgumentsChecks.kt

@@ -17,7 +17,6 @@
 package org.jetbrains.kotlin.resolve.calls.components
 
 import org.jetbrains.kotlin.builtins.*
-import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.model.ArgumentConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.inference.model.TypeVariableForLambdaReturnType
@@ -26,14 +25,12 @@ import org.jetbrains.kotlin.types.UnwrappedType
 import org.jetbrains.kotlin.types.typeUtil.builtIns
 
 fun createPostponedArgumentAndPerformInitialChecks(
-        kotlinCall: KotlinCall,
         csBuilder: ConstraintSystemBuilder,
         argument: PostponableKotlinCallArgument,
-        parameterDescriptor: ValueParameterDescriptor
+        expectedType: UnwrappedType
 ): KotlinCallDiagnostic? {
-    val expectedType = argument.getExpectedType(parameterDescriptor)
     val (postponedArgument, diagnostic) =  when (argument) {
-        is LambdaKotlinCallArgument -> preprocessLambdaArgument(kotlinCall, csBuilder, argument, expectedType)
+        is LambdaKotlinCallArgument -> preprocessLambdaArgument(csBuilder, argument, expectedType)
         is CallableReferenceKotlinCallArgument -> preprocessCallableReference(csBuilder, argument, expectedType)
         is CollectionLiteralKotlinCallArgument -> preprocessCollectionLiteralArgument(csBuilder, argument, expectedType)
         else -> unexpectedArgument(argument)
@@ -45,7 +42,6 @@ fun createPostponedArgumentAndPerformInitialChecks(
 
 // if expected type isn't function type, then may be it is Function<R>, Any or just `T`
 private fun preprocessLambdaArgument(
-        kotlinCall: KotlinCall,
         csBuilder: ConstraintSystemBuilder,
         argument: LambdaKotlinCallArgument,
         expectedType: UnwrappedType
@@ -84,7 +80,7 @@ private fun preprocessLambdaArgument(
         // what about case where expected type is type variable? In old TY such cases was not supported. => do nothing for now. todo design
     }
 
-    val resolvedArgument = PostponedLambdaArgument(kotlinCall, argument, isSuspend, receiverType, parameters, returnType)
+    val resolvedArgument = PostponedLambdaArgument(argument, isSuspend, receiverType, parameters, returnType)
 
     csBuilder.addSubtypeConstraint(resolvedArgument.type, expectedType, ArgumentConstraintPosition(argument))
 

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt

@@ -234,11 +234,12 @@ internal object CheckArguments : ResolutionPart {
             val resolvedCallArgument = argumentMappingByOriginal[parameterDescriptor.original] ?: continue
             for (argument in resolvedCallArgument.arguments) {
 
+                val expectedType = argument.getExpectedType(parameterDescriptor)
                 val diagnostic = when (argument) {
                     is SimpleKotlinCallArgument ->
-                        checkSimpleArgument(csBuilder, argument, argument.getExpectedType(parameterDescriptor))
+                        checkSimpleArgument(csBuilder, argument, expectedType)
                     is PostponableKotlinCallArgument ->
-                        createPostponedArgumentAndPerformInitialChecks(kotlinCall, csBuilder, argument, parameterDescriptor)
+                        createPostponedArgumentAndPerformInitialChecks(csBuilder, argument, expectedType)
                     else -> unexpectedArgument(argument)
                 }
                 diagnostics.addIfNotNull(diagnostic)

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/PostponedKotlinCallArguments.kt

@@ -30,7 +30,6 @@ sealed class PostponedKotlinCallArgument {
 }
 
 class PostponedLambdaArgument(
-        val outerCall: KotlinCall,
         override val argument: LambdaKotlinCallArgument,
         val isSuspend: Boolean,
         val receiver: UnwrappedType?,