[NI] Improve inference for lambdas that return lambdas

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

@@ -19,19 +19,18 @@ import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.TemporaryBindingTrace
+import org.jetbrains.kotlin.resolve.TypeResolver
 import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
 import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
 import org.jetbrains.kotlin.resolve.calls.components.KotlinResolutionCallbacks
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency
-import org.jetbrains.kotlin.resolve.calls.model.LambdaKotlinCallArgument
+import org.jetbrains.kotlin.resolve.calls.model.*
-import org.jetbrains.kotlin.resolve.calls.model.ReceiverKotlinCallArgument
-import org.jetbrains.kotlin.resolve.calls.model.ResolvedCallAtom
-import org.jetbrains.kotlin.resolve.calls.model.SimpleKotlinCallArgument
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.calls.util.CallMaker
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
+import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValueWithSmartCastInfo
 import org.jetbrains.kotlin.types.TypeApproximator
@@ -44,6 +43,13 @@ import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
+data class LambdaContextInfo(
+    var typeInfo: KotlinTypeInfo? = null,
+    var dataFlowInfoAfter: DataFlowInfo? = null,
+    var lexicalScope: LexicalScope? = null,
+    var trace: BindingTrace? = null
+)
+
 class KotlinResolutionCallbacksImpl(
     val trace: BindingTrace,
     val expressionTypingServices: ExpressionTypingServices,
@@ -51,13 +57,14 @@ class KotlinResolutionCallbacksImpl(
     val argumentTypeResolver: ArgumentTypeResolver,
     val languageVersionSettings: LanguageVersionSettings,
     val kotlinToResolvedCallTransformer: KotlinToResolvedCallTransformer,
-    val constantExpressionEvaluator: ConstantExpressionEvaluator,
     val dataFlowValueFactory: DataFlowValueFactory,
-    override val inferenceSession: InferenceSession
+    override val inferenceSession: InferenceSession,
+    val constantExpressionEvaluator: ConstantExpressionEvaluator,
+    val typeResolver: TypeResolver
 ) : KotlinResolutionCallbacks {
     class LambdaInfo(val expectedType: UnwrappedType, val contextDependency: ContextDependency) {
-        var dataFlowInfoAfter: DataFlowInfo? = null
+        val returnStatements = ArrayList<Pair<KtReturnExpression, LambdaContextInfo?>>()
-        val returnStatements = ArrayList<Pair<KtReturnExpression, KotlinTypeInfo?>>()
+        val lastExpressionInfo = LambdaContextInfo()
 
         companion object {
             val STUB_EMPTY = LambdaInfo(TypeUtils.NO_EXPECTED_TYPE, ContextDependency.INDEPENDENT)
@@ -70,16 +77,34 @@ class KotlinResolutionCallbacksImpl(
         receiverType: UnwrappedType?,
         parameters: List<UnwrappedType>,
         expectedReturnType: UnwrappedType?
-    ): List<SimpleKotlinCallArgument> {
+    ): List<KotlinCallArgument> {
         val psiCallArgument = lambdaArgument.psiCallArgument as PSIFunctionKotlinCallArgument
         val outerCallContext = psiCallArgument.outerCallContext
 
-        fun createCallArgument(ktExpression: KtExpression, typeInfo: KotlinTypeInfo) =
+        fun createCallArgument(
-            createSimplePSICallArgument(
+            ktExpression: KtExpression,
+            typeInfo: KotlinTypeInfo,
+            scope: LexicalScope?,
+            newTrace: BindingTrace?
+        ): PSIKotlinCallArgument? {
+            var newContext = outerCallContext
+            if (scope != null) newContext = newContext.replaceScope(scope)
+            if (newTrace != null) newContext = newContext.replaceBindingTrace(newTrace)
+
+            processFunctionalExpression(
+                newContext, ktExpression, typeInfo.dataFlowInfo, CallMaker.makeExternalValueArgument(ktExpression),
+                null, outerCallContext.scope.ownerDescriptor.builtIns, typeResolver
+            )?.let {
+                it.setResultDataFlowInfoIfRelevant(typeInfo.dataFlowInfo)
+                return it
+            }
+
+            return createSimplePSICallArgument(
                 trace.bindingContext, outerCallContext.statementFilter, outerCallContext.scope.ownerDescriptor,
                 CallMaker.makeExternalValueArgument(ktExpression), DataFlowInfo.EMPTY, typeInfo, languageVersionSettings,
                 dataFlowValueFactory
             )
+        }
 
         val lambdaInfo = LambdaInfo(
             expectedReturnType ?: TypeUtils.NO_EXPECTED_TYPE,
@@ -107,12 +132,14 @@ class KotlinResolutionCallbacksImpl(
         trace.record(BindingContext.NEW_INFERENCE_LAMBDA_INFO, psiCallArgument.ktFunction, LambdaInfo.STUB_EMPTY)
 
         var hasReturnWithoutExpression = false
-        val returnArguments = lambdaInfo.returnStatements.mapNotNullTo(ArrayList()) { (expression, typeInfo) ->
+        val returnArguments = lambdaInfo.returnStatements.mapNotNullTo(ArrayList()) { (expression, contextInfo) ->
             val returnedExpression = expression.returnedExpression
             if (returnedExpression != null) {
                 createCallArgument(
                     returnedExpression,
-                    typeInfo ?: throw AssertionError("typeInfo should be non-null for return with expression")
+                    contextInfo?.typeInfo ?: throw AssertionError("typeInfo should be non-null for return with expression"),
+                    contextInfo.lexicalScope,
+                    contextInfo.trace
                 )
             } else {
                 hasReturnWithoutExpression = true
@@ -125,8 +152,9 @@ class KotlinResolutionCallbacksImpl(
 
             // todo lastExpression can be if without else
             val lastExpressionType = if (hasReturnWithoutExpression) null else trace.getType(lastExpression)
-            val lastExpressionTypeInfo = KotlinTypeInfo(lastExpressionType, lambdaInfo.dataFlowInfoAfter ?: functionTypeInfo.dataFlowInfo)
+            val contextInfo = lambdaInfo.lastExpressionInfo
-            createCallArgument(lastExpression, lastExpressionTypeInfo)
+            val lastExpressionTypeInfo = KotlinTypeInfo(lastExpressionType, contextInfo.dataFlowInfoAfter ?: functionTypeInfo.dataFlowInfo)
+            createCallArgument(lastExpression, lastExpressionTypeInfo, contextInfo.lexicalScope, contextInfo.trace)
         }
 
         returnArguments.addIfNotNull(lastExpressionArgument)
@@ -185,4 +213,4 @@ class KotlinResolutionCallbacksImpl(
         if (receiver == null) return callElement
         return PsiTreeUtil.findCommonParent(callElement, receiver.psiExpression)?.safeAs() ?: callElement
     }
-}
+}

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

@@ -9,15 +9,20 @@ import com.intellij.psi.impl.source.tree.LeafPsiElement
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.StatementFilter
+import org.jetbrains.kotlin.resolve.TypeResolver
+import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
 import org.jetbrains.kotlin.resolve.calls.callUtil.getCall
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
+import org.jetbrains.kotlin.resolve.lazy.ForceResolveUtil
 import org.jetbrains.kotlin.resolve.scopes.receivers.*
 import org.jetbrains.kotlin.types.ErrorUtils
 import org.jetbrains.kotlin.types.UnwrappedType
@@ -190,6 +195,75 @@ internal fun KotlinCallArgument.setResultDataFlowInfoIfRelevant(resultDataFlowIn
     }
 }
 
+fun processFunctionalExpression(
+    outerCallContext: BasicCallResolutionContext,
+    argumentExpression: KtExpression,
+    startDataFlowInfo: DataFlowInfo,
+    valueArgument: ValueArgument,
+    argumentName: Name?,
+    builtIns: KotlinBuiltIns,
+    typeResolver: TypeResolver
+): PSIKotlinCallArgument? {
+    val expression = ArgumentTypeResolver.getFunctionLiteralArgumentIfAny(argumentExpression, outerCallContext) ?: return null
+    val postponedExpression = if (expression is KtFunctionLiteral) expression.getParentOfType<KtLambdaExpression>(true) else expression
+
+    val lambdaArgument: PSIKotlinCallArgument? = when (postponedExpression) {
+        is KtLambdaExpression ->
+            LambdaKotlinCallArgumentImpl(
+                outerCallContext, valueArgument, startDataFlowInfo, argumentName, postponedExpression, argumentExpression,
+                resolveParametersTypes(outerCallContext, postponedExpression.functionLiteral, typeResolver)
+            )
+
+        is KtNamedFunction -> {
+            val receiverType = resolveType(outerCallContext, postponedExpression.receiverTypeReference, typeResolver)
+            val parametersTypes = resolveParametersTypes(outerCallContext, postponedExpression, typeResolver) ?: emptyArray()
+            val returnType = resolveType(outerCallContext, postponedExpression.typeReference, typeResolver)
+                    ?: 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
+}
+
+fun checkNoSpread(context: BasicCallResolutionContext, valueArgument: ValueArgument) {
+    valueArgument.getSpreadElement()?.let {
+        context.trace.report(Errors.SPREAD_OF_LAMBDA_OR_CALLABLE_REFERENCE.on(it))
+    }
+}
+
+private fun resolveParametersTypes(
+    context: BasicCallResolutionContext,
+    ktFunction: KtFunction,
+    typeResolver: TypeResolver
+): Array<UnwrappedType?>? {
+    val parameterList = ktFunction.valueParameterList ?: return null
+
+    return Array(parameterList.parameters.size) {
+        parameterList.parameters[it]?.typeReference?.let { resolveType(context, it, typeResolver) }
+    }
+}
+
+internal fun resolveType(
+    context: BasicCallResolutionContext,
+    typeReference: KtTypeReference?,
+    typeResolver: TypeResolver
+): UnwrappedType? {
+    if (typeReference == null) return null
+
+    val type = typeResolver.resolveType(context.scope, typeReference, context.trace, checkBounds = true)
+    ForceResolveUtil.forceResolveAllContents(type)
+    return type.unwrap()
+}
+
 
 // context here is context for value argument analysis
 internal fun createSimplePSICallArgument(

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

@@ -5,7 +5,6 @@
 
 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
@@ -14,7 +13,6 @@ import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.referenceExpression
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
@@ -38,7 +36,6 @@ import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
 import org.jetbrains.kotlin.resolve.calls.util.CallMaker
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
-import org.jetbrains.kotlin.resolve.lazy.ForceResolveUtil
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.resolve.scopes.SyntheticScopes
@@ -157,7 +154,7 @@ class PSICallResolver(
         KotlinResolutionCallbacksImpl(
             trace, expressionTypingServices, typeApproximator,
             argumentTypeResolver, languageVersionSettings, kotlinToResolvedCallTransformer,
-            constantExpressionEvaluator, dataFlowValueFactory, inferenceSession
+            dataFlowValueFactory, inferenceSession, constantExpressionEvaluator, typeResolver
         )
 
     private fun calculateExpectedType(context: BasicCallResolutionContext): UnwrappedType? {
@@ -540,14 +537,6 @@ class PSICallResolver(
             else -> error("Incorrect receiver: $oldReceiver")
         }
 
-    private fun resolveType(context: BasicCallResolutionContext, typeReference: KtTypeReference?): UnwrappedType? {
-        if (typeReference == null) return null
-
-        val type = typeResolver.resolveType(context.scope, typeReference, context.trace, checkBounds = true)
-        ForceResolveUtil.forceResolveAllContents(type)
-        return type.unwrap()
-    }
-
     private fun resolveTypeArguments(context: BasicCallResolutionContext, typeArguments: List<KtTypeProjection>): List<TypeArgument> =
         typeArguments.map { projection ->
             if (projection.projectionKind != KtProjectionKind.NONE) {
@@ -555,7 +544,7 @@ class PSICallResolver(
             }
             ModifierCheckerCore.check(projection, context.trace, null, languageVersionSettings)
 
-            resolveType(context, projection.typeReference)?.let { SimpleTypeArgumentImpl(projection.typeReference!!, it) }
+            resolveType(context, projection.typeReference, typeResolver)?.let { SimpleTypeArgumentImpl(projection.typeReference!!, it) }
                     ?: TypeArgumentPlaceholder
         }
 
@@ -584,7 +573,10 @@ class PSICallResolver(
 
         val argumentName = valueArgument.getArgumentName()?.asName
 
-        processFunctionalExpression(outerCallContext, argumentExpression, startDataFlowInfo, valueArgument, argumentName, builtIns)?.let {
+        processFunctionalExpression(
+            outerCallContext, argumentExpression, startDataFlowInfo,
+            valueArgument, argumentName, builtIns, typeResolver
+        )?.let {
             return it
         }
 
@@ -647,56 +639,4 @@ class PSICallResolver(
         val typeInfo = expressionTypingServices.getTypeInfo(argumentExpression, context)
         return createSimplePSICallArgument(context, valueArgument, typeInfo) ?: parseErrorArgument
     }
-
-    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 lambdaArgument: PSIKotlinCallArgument? = when (postponedExpression) {
-            is KtLambdaExpression ->
-                LambdaKotlinCallArgumentImpl(
-                    outerCallContext, valueArgument, startDataFlowInfo, argumentName, postponedExpression,
-                    argumentExpression, resolveParametersTypes(outerCallContext, postponedExpression.functionLiteral)
-                )
-
-            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) {
-        valueArgument.getSpreadElement()?.let {
-            context.trace.report(Errors.SPREAD_OF_LAMBDA_OR_CALLABLE_REFERENCE.on(it))
-        }
-    }
-
-    private fun resolveParametersTypes(context: BasicCallResolutionContext, ktFunction: KtFunction): Array<UnwrappedType?>? {
-        val parameterList = ktFunction.valueParameterList ?: return null
-
-        return Array(parameterList.parameters.size) {
-            parameterList.parameters[it]?.typeReference?.let { resolveType(context, it) }
-        }
-    }
 }

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

@@ -1,17 +1,6 @@
 /*
- * Copyright 2010-2017 JetBrains s.r.o.
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
- *
+ * that can be found in the license/LICENSE.txt file.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
  */
 
 package org.jetbrains.kotlin.types.expressions;
@@ -37,8 +26,8 @@ import org.jetbrains.kotlin.resolve.calls.model.MutableDataFlowInfoForArguments;
 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.calls.tower.LambdaContextInfo;
 import org.jetbrains.kotlin.resolve.descriptorUtil.DescriptorUtilsKt;
 import org.jetbrains.kotlin.resolve.inline.InlineUtil;
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope;
@@ -687,9 +676,21 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
 
         if (returnedExpression != null) {
             if (newInferenceLambdaInfo != null) {
-                KotlinTypeInfo result = facade.getTypeInfo(returnedExpression, context.replaceExpectedType(newInferenceLambdaInfo.getExpectedType())
+                LambdaContextInfo contextInfo;
-                        .replaceContextDependency(newInferenceLambdaInfo.getContextDependency()));
+                if (returnedExpression instanceof KtLambdaExpression) {
-                newInferenceLambdaInfo.getReturnStatements().add(new kotlin.Pair<>(expression, result));
+                    contextInfo = new LambdaContextInfo(
+                            new KotlinTypeInfo(DONT_CARE, context.dataFlowInfo),
+                            null,
+                            context.scope,
+                            context.trace
+                    );
+                } else {
+                    KotlinTypeInfo result = facade
+                            .getTypeInfo(returnedExpression, context.replaceExpectedType(newInferenceLambdaInfo.getExpectedType())
+                            .replaceContextDependency(newInferenceLambdaInfo.getContextDependency()));
+                    contextInfo = new LambdaContextInfo(result, null, context.scope, context.trace);
+                }
+                newInferenceLambdaInfo.getReturnStatements().add(new kotlin.Pair<>(expression, contextInfo));
             }
             else {
                 facade.getTypeInfo(returnedExpression, context.replaceExpectedType(expectedType).replaceContextDependency(INDEPENDENT));

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

@@ -1,21 +1,11 @@
 /*
- * Copyright 2010-2017 JetBrains s.r.o.
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
- *
+ * that can be found in the license/LICENSE.txt file.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
  */
 
 package org.jetbrains.kotlin.types.expressions;
 
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.tree.IElementType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -27,11 +17,13 @@ import org.jetbrains.kotlin.descriptors.FunctionDescriptor;
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor;
 import org.jetbrains.kotlin.lexer.KtTokens;
 import org.jetbrains.kotlin.psi.*;
+import org.jetbrains.kotlin.psi.psiUtil.PsiUtilsKt;
 import org.jetbrains.kotlin.resolve.*;
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency;
 import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue;
+import org.jetbrains.kotlin.resolve.calls.tower.KotlinResolutionCallbacksImpl;
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope;
 import org.jetbrains.kotlin.resolve.scopes.LexicalScopeKind;
 import org.jetbrains.kotlin.resolve.scopes.LexicalWritableScope;
@@ -44,8 +36,7 @@ import org.jetbrains.kotlin.util.slicedMap.WritableSlice;
 import java.util.Iterator;
 import java.util.List;
 
-import static org.jetbrains.kotlin.types.TypeUtils.NO_EXPECTED_TYPE;
+import static org.jetbrains.kotlin.types.TypeUtils.*;
-import static org.jetbrains.kotlin.types.TypeUtils.UNIT_EXPECTED_TYPE;
 import static org.jetbrains.kotlin.types.expressions.CoercionStrategy.COERCION_TO_UNIT;
 
 public class ExpressionTypingServices {
@@ -344,6 +335,20 @@ public class ExpressionTypingServices {
 
             return blockLevelVisitor.getTypeInfo(statementExpression, context.replaceExpectedType(expectedType).replaceContextDependency(dependency), true);
         }
+        if (context.languageVersionSettings.supportsFeature(LanguageFeature.NewInference) &&
+            statementExpression instanceof KtLambdaExpression) {
+            PsiElement parent = PsiUtilsKt.getNonStrictParentOfType(statementExpression, KtFunctionLiteral.class);
+            if (parent != null) {
+                KtFunctionLiteral functionLiteral = (KtFunctionLiteral) parent;
+                KotlinResolutionCallbacksImpl.LambdaInfo info =
+                        context.trace.getBindingContext().get(BindingContext.NEW_INFERENCE_LAMBDA_INFO, functionLiteral);
+                if (info != null) {
+                    info.getLastExpressionInfo().setLexicalScope(context.scope);
+                    info.getLastExpressionInfo().setTrace(context.trace);
+                    return new KotlinTypeInfo(DONT_CARE, context.dataFlowInfo);
+                }
+            }
+        }
         KotlinTypeInfo result = blockLevelVisitor.getTypeInfo(statementExpression, context, true);
         if (coercionStrategyForLastExpression == COERCION_TO_UNIT) {
             boolean mightBeUnit = false;

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

@@ -1,17 +1,6 @@
 /*
- * Copyright 2010-2017 JetBrains s.r.o.
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
- *
+ * that can be found in the license/LICENSE.txt file.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
  */
 
 package org.jetbrains.kotlin.types.expressions
@@ -260,7 +249,7 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
         val typeOfBodyExpression = blockReturnedType.type
 
         newInferenceLambdaInfo?.let {
-            it.dataFlowInfoAfter = blockReturnedType.dataFlowInfo
+            it.lastExpressionInfo.dataFlowInfoAfter = blockReturnedType.dataFlowInfo
             return null
         }
 

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

@@ -32,7 +32,7 @@ interface KotlinResolutionCallbacks {
         receiverType: UnwrappedType?,
         parameters: List<UnwrappedType>,
         expectedReturnType: UnwrappedType? // null means, that return type is not proper i.e. it depends on some type variables
-    ): List<SimpleKotlinCallArgument>
+    ): List<KotlinCallArgument>
 
     fun bindStubResolvedCallForCandidate(candidate: ResolvedCallAtom)
 

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

@@ -1,17 +1,6 @@
 /*
- * Copyright 2010-2017 JetBrains s.r.o.
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
- *
+ * that can be found in the license/LICENSE.txt file.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
  */
 
 package org.jetbrains.kotlin.resolve.calls.components
@@ -97,7 +86,7 @@ class PostponedArgumentsAnalyzer(
         returnArguments.forEach { c.addSubsystemFromArgument(it) }
 
         val subResolvedKtPrimitives = returnArguments.map {
-            checkSimpleArgument(c.getBuilder(), it, lambda.returnType.let(::substitute), diagnosticHolder, isReceiver = false)
+            resolveKtPrimitive(c.getBuilder(), it, lambda.returnType.let(::substitute), diagnosticHolder, isReceiver = false)
         }
 
         if (returnArguments.isEmpty()) {

compiler/testData/diagnostics/tests/inference/constraints/returnLambdaFromLambda.kt

@@ -0,0 +1,22 @@
+// !DIAGNOSTICS: -UNUSED_VARIABLE
+
+fun testLambda() {
+    val basicTest: (Int) -> Int = myRun {
+        val x: Any? = null
+        if (x is String) return@myRun { it -> <!DEBUG_INFO_SMARTCAST!>x<!>.length + it }
+        if (x !is Int) return@myRun { it -> it }
+
+        { it -> <!DEBUG_INFO_SMARTCAST!>x<!> + it }
+    }
+
+    val twoLambda: (Int) -> Int = myRun {
+        val x: Int = 1
+        run {
+            val y: Int = 2
+            { x + y }
+        }
+    }
+
+}
+
+inline fun <R> myRun(block: () -> R): R = block()

compiler/testData/diagnostics/tests/inference/constraints/returnLambdaFromLambda.txt

@@ -0,0 +1,4 @@
+package
+
+public inline fun </*0*/ R> myRun(/*0*/ block: () -> R): R
+public fun testLambda(): kotlin.Unit

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -10903,6 +10903,12 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                     doTest(fileName);
                 }
 
+                @TestMetadata("returnLambdaFromLambda.kt")
+                public void testReturnLambdaFromLambda() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/inference/constraints/returnLambdaFromLambda.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("subtypeConstraintOnNullableType.kt")
                 public void testSubtypeConstraintOnNullableType() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/inference/constraints/subtypeConstraintOnNullableType.kt");

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsUsingJavacTestGenerated.java

@@ -10903,6 +10903,12 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                     doTest(fileName);
                 }
 
+                @TestMetadata("returnLambdaFromLambda.kt")
+                public void testReturnLambdaFromLambda() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/inference/constraints/returnLambdaFromLambda.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("subtypeConstraintOnNullableType.kt")
                 public void testSubtypeConstraintOnNullableType() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/inference/constraints/subtypeConstraintOnNullableType.kt");