diff --git a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt
index 5e77e7ce7d1..0761dbabe24 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt
@@ -109,9 +109,7 @@ class KotlinResolutionCallbacksImpl(
         trace.record(BindingContext.NEW_INFERENCE_LAMBDA_INFO, ktFunction, LambdaInfo.STUB_EMPTY)
 
         val lastExpressionArgument = getLastDeparentesizedExpression(psiCallArgument)?.let { lastExpression ->
-            val lastExpressionType = functionTypeInfo.type?.let {
-                if (it.isFunctionType) it.getReturnTypeFromFunctionType() else it
-            }
+            val lastExpressionType = trace.getType(lastExpression)
             val lastExpressionTypeInfo = KotlinTypeInfo(lastExpressionType, lambdaInfo.dataFlowInfoAfter ?: functionTypeInfo.dataFlowInfo)
             createCallArgument(lastExpression, lastExpressionTypeInfo)
         }
diff --git a/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt b/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt
index 8ec2a1ee4f4..e956879ac34 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt
@@ -225,7 +225,10 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
         val blockReturnedType = components.expressionTypingServices.getBlockReturnedType(functionLiteral.bodyExpression!!, COERCION_TO_UNIT, newContext)
         val typeOfBodyExpression = blockReturnedType.type
 
-        context.trace[BindingContext.NEW_INFERENCE_LAMBDA_INFO, expression.functionLiteral]?.dataFlowInfoAfter = blockReturnedType.dataFlowInfo
+        context.trace[BindingContext.NEW_INFERENCE_LAMBDA_INFO, expression.functionLiteral]?.let {
+            it.dataFlowInfoAfter = blockReturnedType.dataFlowInfo
+            return null
+        }
 
         return computeReturnTypeBasedOnReturnExpressions(functionLiteral, context, typeOfBodyExpression)
     }