[NI] Don't forget to analyze whole block for last postponed expression

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

@@ -106,10 +106,11 @@ class LambdaKotlinCallArgumentImpl(
         dataFlowInfoBeforeThisArgument: DataFlowInfo,
         argumentName: Name?,
         val ktLambdaExpression: KtLambdaExpression,
+        val containingBlockForLambda: KtExpression,
         override val parametersTypes: Array<UnwrappedType?>?
 ) : PSIFunctionKotlinCallArgument(outerCallContext, valueArgument, dataFlowInfoBeforeThisArgument, argumentName) {
     override val ktFunction get() = ktLambdaExpression.functionLiteral
-    override val expression get() = ktLambdaExpression
+    override val expression get() = containingBlockForLambda
 }
 
 class FunctionExpressionImpl(
@@ -117,12 +118,13 @@ class FunctionExpressionImpl(
         valueArgument: ValueArgument,
         dataFlowInfoBeforeThisArgument: DataFlowInfo,
         argumentName: Name?,
+        val containingBlockForFunction: KtExpression,
         override val ktFunction: KtNamedFunction,
         override val receiverType: UnwrappedType?,
         override val parametersTypes: Array<UnwrappedType?>,
         override val returnType: UnwrappedType?
 ) : FunctionExpression, PSIFunctionKotlinCallArgument(outerCallContext, valueArgument, dataFlowInfoBeforeThisArgument, argumentName) {
-    override val expression get() = ktFunction
+    override val expression get() = containingBlockForFunction
 }
 
 class CallableReferenceKotlinCallArgumentImpl(

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

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.resolve.calls.tower
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.contracts.EffectSystem
@@ -521,27 +522,14 @@ class PSICallResolver(
     ): PSIKotlinCallArgument {
         val builtIns = outerCallContext.scope.ownerDescriptor.builtIns
         val parseErrorArgument = ParseErrorKotlinCallArgument(valueArgument, startDataFlowInfo, builtIns)
-        val ktExpression = extractArgumentExpression(outerCallContext, valueArgument) ?: parseErrorArgument
+        val argumentExpression = valueArgument.getArgumentExpression() ?: return parseErrorArgument
+
+        val ktExpression = KtPsiUtil.deparenthesize(argumentExpression) ?: parseErrorArgument
 
         val argumentName = valueArgument.getArgumentName()?.asName
 
-        val lambdaArgument: PSIKotlinCallArgument? = when (ktExpression) {
-            is KtLambdaExpression ->
-                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, argumentName, ktExpression, receiverType, parametersTypes, returnType)
-            }
-
-            else -> null
-        }
-        if (lambdaArgument != null) {
-            checkNoSpread(outerCallContext, valueArgument)
-            return lambdaArgument
+        processFunctionalExpression(outerCallContext, argumentExpression, startDataFlowInfo, valueArgument, argumentName, builtIns)?.let {
+            return it
         }
 
         if (ktExpression is KtCollectionLiteralExpression) {
@@ -596,24 +584,43 @@ class PSICallResolver(
                                                            ktExpression, argumentName, lhsNewResult, name)
         }
 
-        val argumentExpression = valueArgument.getArgumentExpression() ?: return parseErrorArgument
-
         // argumentExpression instead of ktExpression is hack -- type info should be stored also for parenthesized expression
         val typeInfo = expressionTypingServices.getTypeInfo(argumentExpression, context)
         return createSimplePSICallArgument(context, valueArgument, typeInfo) ?: parseErrorArgument
     }
 
-    private fun extractArgumentExpression(outerCallContext: BasicCallResolutionContext, valueArgument: ValueArgument): KtExpression? {
-        val argumentExpression = valueArgument.getArgumentExpression() ?: return null
+    private fun processFunctionalExpression(
+            outerCallContext: BasicCallResolutionContext,
+            argumentExpression: KtExpression,
+            startDataFlowInfo: DataFlowInfo,
+            valueArgument: ValueArgument,
+            argumentName: Name?,
+            builtIns: KotlinBuiltIns
+    ): PSIKotlinCallArgument? {
+        val expression = ArgumentTypeResolver.getFunctionLiteralArgumentIfAny(argumentExpression, outerCallContext) ?: return null
+        val postponedExpression = if (expression is KtFunctionLiteral) expression.getParentOfType<KtLambdaExpression>(true) else expression
 
-        val ktExpression = ArgumentTypeResolver.getFunctionLiteralArgumentIfAny(argumentExpression, outerCallContext) ?:
-                           ArgumentTypeResolver.getCallableReferenceExpressionIfAny(argumentExpression, outerCallContext) ?:
-                           KtPsiUtil.deparenthesize(argumentExpression)
+        val lambdaArgument: PSIKotlinCallArgument? = when (postponedExpression) {
+            is KtLambdaExpression ->
+                LambdaKotlinCallArgumentImpl(outerCallContext, valueArgument, startDataFlowInfo, argumentName, postponedExpression,
+                                             argumentExpression, resolveParametersTypes(outerCallContext, postponedExpression.functionLiteral))
 
-        return when (ktExpression) {
-            is KtFunctionLiteral -> ktExpression.getParentOfType<KtLambdaExpression>(true)
-            else -> ktExpression
+            is KtNamedFunction -> {
+                val receiverType = resolveType(outerCallContext, postponedExpression.receiverTypeReference)
+                val parametersTypes = resolveParametersTypes(outerCallContext, postponedExpression) ?: emptyArray()
+                val returnType = resolveType(outerCallContext, postponedExpression.typeReference) ?:
+                                 if (postponedExpression.hasBlockBody()) builtIns.unitType else null
+
+                FunctionExpressionImpl(outerCallContext, valueArgument, startDataFlowInfo, argumentName,
+                                       argumentExpression, postponedExpression, receiverType, parametersTypes, returnType)
+            }
+
+            else -> return null
         }
+
+        checkNoSpread(outerCallContext, valueArgument)
+
+        return lambdaArgument
     }
 
     private fun checkNoSpread(context: BasicCallResolutionContext, valueArgument: ValueArgument) {
@@ -629,6 +636,4 @@ class PSICallResolver(
             parameterList.parameters[it]?.typeReference?.let { resolveType(context, it) }
         }
     }
-
-
-}
+}