[NI] Change lambda analysis -- create arguments for return statements.

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContext.java

@@ -36,6 +36,7 @@ import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.ExplicitSmartCasts;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.ImplicitSmartCasts;
+import org.jetbrains.kotlin.resolve.calls.tower.KotlinResolutionCallbacksImpl;
 import org.jetbrains.kotlin.resolve.constants.CompileTimeConstant;
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics;
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope;
@@ -257,6 +258,8 @@ public interface BindingContext {
     WritableSlice<KtFile, PackageFragmentDescriptor> FILE_TO_PACKAGE_FRAGMENT = Slices.createSimpleSlice();
     WritableSlice<FqName, Collection<KtFile>> PACKAGE_TO_FILES = Slices.createSimpleSlice();
 
+    WritableSlice<KtFunction, KotlinResolutionCallbacksImpl.LambdaInfo> NEW_INFERENCE_LAMBDA_INFO = new BasicWritableSlice<>(DO_NOTHING);
+
     @SuppressWarnings("UnusedDeclaration")
     @Deprecated // This field is needed only for the side effects of its initializer
             Void _static_initializer = BasicWritableSlice.initSliceDebugNames(BindingContext.class);

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

@@ -25,7 +25,9 @@ import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.descriptors.impl.FunctionDescriptorImpl
 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
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
@@ -45,6 +47,7 @@ import org.jetbrains.kotlin.types.expressions.DoubleColonExpressionResolver
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import org.jetbrains.kotlin.types.expressions.KotlinTypeInfo
 import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
+import org.jetbrains.kotlin.utils.addIfNotNull
 
 class KotlinResolutionCallbacksImpl(
         val topLevelCallContext: BasicCallResolutionContext,
@@ -56,6 +59,15 @@ class KotlinResolutionCallbacksImpl(
 ): KotlinResolutionCallbacks {
     val trace: BindingTrace = topLevelCallContext.trace
 
+    class LambdaInfo(val expectedType: UnwrappedType, val contextDependency: ContextDependency) {
+        var dataFlowInfoAfter: DataFlowInfo? = null
+        val returnStatements = ArrayList<Pair<KtReturnExpression, KotlinTypeInfo>>()
+
+        companion object {
+            val STUB_EMPTY = LambdaInfo(TypeUtils.NO_EXPECTED_TYPE, ContextDependency.INDEPENDENT)
+        }
+    }
+
     override fun analyzeAndGetLambdaResultArguments(
             topLevelCall: KotlinCall,
             lambdaArgument: LambdaKotlinCallArgument,
@@ -67,25 +79,54 @@ class KotlinResolutionCallbacksImpl(
         val psiCallArgument = lambdaArgument.psiCallArgument
         val outerCallContext = (psiCallArgument as? LambdaKotlinCallArgumentImpl)?.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)
+
         val builtIns = outerCallContext.scope.ownerDescriptor.builtIns
         val expectedType = createFunctionType(builtIns, Annotations.EMPTY, receiverType, parameters, null,
-                           expectedReturnType ?: TypeUtils.NO_EXPECTED_TYPE, isSuspend)
+                           lambdaInfo.expectedType, isSuspend)
 
         val approximatesExpectedType = typeApproximator.approximateToSubType(expectedType, TypeApproximatorConfiguration.LocalDeclaration) ?: expectedType
 
         val actualContext = outerCallContext.replaceBindingTrace(trace).
-                replaceContextDependency(if (expectedReturnType == null) ContextDependency.DEPENDENT else ContextDependency.INDEPENDENT).replaceExpectedType(approximatesExpectedType)
+                replaceContextDependency(lambdaInfo.contextDependency).replaceExpectedType(approximatesExpectedType)
 
 
         val functionTypeInfo = expressionTypingServices.getTypeInfo(expression, actualContext)
-        val lastExpressionType = functionTypeInfo.type?.let {
+        trace.record(BindingContext.NEW_INFERENCE_LAMBDA_INFO, ktFunction, LambdaInfo.STUB_EMPTY)
-            if (it.isFunctionType) it.getReturnTypeFromFunctionType() else it
-        }
-        val lastExpressionTypeInfo = KotlinTypeInfo(lastExpressionType, functionTypeInfo.dataFlowInfo)
 
+        val lastExpressionArgument = getLastDeparentesizedExpression(psiCallArgument)?.let { lastExpression ->
+            val lastExpressionType = functionTypeInfo.type?.let {
+                if (it.isFunctionType) it.getReturnTypeFromFunctionType() else it
+            }
+            val lastExpressionTypeInfo = KotlinTypeInfo(lastExpressionType, lambdaInfo.dataFlowInfoAfter ?: functionTypeInfo.dataFlowInfo)
+            createCallArgument(lastExpression, lastExpressionTypeInfo)
+        }
+
+        val returnArguments = lambdaInfo.returnStatements.mapNotNullTo(ArrayList()) { (expression, typeInfo) ->
+            expression.returnedExpression?.let { createCallArgument(it, typeInfo) }
+        }
+
+        returnArguments.addIfNotNull(lastExpressionArgument)
+
+        return returnArguments
+    }
+
+    private fun getLastDeparentesizedExpression(psiCallArgument: PSIKotlinCallArgument): KtExpression? {
         val lastExpression: KtExpression?
         if (psiCallArgument is LambdaKotlinCallArgumentImpl) {
             lastExpression = psiCallArgument.ktLambdaExpression.bodyExpression?.statements?.lastOrNull()
@@ -94,11 +135,7 @@ class KotlinResolutionCallbacksImpl(
             lastExpression = (psiCallArgument as FunctionExpressionImpl).ktFunction.bodyExpression?.lastBlockStatementOrThis()
         }
 
-        val deparentesized = KtPsiUtil.deparenthesize(lastExpression) ?: return emptyList()
+        return KtPsiUtil.deparenthesize(lastExpression)
-
-        val simpleArgument = createSimplePSICallArgument(actualContext, CallMaker.makeExternalValueArgument(deparentesized), lastExpressionTypeInfo)
-
-        return listOfNotNull(simpleArgument)
     }
 
     override fun bindStubResolvedCallForCandidate(candidate: KotlinResolutionCandidate) {

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

@@ -18,9 +18,11 @@ package org.jetbrains.kotlin.resolve.calls.tower
 
 import com.intellij.psi.impl.source.tree.LeafPsiElement
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.StatementFilter
 import org.jetbrains.kotlin.resolve.calls.callUtil.getCall
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.model.*
@@ -147,30 +149,45 @@ class FakeValueArgumentForLeftCallableReference(val ktExpression: KtCallableRefe
     override fun isExternal(): Boolean = false
 }
 
+// context here is context for value argument analysis
 internal fun createSimplePSICallArgument(
-        context: BasicCallResolutionContext,
+        contextForArgument: BasicCallResolutionContext,
         valueArgument: ValueArgument,
-        typeInfo: KotlinTypeInfo
+        typeInfoForArgument: KotlinTypeInfo
+) = createSimplePSICallArgument(contextForArgument.trace.bindingContext, contextForArgument.statementFilter,
+                                contextForArgument.scope.ownerDescriptor, valueArgument,
+                                contextForArgument.dataFlowInfo, typeInfoForArgument)
+
+internal fun createSimplePSICallArgument(
+        bindingContext: BindingContext,
+        statementFilter: StatementFilter,
+        ownerDescriptor: DeclarationDescriptor,
+        valueArgument: ValueArgument,
+        dataFlowInfoBeforeThisArgument: DataFlowInfo,
+        typeInfoForArgument: KotlinTypeInfo
 ): PSIKotlinCallArgument? {
-    val ktExpression = KtPsiUtil.getLastElementDeparenthesized(valueArgument.getArgumentExpression(), context.statementFilter) ?: return null
+
-    val onlyResolvedCall = ktExpression.getCall(context.trace.bindingContext)?.let {
+    val ktExpression = KtPsiUtil.getLastElementDeparenthesized(valueArgument.getArgumentExpression(), statementFilter) ?: return null
-        context.trace.bindingContext.get(BindingContext.ONLY_RESOLVED_CALL, it)
+    val onlyResolvedCall = ktExpression.getCall(bindingContext)?.let {
+        bindingContext.get(BindingContext.ONLY_RESOLVED_CALL, it)
     }
-    val baseType = onlyResolvedCall?.currentReturnType ?: typeInfo.type?.unwrap() ?: return null
+    val baseType = onlyResolvedCall?.currentReturnType ?: typeInfoForArgument.type?.unwrap() ?: return null
     val preparedType = prepareArgumentTypeRegardingCaptureTypes(baseType) ?: baseType
 
     // we should use DFI after this argument, because there can be some useful smartcast. Popular case: if branches.
-    val receiverToCast = context.replaceDataFlowInfo(typeInfo.dataFlowInfo).transformToReceiverWithSmartCastInfo(
+    val receiverToCast = transformToReceiverWithSmartCastInfo(
-            ExpressionReceiver.create(ktExpression, baseType, context.trace.bindingContext)
+            ownerDescriptor, bindingContext,
+            typeInfoForArgument.dataFlowInfo,
+            ExpressionReceiver.create(ktExpression, baseType, bindingContext)
     ).let {
         ReceiverValueWithSmartCastInfo(it.receiverValue.replaceType(preparedType), it.possibleTypes, it.isStable)
     }
 
     return if (onlyResolvedCall == null) {
-        ExpressionKotlinCallArgumentImpl(valueArgument, context.dataFlowInfo, typeInfo.dataFlowInfo, receiverToCast)
+        ExpressionKotlinCallArgumentImpl(valueArgument, dataFlowInfoBeforeThisArgument, typeInfoForArgument.dataFlowInfo, receiverToCast)
     }
     else {
-        SubKotlinCallArgumentImpl(valueArgument, context.dataFlowInfo, typeInfo.dataFlowInfo, receiverToCast, onlyResolvedCall)
+        SubKotlinCallArgumentImpl(valueArgument, dataFlowInfoBeforeThisArgument, typeInfoForArgument.dataFlowInfo, receiverToCast, onlyResolvedCall)
     }
 
 }

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

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.resolve.calls.tower
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.diagnostics.Errors
@@ -43,6 +44,7 @@ import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCallIm
 import org.jetbrains.kotlin.resolve.calls.results.OverloadResolutionResultsImpl
 import org.jetbrains.kotlin.resolve.calls.results.ResolutionResultsHandler
 import org.jetbrains.kotlin.resolve.calls.results.ResolutionStatus
+import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.calls.tasks.*
 import org.jetbrains.kotlin.resolve.descriptorUtil.hasDynamicExtensionAnnotation
@@ -457,8 +459,16 @@ class NewResolutionOldInference(
 
 }
 
-fun ResolutionContext<*>.transformToReceiverWithSmartCastInfo(receiver: ReceiverValue): ReceiverValueWithSmartCastInfo {
+fun ResolutionContext<*>.transformToReceiverWithSmartCastInfo(receiver: ReceiverValue) =
-    val dataFlowValue = DataFlowValueFactory.createDataFlowValue(receiver, this)
+        transformToReceiverWithSmartCastInfo(scope.ownerDescriptor, trace.bindingContext, dataFlowInfo, receiver)
+
+fun transformToReceiverWithSmartCastInfo(
+        containingDescriptor: DeclarationDescriptor,
+        bindingContext: BindingContext,
+        dataFlowInfo: DataFlowInfo,
+        receiver: ReceiverValue
+): ReceiverValueWithSmartCastInfo {
+    val dataFlowValue = DataFlowValueFactory.createDataFlowValue(receiver, bindingContext, containingDescriptor)
     return ReceiverValueWithSmartCastInfo(receiver, dataFlowInfo.getCollectedTypes(dataFlowValue), dataFlowValue.isStable)
 }
 

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ControlStructureTypingVisitor.java

@@ -36,6 +36,7 @@ import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory;
+import org.jetbrains.kotlin.resolve.calls.tower.KotlinResolutionCallbacksImpl;
 import org.jetbrains.kotlin.resolve.descriptorUtil.DescriptorUtilsKt;
 import org.jetbrains.kotlin.resolve.inline.InlineUtil;
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope;
@@ -611,6 +612,8 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
 
         KtExpression returnedExpression = expression.getReturnedExpression();
 
+        KotlinResolutionCallbacksImpl.LambdaInfo newInferenceLambdaInfo = null;
+
         KotlinType expectedType = NO_EXPECTED_TYPE;
         KotlinType resultType = components.builtIns.getNothingType();
         KtDeclaration parentDeclaration = context.getContextParentOfType(expression, KtDeclaration.class);
@@ -642,6 +645,7 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
                 }
 
                 expectedType = getFunctionExpectedReturnType(containingFunctionDescriptor, (KtElement) containingFunInfo.getSecond(), context);
+                newInferenceLambdaInfo = getNewInferenceLambdaInfo(context, (KtElement) containingFunInfo.getSecond());
             }
             else {
                 // Outside a function
@@ -653,6 +657,7 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
             SimpleFunctionDescriptor functionDescriptor = context.trace.get(FUNCTION, labelTargetElement);
             if (functionDescriptor != null) {
                 expectedType = getFunctionExpectedReturnType(functionDescriptor, labelTargetElement, context);
+                newInferenceLambdaInfo = getNewInferenceLambdaInfo(context, labelTargetElement);
                 if (!InlineUtil.checkNonLocalReturnUsage(functionDescriptor, expression, context)) {
                     // Qualified, non-local
                     context.trace.report(RETURN_NOT_ALLOWED.on(expression));
@@ -662,7 +667,14 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
         }
 
         if (returnedExpression != null) {
-            facade.getTypeInfo(returnedExpression, context.replaceExpectedType(expectedType).replaceContextDependency(INDEPENDENT));
+            if (newInferenceLambdaInfo != null) {
+                KotlinTypeInfo result = facade.getTypeInfo(returnedExpression, context.replaceExpectedType(newInferenceLambdaInfo.getExpectedType())
+                        .replaceContextDependency(newInferenceLambdaInfo.getContextDependency()));
+                newInferenceLambdaInfo.getReturnStatements().add(new kotlin.Pair<>(expression, result));
+            }
+            else {
+                facade.getTypeInfo(returnedExpression, context.replaceExpectedType(expectedType).replaceContextDependency(INDEPENDENT));
+            }
         }
         else {
             // for lambda with implicit return type Unit
@@ -692,6 +704,17 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
                 replaceJumpOutPossible(true);
     }
 
+    @Nullable
+    private static KotlinResolutionCallbacksImpl.LambdaInfo getNewInferenceLambdaInfo(
+            @NotNull ExpressionTypingContext context,
+            @NotNull KtElement function
+    ) {
+        if (function instanceof KtFunction) {
+            return context.trace.get(BindingContext.NEW_INFERENCE_LAMBDA_INFO, (KtFunction) function);
+        }
+        return null;
+    }
+
     @NotNull
     private static KotlinType getFunctionExpectedReturnType(
             @NotNull FunctionDescriptor descriptor,

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt

@@ -225,6 +225,8 @@ 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
+
         return computeReturnTypeBasedOnReturnExpressions(functionLiteral, context, typeOfBodyExpression)
     }