KtExpression.analyzeInContext uses correct context when analyzing return's

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/checkers/CallChecker.kt

@@ -35,15 +35,20 @@ interface CallChecker {
 }
 
 class CallCheckerContext(
+        val resolutionContext: ResolutionContext<*>,
         val trace: BindingTrace,
-        val scope: LexicalScope,
+        val languageFeatureSettings: LanguageFeatureSettings
-        val languageFeatureSettings: LanguageFeatureSettings,
-        val dataFlowInfo: DataFlowInfo,
-        val isAnnotationContext: Boolean
 ) {
-    constructor(c: ResolutionContext<*>, languageFeatureSettings: LanguageFeatureSettings) : this(
+    val scope: LexicalScope
-            c.trace, c.scope, languageFeatureSettings, c.dataFlowInfo, c.isAnnotationContext
+        get() = resolutionContext.scope
-    )
+
+    val dataFlowInfo: DataFlowInfo
+        get() = resolutionContext.dataFlowInfo
+
+    val isAnnotationContext: Boolean
+        get() = resolutionContext.isAnnotationContext
+
+    constructor(c: ResolutionContext<*>, languageFeatureSettings: LanguageFeatureSettings) : this(c, c.trace, languageFeatureSettings)
 }
 
 // Use this utility to avoid premature computation of deferred return type of a resolved callable descriptor.

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/checkers/InlineChecker.java

@@ -285,7 +285,7 @@ class InlineChecker implements CallChecker {
     ) {
         if (!allowsNonLocalReturns(inlinableParameterDescriptor)) return;
 
-        if (!checkNonLocalReturnUsage(descriptor, parameterUsage, context.getTrace())) {
+        if (!checkNonLocalReturnUsage(descriptor, parameterUsage, context.getResolutionContext())) {
             context.getTrace().report(NON_LOCAL_RETURN_NOT_ALLOWED.on(parameterUsage, parameterUsage));
         }
     }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/checkers/coroutineCallChecker.kt

@@ -33,7 +33,7 @@ object CoroutineSuspendCallChecker : CallChecker {
         val dispatchReceiverOwner = (resolvedCall.dispatchReceiver as? CoroutineReceiverValue)?.declarationDescriptor ?: return
         val callElement = resolvedCall.call.callElement as KtExpression
 
-        if (!InlineUtil.checkNonLocalReturnUsage(dispatchReceiverOwner, callElement, context.trace)) {
+        if (!InlineUtil.checkNonLocalReturnUsage(dispatchReceiverOwner, callElement, context.resolutionContext)) {
             context.trace.report(Errors.NON_LOCAL_SUSPENSION_POINT.on(reportOn))
         }
     }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/context/BasicCallResolutionContext.java

@@ -16,9 +16,11 @@
 
 package org.jetbrains.kotlin.resolve.calls.context;
 
+import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.psi.Call;
+import org.jetbrains.kotlin.psi.KtExpression;
 import org.jetbrains.kotlin.resolve.BindingTrace;
 import org.jetbrains.kotlin.resolve.StatementFilter;
 import org.jetbrains.kotlin.resolve.calls.model.MutableDataFlowInfoForArguments;
@@ -41,10 +43,12 @@ public class BasicCallResolutionContext extends CallResolutionContext<BasicCallR
             boolean isAnnotationContext,
             boolean isDebuggerContext,
             boolean collectAllCandidates,
-            @NotNull CallPosition callPosition
+            @NotNull CallPosition callPosition,
+            @NotNull Function1<KtExpression, KtExpression> expressionContextProvider
     ) {
         super(trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments, resolutionResultsCache,
-              dataFlowInfoForArguments, statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates, callPosition);
+              dataFlowInfoForArguments, statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates,
+              callPosition, expressionContextProvider);
     }
 
     @NotNull
@@ -61,7 +65,7 @@ public class BasicCallResolutionContext extends CallResolutionContext<BasicCallR
         return new BasicCallResolutionContext(trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments,
                                               new ResolutionResultsCacheImpl(), null,
                                               StatementFilter.NONE, isAnnotationContext, false, false,
-                                              CallPosition.Unknown.INSTANCE);
+                                              CallPosition.Unknown.INSTANCE, DEFAULT_EXPRESSION_CONTEXT_PROVIDER);
     }
 
     @NotNull
@@ -72,7 +76,8 @@ public class BasicCallResolutionContext extends CallResolutionContext<BasicCallR
         return new BasicCallResolutionContext(
                 context.trace, context.scope, call, context.expectedType, context.dataFlowInfo, context.contextDependency, checkArguments,
                 context.resolutionResultsCache, dataFlowInfoForArguments,
-                context.statementFilter, context.isAnnotationContext, context.isDebuggerContext, context.collectAllCandidates, context.callPosition);
+                context.statementFilter, context.isAnnotationContext, context.isDebuggerContext, context.collectAllCandidates,
+                context.callPosition, context.expressionContextProvider);
     }
 
     @NotNull
@@ -92,17 +97,20 @@ public class BasicCallResolutionContext extends CallResolutionContext<BasicCallR
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @NotNull StatementFilter statementFilter,
             boolean collectAllCandidates,
-            @NotNull CallPosition callPosition
+            @NotNull CallPosition callPosition,
+            @NotNull Function1<KtExpression, KtExpression> expressionContextProvider
     ) {
         return new BasicCallResolutionContext(
                 trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments, resolutionResultsCache,
-                dataFlowInfoForArguments, statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates, callPosition);
+                dataFlowInfoForArguments, statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates,
+                callPosition, expressionContextProvider);
     }
 
     @NotNull
     public BasicCallResolutionContext replaceCall(@NotNull Call newCall) {
         return new BasicCallResolutionContext(
                 trace, scope, newCall, expectedType, dataFlowInfo, contextDependency, checkArguments, resolutionResultsCache,
-                dataFlowInfoForArguments, statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates, callPosition);
+                dataFlowInfoForArguments, statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates,
+                callPosition, expressionContextProvider);
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/context/CallCandidateResolutionContext.java

@@ -16,10 +16,12 @@
 
 package org.jetbrains.kotlin.resolve.calls.context;
 
+import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.CallableDescriptor;
 import org.jetbrains.kotlin.psi.Call;
+import org.jetbrains.kotlin.psi.KtExpression;
 import org.jetbrains.kotlin.resolve.BindingTrace;
 import org.jetbrains.kotlin.resolve.StatementFilter;
 import org.jetbrains.kotlin.resolve.calls.model.MutableDataFlowInfoForArguments;
@@ -54,11 +56,12 @@ public final class CallCandidateResolutionContext<D extends CallableDescriptor>
             boolean isAnnotationContext,
             boolean isDebuggerContext,
             boolean collectAllCandidates,
-            @NotNull CallPosition callPosition
+            @NotNull CallPosition callPosition,
+            @NotNull Function1<KtExpression, KtExpression> expressionContextProvider
     ) {
         super(trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments, resolutionResultsCache,
               dataFlowInfoForArguments, statementFilter, isAnnotationContext, isDebuggerContext,
-              collectAllCandidates, callPosition);
+              collectAllCandidates, callPosition, expressionContextProvider);
         this.candidateCall = candidateCall;
         this.tracing = tracing;
         this.candidateResolveMode = candidateResolveMode;
@@ -74,7 +77,8 @@ public final class CallCandidateResolutionContext<D extends CallableDescriptor>
                 context.dataFlowInfo, context.contextDependency, context.checkArguments,
                 context.resolutionResultsCache, context.dataFlowInfoForArguments,
                 context.statementFilter,
-                candidateResolveMode, context.isAnnotationContext, context.isDebuggerContext, context.collectAllCandidates, context.callPosition);
+                candidateResolveMode, context.isAnnotationContext, context.isDebuggerContext, context.collectAllCandidates,
+                context.callPosition, context.expressionContextProvider);
     }
 
     @NotNull
@@ -86,7 +90,7 @@ public final class CallCandidateResolutionContext<D extends CallableDescriptor>
                 context.dataFlowInfo, context.contextDependency, context.checkArguments, context.resolutionResultsCache,
                 context.dataFlowInfoForArguments, context.statementFilter,
                 CandidateResolveMode.FULLY, context.isAnnotationContext, context.isDebuggerContext, context.collectAllCandidates,
-                context.callPosition);
+                context.callPosition, context.expressionContextProvider);
     }
 
     @Override
@@ -99,11 +103,12 @@ public final class CallCandidateResolutionContext<D extends CallableDescriptor>
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @NotNull StatementFilter statementFilter,
             boolean collectAllCandidates,
-            @NotNull CallPosition callPosition
+            @NotNull CallPosition callPosition,
+            @NotNull Function1<KtExpression, KtExpression> expressionContextProvider
     ) {
         return new CallCandidateResolutionContext<D>(
                 candidateCall, tracing, trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments,
                 resolutionResultsCache, dataFlowInfoForArguments, statementFilter,
-                candidateResolveMode, isAnnotationContext, isDebuggerContext, collectAllCandidates, callPosition);
+                candidateResolveMode, isAnnotationContext, isDebuggerContext, collectAllCandidates, callPosition, expressionContextProvider);
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/context/CallResolutionContext.java

@@ -16,9 +16,11 @@
 
 package org.jetbrains.kotlin.resolve.calls.context;
 
+import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.psi.Call;
+import org.jetbrains.kotlin.psi.KtExpression;
 import org.jetbrains.kotlin.resolve.BindingTrace;
 import org.jetbrains.kotlin.resolve.StatementFilter;
 import org.jetbrains.kotlin.resolve.calls.model.DataFlowInfoForArgumentsImpl;
@@ -50,10 +52,11 @@ public abstract class CallResolutionContext<Context extends CallResolutionContex
             boolean isAnnotationContext,
             boolean isDebuggerContext,
             boolean collectAllCandidates,
-            @NotNull CallPosition callPosition
+            @NotNull CallPosition callPosition,
+            @NotNull Function1<KtExpression, KtExpression> expressionContextProvider
     ) {
         super(trace, scope, expectedType, dataFlowInfo, contextDependency, resolutionResultsCache,
-              statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates, callPosition);
+              statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates, callPosition, expressionContextProvider);
         this.call = call;
         this.checkArguments = checkArguments;
         if (dataFlowInfoForArguments != null) {

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/context/ResolutionContext.java

@@ -16,8 +16,12 @@
 
 package org.jetbrains.kotlin.resolve.calls.context;
 
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.psi.KtExpression;
 import org.jetbrains.kotlin.resolve.BindingTrace;
 import org.jetbrains.kotlin.resolve.StatementFilter;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
@@ -56,6 +60,23 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
     @NotNull
     public final CallPosition callPosition;
 
+    /**
+     * Used for analyzing expression in the given context.
+     * Should be used for going through parents to find containing function, loop etc.
+     * The provider should return specific context expression (which can be used instead of parent)
+     * for the given expression or null otherwise.
+     * @see #getContextParentOfType
+     */
+    @NotNull
+    public final Function1<KtExpression, KtExpression> expressionContextProvider;
+
+    public static final Function1<KtExpression, KtExpression> DEFAULT_EXPRESSION_CONTEXT_PROVIDER = new Function1<KtExpression, KtExpression>() {
+        @Override
+        public KtExpression invoke(KtExpression expression) {
+            return null;
+        }
+    };
+
     protected ResolutionContext(
             @NotNull BindingTrace trace,
             @NotNull LexicalScope scope,
@@ -67,7 +88,8 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
             boolean isAnnotationContext,
             boolean isDebuggerContext,
             boolean collectAllCandidates,
-            @NotNull CallPosition callPosition
+            @NotNull CallPosition callPosition,
+            @NotNull Function1<KtExpression, KtExpression> expressionContextProvider
     ) {
         this.trace = trace;
         this.scope = scope;
@@ -80,6 +102,7 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
         this.isDebuggerContext = isDebuggerContext;
         this.collectAllCandidates = collectAllCandidates;
         this.callPosition = callPosition;
+        this.expressionContextProvider = expressionContextProvider;
     }
 
     protected abstract Context create(
@@ -91,7 +114,8 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @NotNull StatementFilter statementFilter,
             boolean collectAllCandidates,
-            @NotNull CallPosition callPosition
+            @NotNull CallPosition callPosition,
+            @NotNull Function1<KtExpression, KtExpression> expressionContextProvider
     );
 
     @NotNull
@@ -104,14 +128,14 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
     public Context replaceBindingTrace(@NotNull BindingTrace trace) {
         if (this.trace == trace) return self();
         return create(trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, statementFilter,
-                      collectAllCandidates, callPosition);
+                      collectAllCandidates, callPosition, expressionContextProvider);
     }
 
     @NotNull
     public Context replaceDataFlowInfo(@NotNull DataFlowInfo newDataFlowInfo) {
         if (newDataFlowInfo == dataFlowInfo) return self();
         return create(trace, scope, newDataFlowInfo, expectedType, contextDependency, resolutionResultsCache, statementFilter,
-                      collectAllCandidates, callPosition);
+                      collectAllCandidates, callPosition, expressionContextProvider);
     }
 
     @NotNull
@@ -119,28 +143,28 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
         if (newExpectedType == null) return replaceExpectedType(TypeUtils.NO_EXPECTED_TYPE);
         if (expectedType == newExpectedType) return self();
         return create(trace, scope, dataFlowInfo, newExpectedType, contextDependency, resolutionResultsCache, statementFilter,
-                      collectAllCandidates, callPosition);
+                      collectAllCandidates, callPosition, expressionContextProvider);
     }
 
     @NotNull
     public Context replaceScope(@NotNull LexicalScope newScope) {
         if (newScope == scope) return self();
         return create(trace, newScope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, statementFilter,
-                      collectAllCandidates, callPosition);
+                      collectAllCandidates, callPosition, expressionContextProvider);
     }
 
     @NotNull
     public Context replaceContextDependency(@NotNull ContextDependency newContextDependency) {
         if (newContextDependency == contextDependency) return self();
         return create(trace, scope, dataFlowInfo, expectedType, newContextDependency, resolutionResultsCache, statementFilter,
-                      collectAllCandidates, callPosition);
+                      collectAllCandidates, callPosition, expressionContextProvider);
     }
 
     @NotNull
     public Context replaceResolutionResultsCache(@NotNull ResolutionResultsCache newResolutionResultsCache) {
         if (newResolutionResultsCache == resolutionResultsCache) return self();
         return create(trace, scope, dataFlowInfo, expectedType, contextDependency, newResolutionResultsCache, statementFilter,
-                      collectAllCandidates, callPosition);
+                      collectAllCandidates, callPosition, expressionContextProvider);
     }
 
     @NotNull
@@ -151,18 +175,50 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
     @NotNull
     public Context replaceCollectAllCandidates(boolean newCollectAllCandidates) {
         return create(trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, statementFilter,
-                      newCollectAllCandidates, callPosition);
+                      newCollectAllCandidates, callPosition, expressionContextProvider);
     }
 
     @NotNull
     public Context replaceStatementFilter(@NotNull StatementFilter statementFilter) {
         return create(trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, statementFilter,
-                      collectAllCandidates, callPosition);
+                      collectAllCandidates, callPosition, expressionContextProvider);
     }
 
     @NotNull
     public Context replaceCallPosition(@NotNull CallPosition callPosition) {
         return create(trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, statementFilter,
-                      collectAllCandidates, callPosition);
+                      collectAllCandidates, callPosition, expressionContextProvider);
+    }
+
+    @NotNull
+    public Context replaceExpressionContextProvider(@NotNull Function1<KtExpression, KtExpression> expressionContextProvider) {
+        return create(trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, statementFilter,
+                      collectAllCandidates, callPosition, expressionContextProvider);
+    }
+
+    @Nullable
+    public <T extends PsiElement> T getContextParentOfType(@NotNull KtExpression expression, @NotNull Class<? extends T>... classes) {
+        PsiElement current = expression.getParent();
+        while (current != null) {
+            for (Class<? extends T> klass : classes) {
+                if (klass.isInstance(current)) {
+                    //noinspection unchecked
+                    return (T) current;
+                }
+            }
+
+            if (current instanceof PsiFile) return null;
+
+            if (current instanceof KtExpression) {
+                KtExpression context = expressionContextProvider.invoke((KtExpression) current);
+                if (context != null) {
+                    current = context;
+                    continue;
+                }
+            }
+
+            current = current.getParent();
+        }
+        return null;
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/inline/InlineUtil.java

@@ -17,7 +17,6 @@
 package org.jetbrains.kotlin.resolve.inline;
 
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.FunctionTypesKt;
@@ -25,9 +24,9 @@ import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.BindingContext;
-import org.jetbrains.kotlin.resolve.BindingTrace;
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils;
 import org.jetbrains.kotlin.resolve.calls.callUtil.CallUtilKt;
+import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext;
 import org.jetbrains.kotlin.resolve.calls.model.ArgumentMapping;
 import org.jetbrains.kotlin.resolve.calls.model.ArgumentMatch;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
@@ -81,19 +80,19 @@ public class InlineUtil {
     public static boolean checkNonLocalReturnUsage(
             @NotNull DeclarationDescriptor fromFunction,
             @NotNull KtExpression startExpression,
-            @NotNull BindingTrace trace
+            @NotNull ResolutionContext<?> context
     ) {
-        PsiElement containingFunction = PsiTreeUtil.getParentOfType(startExpression, KtClassOrObject.class, KtDeclarationWithBody.class);
+        PsiElement containingFunction = context.getContextParentOfType(startExpression, KtClassOrObject.class, KtDeclarationWithBody.class);
         if (containingFunction == null) {
             return false;
         }
 
-        DeclarationDescriptor containingFunctionDescriptor = trace.get(BindingContext.DECLARATION_TO_DESCRIPTOR, containingFunction);
+        DeclarationDescriptor containingFunctionDescriptor = context.trace.get(BindingContext.DECLARATION_TO_DESCRIPTOR, containingFunction);
         if (containingFunctionDescriptor == null) {
             return false;
         }
 
-        BindingContext bindingContext = trace.getBindingContext();
+        BindingContext bindingContext = context.trace.getBindingContext();
 
         while (canBeInlineArgument(containingFunction) && fromFunction != containingFunctionDescriptor) {
             if (!isInlinedArgument((KtFunction) containingFunction, bindingContext, true)) {

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

@@ -909,9 +909,7 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
                 if (resolvedCall != null) {
                     // Call must be validated with the actual, not temporary trace in order to report operator diagnostic
                     // Only unary assignment expressions (++, --) and +=/... must be checked, normal assignments have the proper trace
-                    CallCheckerContext callCheckerContext = new CallCheckerContext(
+                    CallCheckerContext callCheckerContext = new CallCheckerContext(context, trace, components.languageFeatureSettings);
-                            trace, context.scope, components.languageFeatureSettings, context.dataFlowInfo, context.isAnnotationContext
-                    );
                     for (CallChecker checker : components.callCheckers) {
                         checker.check(resolvedCall, expression, callCheckerContext);
                     }

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

@@ -19,7 +19,6 @@ package org.jetbrains.kotlin.types.expressions;
 import com.google.common.collect.Lists;
 import com.intellij.openapi.util.Pair;
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
@@ -555,7 +554,7 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
 
         KotlinType expectedType = NO_EXPECTED_TYPE;
         KotlinType resultType = components.builtIns.getNothingType();
-        KtDeclaration parentDeclaration = PsiTreeUtil.getParentOfType(expression, KtDeclaration.class);
+        KtDeclaration parentDeclaration = context.getContextParentOfType(expression, KtDeclaration.class);
 
         if (parentDeclaration instanceof KtParameter) {
             // In a default value for parameter
@@ -565,7 +564,7 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
         if (expression.getTargetLabel() == null) {
             while (parentDeclaration instanceof KtDestructuringDeclaration) {
                 //TODO: It's hacking fix for KT-5100: Strange "Return is not allowed here" for multi-declaration initializer with elvis expression
-                parentDeclaration = PsiTreeUtil.getParentOfType(parentDeclaration, KtDeclaration.class);
+                parentDeclaration = context.getContextParentOfType(parentDeclaration, KtDeclaration.class);
             }
 
             // Parent declaration can be null in code fragments or in some bad error expressions
@@ -576,7 +575,7 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
             FunctionDescriptor containingFunctionDescriptor = containingFunInfo.getFirst();
 
             if (containingFunctionDescriptor != null) {
-                if (!InlineUtil.checkNonLocalReturnUsage(containingFunctionDescriptor, expression, context.trace) ||
+                if (!InlineUtil.checkNonLocalReturnUsage(containingFunctionDescriptor, expression, context) ||
                     isClassInitializer(containingFunInfo)) {
                     // Unqualified, in a function literal
                     context.trace.report(RETURN_NOT_ALLOWED.on(expression));
@@ -595,7 +594,7 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
             SimpleFunctionDescriptor functionDescriptor = context.trace.get(FUNCTION, labelTargetElement);
             if (functionDescriptor != null) {
                 expectedType = getFunctionExpectedReturnType(functionDescriptor, labelTargetElement, context);
-                if (!InlineUtil.checkNonLocalReturnUsage(functionDescriptor, expression, context.trace)) {
+                if (!InlineUtil.checkNonLocalReturnUsage(functionDescriptor, expression, context)) {
                     // Qualified, non-local
                     context.trace.report(RETURN_NOT_ALLOWED.on(expression));
                     resultType = ErrorUtils.createErrorType(RETURN_NOT_ALLOWED_MESSAGE);

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

@@ -16,7 +16,9 @@
 
 package org.jetbrains.kotlin.types.expressions;
 
+import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.psi.KtExpression;
 import org.jetbrains.kotlin.resolve.BindingTrace;
 import org.jetbrains.kotlin.resolve.StatementFilter;
 import org.jetbrains.kotlin.resolve.calls.context.*;
@@ -45,7 +47,7 @@ public class ExpressionTypingContext extends ResolutionContext<ExpressionTypingC
                 context.contextDependency, context.resolutionResultsCache,
                 context.statementFilter,
                 context.isAnnotationContext, context.isDebuggerContext, context.collectAllCandidates,
-                context.callPosition);
+                context.callPosition, context.expressionContextProvider);
     }
 
     @NotNull
@@ -55,7 +57,7 @@ public class ExpressionTypingContext extends ResolutionContext<ExpressionTypingC
                 context.contextDependency, context.resolutionResultsCache,
                 context.statementFilter,
                 context.isAnnotationContext, isDebuggerContext, context.collectAllCandidates,
-                context.callPosition);
+                context.callPosition, context.expressionContextProvider);
     }
 
     @NotNull
@@ -71,7 +73,7 @@ public class ExpressionTypingContext extends ResolutionContext<ExpressionTypingC
     ) {
         return new ExpressionTypingContext(
                 trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache,
-                statementFilter, isAnnotationContext, false, false, CallPosition.Unknown.INSTANCE);
+                statementFilter, isAnnotationContext, false, false, CallPosition.Unknown.INSTANCE, DEFAULT_EXPRESSION_CONTEXT_PROVIDER);
     }
 
     private ExpressionTypingContext(
@@ -85,10 +87,11 @@ public class ExpressionTypingContext extends ResolutionContext<ExpressionTypingC
             boolean isAnnotationContext,
             boolean isDebuggerContext,
             boolean collectAllCandidates,
-            @NotNull CallPosition callPosition
+            @NotNull CallPosition callPosition,
+            @NotNull Function1<KtExpression, KtExpression> expressionContextProvider
     ) {
         super(trace, scope, expectedType, dataFlowInfo, contextDependency, resolutionResultsCache,
-              statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates, callPosition);
+              statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates, callPosition, expressionContextProvider);
     }
 
     @Override
@@ -101,10 +104,12 @@ public class ExpressionTypingContext extends ResolutionContext<ExpressionTypingC
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @NotNull StatementFilter statementFilter,
             boolean collectAllCandidates,
-            @NotNull CallPosition callPosition
+            @NotNull CallPosition callPosition,
+            @NotNull Function1<KtExpression, KtExpression> expressionContextProvider
     ) {
         return new ExpressionTypingContext(trace, scope, dataFlowInfo,
                                            expectedType, contextDependency, resolutionResultsCache,
-                                           statementFilter, isAnnotationContext, isDebuggerContext, collectAllCandidates, callPosition);
+                                           statementFilter, isAnnotationContext, isDebuggerContext,
+                                           collectAllCandidates, callPosition, expressionContextProvider);
     }
 }

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

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.types.expressions;
 
 import com.intellij.psi.tree.IElementType;
+import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
@@ -91,10 +92,31 @@ public class ExpressionTypingServices {
             @NotNull DataFlowInfo dataFlowInfo,
             @NotNull BindingTrace trace,
             boolean isStatement
+    ) {
+        return getTypeInfo(scope, expression, expectedType, dataFlowInfo, trace, isStatement, expression);
+    }
+
+    @NotNull
+    public KotlinTypeInfo getTypeInfo(
+            @NotNull LexicalScope scope,
+            @NotNull final KtExpression expression,
+            @NotNull KotlinType expectedType,
+            @NotNull DataFlowInfo dataFlowInfo,
+            @NotNull BindingTrace trace,
+            boolean isStatement,
+            @NotNull final KtExpression contextExpression
     ) {
         ExpressionTypingContext context = ExpressionTypingContext.newContext(
                 trace, scope, dataFlowInfo, expectedType
         );
+        if (contextExpression != expression) {
+            context = context.replaceExpressionContextProvider(new Function1<KtExpression, KtExpression>() {
+                @Override
+                public KtExpression invoke(KtExpression arg) {
+                    return arg == expression ? contextExpression : null;
+                }
+            });
+        }
         return expressionTypingFacade.getTypeInfo(expression, context, isStatement);
     }
 

idea/ide-common/src/org/jetbrains/kotlin/idea/analysis/AnalyzerUtil.kt

@@ -40,7 +40,7 @@ import org.jetbrains.kotlin.types.expressions.PreliminaryDeclarationVisitor
 ): KotlinTypeInfo {
     PreliminaryDeclarationVisitor.createForExpression(this, trace)
     return contextExpression.getResolutionFacade().frontendService<ExpressionTypingServices>()
-            .getTypeInfo(scope, this, expectedType, dataFlowInfo, trace, isStatement)
+            .getTypeInfo(scope, this, expectedType, dataFlowInfo, trace, isStatement, contextExpression)
 }
 
 @JvmOverloads fun KtExpression.analyzeInContext(