Used partial resolve filter to find the last statement in a block

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

@@ -31,6 +31,10 @@ import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver;
 import org.jetbrains.kotlin.resolve.calls.CallResolver;
 import org.jetbrains.kotlin.resolve.calls.callUtil.CallUtilPackage;
+import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker;
+import org.jetbrains.kotlin.resolve.calls.checkers.CompositeChecker;
+import org.jetbrains.kotlin.resolve.calls.context.ContextDependency;
+import org.jetbrains.kotlin.resolve.calls.context.SimpleResolutionContext;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedValueArgument;
 import org.jetbrains.kotlin.resolve.calls.results.OverloadResolutionResults;
@@ -346,7 +350,11 @@ public class AnnotationResolver {
                 CompileTimeConstant<?> constant = ConstantExpressionEvaluator.OBJECT$.evaluate(argumentExpression, trace, expectedType);
                 if (constant instanceof IntegerValueTypeConstant) {
                     JetType defaultType = ((IntegerValueTypeConstant) constant).getType(expectedType);
-                    ArgumentTypeResolver.updateNumberType(defaultType, argumentExpression, trace);
+                    SimpleResolutionContext context =
+                            new SimpleResolutionContext(trace, JetScope.Empty.INSTANCE$, NO_EXPECTED_TYPE, DataFlowInfo.EMPTY,
+                                                        ContextDependency.INDEPENDENT,
+                                                        new CompositeChecker(Lists.<CallChecker>newArrayList()), PartialBodyResolveProvider.NONE);
+                    ArgumentTypeResolver.updateNumberType(defaultType, argumentExpression, context);
                 }
                 if (constant != null) {
                     constants.add(constant);

compiler/frontend/src/org/jetbrains/kotlin/resolve/PartialBodyResolveProvider.kt

@@ -17,8 +17,11 @@
 package org.jetbrains.kotlin.resolve
 
 import org.jetbrains.kotlin.psi.JetElement
+import org.jetbrains.kotlin.psi.JetBlockExpression
+import org.jetbrains.kotlin.psi.JetPsiUtil
 
 public open class PartialBodyResolveProvider {
+
     public open val filter: ((JetElement) -> Boolean)?
         get() = null
 
@@ -26,3 +29,10 @@ public open class PartialBodyResolveProvider {
         public val NONE: PartialBodyResolveProvider = PartialBodyResolveProvider()
     }
 }
+
+fun PartialBodyResolveProvider.filterStatements(block: JetBlockExpression): List<JetElement> {
+    if (filter == null || block is JetPsiUtil.JetExpressionWrapper) return block.getStatements()
+    return block.getStatements().filter { filter!!(it) }
+}
+
+fun PartialBodyResolveProvider.getLastStatementInABlock(block: JetBlockExpression) = filterStatements(block).lastOrNull()

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/ArgumentTypeResolver.java

@@ -20,7 +20,6 @@ import com.google.common.collect.Lists;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
-import org.jetbrains.kotlin.descriptors.CallableDescriptor;
 import org.jetbrains.kotlin.descriptors.annotations.Annotations;
 import org.jetbrains.kotlin.diagnostics.Errors;
 import org.jetbrains.kotlin.psi.*;
@@ -119,7 +118,7 @@ public class ArgumentTypeResolver {
 
         for (ValueArgument valueArgument : context.call.getValueArguments()) {
             JetExpression argumentExpression = valueArgument.getArgumentExpression();
-            if (argumentExpression != null && isFunctionLiteralArgument(argumentExpression)) {
+            if (argumentExpression != null && isFunctionLiteralArgument(argumentExpression, context)) {
                 checkArgumentTypeWithNoCallee(context, argumentExpression);
             }
         }
@@ -130,20 +129,26 @@ public class ArgumentTypeResolver {
         updateResultArgumentTypeIfNotDenotable(context, argumentExpression);
     }
 
-    public static boolean isFunctionLiteralArgument(@NotNull JetExpression expression) {
+    public static boolean isFunctionLiteralArgument(
-        return getFunctionLiteralArgumentIfAny(expression) != null;
+            @NotNull JetExpression expression, @NotNull ResolutionContext context
+    ) {
+        return getFunctionLiteralArgumentIfAny(expression, context) != null;
     }
 
     @NotNull
-    public static JetFunctionLiteralExpression getFunctionLiteralArgument(@NotNull JetExpression expression) {
+    public static JetFunctionLiteralExpression getFunctionLiteralArgument(
-        assert isFunctionLiteralArgument(expression);
+            @NotNull JetExpression expression, @NotNull ResolutionContext context
+    ) {
+        assert isFunctionLiteralArgument(expression, context);
         //noinspection ConstantConditions
-        return getFunctionLiteralArgumentIfAny(expression);
+        return getFunctionLiteralArgumentIfAny(expression, context);
     }
 
     @Nullable
-    private static JetFunctionLiteralExpression getFunctionLiteralArgumentIfAny(@NotNull JetExpression expression) {
+    private static JetFunctionLiteralExpression getFunctionLiteralArgumentIfAny(
-        JetExpression deparenthesizedExpression = deparenthesizeArgument(expression);
+            @NotNull JetExpression expression, @NotNull ResolutionContext context
+    ) {
+        JetExpression deparenthesizedExpression = deparenthesizeArgument(expression, context);
         if (deparenthesizedExpression instanceof JetFunctionLiteralExpression) {
             return (JetFunctionLiteralExpression) deparenthesizedExpression;
         }
@@ -151,16 +156,20 @@ public class ArgumentTypeResolver {
     }
 
     @Nullable
-    public static JetExpression deparenthesizeArgument(@Nullable JetExpression expression) {
+    public static JetExpression deparenthesizeArgument(
+            @Nullable JetExpression expression,
+            @NotNull ResolutionContext context
+    ) {
         JetExpression deparenthesizedExpression = JetPsiUtil.deparenthesize(expression, false);
         if (deparenthesizedExpression instanceof JetBlockExpression) {
+            JetBlockExpression blockExpression = (JetBlockExpression) deparenthesizedExpression;
             // todo
             // This case is a temporary hack for 'if' branches.
             // The right way to implement this logic is to interpret 'if' branches as function literals with explicitly-typed signatures
             // (no arguments and no receiver) and therefore analyze them straight away (not in the 'complete' phase).
-            JetElement lastStatementInABlock = JetPsiUtil.getLastStatementInABlock((JetBlockExpression) deparenthesizedExpression);
+            JetElement lastStatementInABlock = ResolvePackage.getLastStatementInABlock(context.partialBodyResolveProvider, blockExpression);
             if (lastStatementInABlock instanceof JetExpression) {
-                return deparenthesizeArgument((JetExpression) lastStatementInABlock);
+                return deparenthesizeArgument((JetExpression) lastStatementInABlock, context);
             }
         }
         return deparenthesizedExpression;
@@ -175,8 +184,8 @@ public class ArgumentTypeResolver {
         if (expression == null) {
             return JetTypeInfo.create(null, context.dataFlowInfo);
         }
-        if (isFunctionLiteralArgument(expression)) {
+        if (isFunctionLiteralArgument(expression, context)) {
-            return getFunctionLiteralTypeInfo(expression, getFunctionLiteralArgument(expression), context, resolveArgumentsMode);
+            return getFunctionLiteralTypeInfo(expression, getFunctionLiteralArgument(expression, context), context, resolveArgumentsMode);
         }
         JetTypeInfo recordedTypeInfo = getRecordedTypeInfo(expression, context.trace.getBindingContext());
         if (recordedTypeInfo != null) {
@@ -240,7 +249,7 @@ public class ArgumentTypeResolver {
         return defaultValue;
     }
 
-    public <D extends CallableDescriptor> void analyzeArgumentsAndRecordTypes(
+    public void analyzeArgumentsAndRecordTypes(
             @NotNull CallResolutionContext<?> context
     ) {
         MutableDataFlowInfoForArguments infoForArguments = context.dataFlowInfoForArguments;
@@ -266,7 +275,7 @@ public class ArgumentTypeResolver {
             if (type.getConstructor() instanceof IntegerValueTypeConstructor) {
                 IntegerValueTypeConstructor constructor = (IntegerValueTypeConstructor) type.getConstructor();
                 JetType primitiveType = TypeUtils.getPrimitiveNumberType(constructor, context.expectedType);
-                updateNumberType(primitiveType, expression, context.trace);
+                updateNumberType(primitiveType, expression, context);
                 return primitiveType;
             }
         }
@@ -276,19 +285,19 @@ public class ArgumentTypeResolver {
     public static void updateNumberType(
             @NotNull JetType numberType,
             @Nullable JetExpression expression,
-            @NotNull BindingTrace trace
+            @NotNull ResolutionContext context
     ) {
         if (expression == null) return;
-        BindingContextUtils.updateRecordedType(numberType, expression, trace, false);
+        BindingContextUtils.updateRecordedType(numberType, expression, context.trace, false);
 
         if (!(expression instanceof JetConstantExpression)) {
-            JetExpression deparenthesized = deparenthesizeArgument(expression);
+            JetExpression deparenthesized = deparenthesizeArgument(expression, context);
             if (deparenthesized != expression) {
-                updateNumberType(numberType, deparenthesized, trace);
+                updateNumberType(numberType, deparenthesized, context);
             }
             return;
         }
 
-        ConstantExpressionEvaluator.OBJECT$.evaluate(expression, trace, numberType);
+        ConstantExpressionEvaluator.OBJECT$.evaluate(expression, context.trace, numberType);
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/CallCompleter.kt

@@ -226,7 +226,7 @@ public class CallCompleter(
         if (valueArgument.isExternal()) return
 
         val expression = valueArgument.getArgumentExpression()
-        val deparenthesized = ArgumentTypeResolver.deparenthesizeArgument(expression)
+        val deparenthesized = ArgumentTypeResolver.deparenthesizeArgument(expression, context as ResolutionContext<*>)
         if (deparenthesized == null) return
 
         val recordedType = context.trace[BindingContext.EXPRESSION_TYPE, expression]
@@ -248,9 +248,9 @@ public class CallCompleter(
 
         // While the expected type is not known, the function literal arguments are not analyzed (to analyze function literal bodies once),
         // but they should be analyzed when the expected type is known (during the call completion).
-        if (ArgumentTypeResolver.isFunctionLiteralArgument(expression)) {
+        if (ArgumentTypeResolver.isFunctionLiteralArgument(expression, context as ResolutionContext<*>)) {
             argumentTypeResolver.getFunctionLiteralTypeInfo(
-                    expression, ArgumentTypeResolver.getFunctionLiteralArgument(expression),
+                    expression, ArgumentTypeResolver.getFunctionLiteralArgument(expression, context as ResolutionContext<*>),
                     context as CallResolutionContext<*>, RESOLVE_FUNCTION_ARGUMENTS)
         }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/CandidateResolver.java

@@ -279,9 +279,9 @@ public class CandidateResolver {
     ) {
         JetExpression argumentExpression = valueArgument.getArgumentExpression();
         if (argumentExpression == null) return;
-        if (!ArgumentTypeResolver.isFunctionLiteralArgument(argumentExpression)) return;
+        if (!ArgumentTypeResolver.isFunctionLiteralArgument(argumentExpression, context)) return;
 
-        JetFunctionLiteralExpression functionLiteralExpression = ArgumentTypeResolver.getFunctionLiteralArgument(argumentExpression);
+        JetFunctionLiteralExpression functionLiteralExpression = ArgumentTypeResolver.getFunctionLiteralArgument(argumentExpression, context);
 
         JetType effectiveExpectedType = getEffectiveExpectedType(valueParameterDescriptor, valueArgument);
         JetType expectedType = constraintSystem.getCurrentSubstitutor().substitute(effectiveExpectedType, Variance.INVARIANT);
@@ -379,8 +379,8 @@ public class CandidateResolver {
                     ? TypeUtils.makeNotNullable(receiverArgument.getType())
                     : receiverArgument.getType();
             if (receiverArgument instanceof ExpressionReceiver) {
-                receiverType = updateResultTypeForSmartCasts(receiverType, ((ExpressionReceiver) receiverArgument).getExpression(),
+                receiverType = updateResultTypeForSmartCasts(
-                                                             context.dataFlowInfo, context.trace);
+                        receiverType, ((ExpressionReceiver) receiverArgument).getExpression(), context);
             }
             constraintSystem.addSubtypeConstraint(receiverType, receiverParameter.getType(), RECEIVER_POSITION.position());
         }
@@ -424,7 +424,8 @@ public class CandidateResolver {
                 argumentExpression, newContext, resolveFunctionArgumentBodies);
         context.candidateCall.getDataFlowInfoForArguments().updateInfo(valueArgument, typeInfoForCall.getDataFlowInfo());
 
-        JetType type = updateResultTypeForSmartCasts(typeInfoForCall.getType(), argumentExpression, dataFlowInfoForArgument, context.trace);
+        JetType type = updateResultTypeForSmartCasts(
+                typeInfoForCall.getType(), argumentExpression, context.replaceDataFlowInfo(dataFlowInfoForArgument));
         constraintSystem.addSubtypeConstraint(
                 type, effectiveExpectedType, VALUE_PARAMETER_POSITION.position(valueParameterDescriptor.getIndex()));
     }
@@ -433,16 +434,15 @@ public class CandidateResolver {
     private static JetType updateResultTypeForSmartCasts(
             @Nullable JetType type,
             @Nullable JetExpression argumentExpression,
-            @NotNull DataFlowInfo dataFlowInfoForArgument,
+            @NotNull ResolutionContext context
-            @NotNull BindingTrace trace
     ) {
-        JetExpression deparenthesizedArgument = deparenthesizeArgument(argumentExpression);
+        JetExpression deparenthesizedArgument = deparenthesizeArgument(argumentExpression, context);
         if (deparenthesizedArgument == null || type == null) return type;
 
-        DataFlowValue dataFlowValue = DataFlowValueFactory.createDataFlowValue(deparenthesizedArgument, type, trace.getBindingContext());
+        DataFlowValue dataFlowValue = DataFlowValueFactory.createDataFlowValue(deparenthesizedArgument, type, context.trace.getBindingContext());
         if (!dataFlowValue.isStableIdentifier()) return type;
 
-        Set<JetType> possibleTypes = dataFlowInfoForArgument.getPossibleTypes(dataFlowValue);
+        Set<JetType> possibleTypes = context.dataFlowInfo.getPossibleTypes(dataFlowValue);
         if (possibleTypes.isEmpty()) return type;
 
         return TypeUtils.intersect(JetTypeChecker.DEFAULT, possibleTypes);

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

@@ -20,6 +20,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.psi.Call;
 import org.jetbrains.kotlin.resolve.BindingTrace;
+import org.jetbrains.kotlin.resolve.PartialBodyResolveProvider;
 import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker;
 import org.jetbrains.kotlin.resolve.calls.model.MutableDataFlowInfoForArguments;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
@@ -38,11 +39,12 @@ public class BasicCallResolutionContext extends CallResolutionContext<BasicCallR
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @Nullable MutableDataFlowInfoForArguments dataFlowInfoForArguments,
             @NotNull CallChecker callChecker,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean isAnnotationContext,
             boolean collectAllCandidates
     ) {
         super(trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments, resolutionResultsCache,
-              dataFlowInfoForArguments, callChecker, isAnnotationContext, collectAllCandidates);
+              dataFlowInfoForArguments, callChecker, partialBodyResolveProvider, isAnnotationContext, collectAllCandidates);
     }
 
     @NotNull
@@ -59,7 +61,7 @@ public class BasicCallResolutionContext extends CallResolutionContext<BasicCallR
     ) {
         return new BasicCallResolutionContext(trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments,
              new ResolutionResultsCacheImpl(), null,
-             callChecker, isAnnotationContext, false);
+             callChecker, PartialBodyResolveProvider.NONE, isAnnotationContext, false);
     }
 
     @NotNull
@@ -67,9 +69,10 @@ public class BasicCallResolutionContext extends CallResolutionContext<BasicCallR
             @NotNull ResolutionContext context, @NotNull Call call, @NotNull CheckValueArgumentsMode checkArguments,
             @Nullable MutableDataFlowInfoForArguments dataFlowInfoForArguments
     ) {
-        return new BasicCallResolutionContext(context.trace, context.scope, call, context.expectedType, context.dataFlowInfo, context.contextDependency, checkArguments,
+        return new BasicCallResolutionContext(
-             context.resolutionResultsCache, dataFlowInfoForArguments, context.callChecker,
+                context.trace, context.scope, call, context.expectedType, context.dataFlowInfo, context.contextDependency, checkArguments,
-             context.isAnnotationContext, context.collectAllCandidates);
+                context.resolutionResultsCache, dataFlowInfoForArguments, context.callChecker, context.partialBodyResolveProvider,
+                context.isAnnotationContext, context.collectAllCandidates);
     }
 
     @NotNull
@@ -87,17 +90,18 @@ public class BasicCallResolutionContext extends CallResolutionContext<BasicCallR
             @NotNull JetType expectedType,
             @NotNull ContextDependency contextDependency,
             @NotNull ResolutionResultsCache resolutionResultsCache,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean collectAllCandidates
     ) {
         return new BasicCallResolutionContext(
                 trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments, resolutionResultsCache,
-                dataFlowInfoForArguments, callChecker, isAnnotationContext, collectAllCandidates);
+                dataFlowInfoForArguments, callChecker, partialBodyResolveProvider, isAnnotationContext, collectAllCandidates);
     }
 
     @NotNull
     public BasicCallResolutionContext replaceCall(@NotNull Call newCall) {
         return new BasicCallResolutionContext(
                 trace, scope, newCall, expectedType, dataFlowInfo, contextDependency, checkArguments, resolutionResultsCache,
-                dataFlowInfoForArguments, callChecker, isAnnotationContext, collectAllCandidates);
+                dataFlowInfoForArguments, callChecker, partialBodyResolveProvider, isAnnotationContext, collectAllCandidates);
     }
 }

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

@@ -21,6 +21,7 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.CallableDescriptor;
 import org.jetbrains.kotlin.psi.Call;
 import org.jetbrains.kotlin.resolve.BindingTrace;
+import org.jetbrains.kotlin.resolve.PartialBodyResolveProvider;
 import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker;
 import org.jetbrains.kotlin.resolve.calls.model.MutableDataFlowInfoForArguments;
 import org.jetbrains.kotlin.resolve.calls.model.MutableResolvedCall;
@@ -51,12 +52,13 @@ public final class CallCandidateResolutionContext<D extends CallableDescriptor>
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @Nullable MutableDataFlowInfoForArguments dataFlowInfoForArguments,
             @NotNull CallChecker callChecker,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             @NotNull ReceiverValue explicitExtensionReceiverForInvoke,
             boolean isAnnotationContext,
             boolean collectAllCandidates
     ) {
         super(trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments, resolutionResultsCache,
-              dataFlowInfoForArguments, callChecker, isAnnotationContext, collectAllCandidates);
+              dataFlowInfoForArguments, callChecker, partialBodyResolveProvider, isAnnotationContext, collectAllCandidates);
         this.candidateCall = candidateCall;
         this.tracing = tracing;
         this.explicitExtensionReceiverForInvoke = explicitExtensionReceiverForInvoke;
@@ -71,7 +73,8 @@ public final class CallCandidateResolutionContext<D extends CallableDescriptor>
                 candidateCall, tracing, trace, context.scope, call, context.expectedType,
                 context.dataFlowInfo, context.contextDependency, context.checkArguments,
                 context.resolutionResultsCache, context.dataFlowInfoForArguments,
-                context.callChecker, explicitExtensionReceiverForInvoke, context.isAnnotationContext, context.collectAllCandidates);
+                context.callChecker, context.partialBodyResolveProvider, explicitExtensionReceiverForInvoke,
+                context.isAnnotationContext, context.collectAllCandidates);
     }
 
     public static <D extends CallableDescriptor> CallCandidateResolutionContext<D> create(
@@ -94,7 +97,7 @@ public final class CallCandidateResolutionContext<D extends CallableDescriptor>
         return new CallCandidateResolutionContext<D>(
                 candidateCall, tracing, context.trace, context.scope, context.call, context.expectedType,
                 context.dataFlowInfo, context.contextDependency, context.checkArguments, context.resolutionResultsCache,
-                context.dataFlowInfoForArguments, context.callChecker, ReceiverValue.NO_RECEIVER,
+                context.dataFlowInfoForArguments, context.callChecker, context.partialBodyResolveProvider, ReceiverValue.NO_RECEIVER,
                 context.isAnnotationContext, context.collectAllCandidates);
     }
 
@@ -106,11 +109,12 @@ public final class CallCandidateResolutionContext<D extends CallableDescriptor>
             @NotNull JetType expectedType,
             @NotNull ContextDependency contextDependency,
             @NotNull ResolutionResultsCache resolutionResultsCache,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean collectAllCandidates
     ) {
         return new CallCandidateResolutionContext<D>(
                 candidateCall, tracing, trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments,
-                resolutionResultsCache, dataFlowInfoForArguments, callChecker, explicitExtensionReceiverForInvoke,
+                resolutionResultsCache, dataFlowInfoForArguments, callChecker, partialBodyResolveProvider,
-                isAnnotationContext, collectAllCandidates);
+                explicitExtensionReceiverForInvoke, isAnnotationContext, collectAllCandidates);
     }
 }

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

@@ -20,6 +20,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.psi.Call;
 import org.jetbrains.kotlin.resolve.BindingTrace;
+import org.jetbrains.kotlin.resolve.PartialBodyResolveProvider;
 import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker;
 import org.jetbrains.kotlin.resolve.calls.model.DataFlowInfoForArgumentsImpl;
 import org.jetbrains.kotlin.resolve.calls.model.MutableDataFlowInfoForArguments;
@@ -47,11 +48,12 @@ public abstract class CallResolutionContext<Context extends CallResolutionContex
             @SuppressWarnings("NullableProblems")
             @Nullable MutableDataFlowInfoForArguments dataFlowInfoForArguments,
             @NotNull CallChecker callChecker,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean isAnnotationContext,
             boolean collectAllCandidates
     ) {
         super(trace, scope, expectedType, dataFlowInfo, contextDependency, resolutionResultsCache, callChecker,
-              isAnnotationContext, collectAllCandidates);
+              partialBodyResolveProvider, isAnnotationContext, collectAllCandidates);
         this.call = call;
         this.checkArguments = checkArguments;
         if (dataFlowInfoForArguments != null) {

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

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.resolve.calls.context;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.resolve.BindingTrace;
+import org.jetbrains.kotlin.resolve.PartialBodyResolveProvider;
 import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
 import org.jetbrains.kotlin.resolve.scopes.JetScope;
@@ -40,6 +41,8 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
     public final ResolutionResultsCache resolutionResultsCache;
     @NotNull
     public final CallChecker callChecker;
+    @NotNull
+    public final PartialBodyResolveProvider partialBodyResolveProvider;
 
     public final boolean isAnnotationContext;
 
@@ -53,6 +56,7 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
             @NotNull ContextDependency contextDependency,
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @NotNull CallChecker callChecker,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean isAnnotationContext,
             boolean collectAllCandidates
     ) {
@@ -63,6 +67,7 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
         this.contextDependency = contextDependency;
         this.resolutionResultsCache = resolutionResultsCache;
         this.callChecker = callChecker;
+        this.partialBodyResolveProvider = partialBodyResolveProvider;
         this.isAnnotationContext = isAnnotationContext;
         this.collectAllCandidates = collectAllCandidates;
     }
@@ -74,6 +79,7 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
             @NotNull JetType expectedType,
             @NotNull ContextDependency contextDependency,
             @NotNull ResolutionResultsCache resolutionResultsCache,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean collectAllCandidates
     );
 
@@ -86,38 +92,44 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
     @NotNull
     public Context replaceBindingTrace(@NotNull BindingTrace trace) {
         if (this.trace == trace) return self();
-        return create(trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, collectAllCandidates);
+        return create(trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, partialBodyResolveProvider,
+                      collectAllCandidates);
     }
 
     @NotNull
     public Context replaceDataFlowInfo(@NotNull DataFlowInfo newDataFlowInfo) {
         if (newDataFlowInfo == dataFlowInfo) return self();
-        return create(trace, scope, newDataFlowInfo, expectedType, contextDependency, resolutionResultsCache, collectAllCandidates);
+        return create(trace, scope, newDataFlowInfo, expectedType, contextDependency, resolutionResultsCache, partialBodyResolveProvider,
+                      collectAllCandidates);
     }
 
     @NotNull
     public Context replaceExpectedType(@Nullable JetType newExpectedType) {
         if (newExpectedType == null) return replaceExpectedType(TypeUtils.NO_EXPECTED_TYPE);
         if (expectedType == newExpectedType) return self();
-        return create(trace, scope, dataFlowInfo, newExpectedType, contextDependency, resolutionResultsCache, collectAllCandidates);
+        return create(trace, scope, dataFlowInfo, newExpectedType, contextDependency, resolutionResultsCache, partialBodyResolveProvider,
+                      collectAllCandidates);
     }
 
     @NotNull
     public Context replaceScope(@NotNull JetScope newScope) {
         if (newScope == scope) return self();
-        return create(trace, newScope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, collectAllCandidates);
+        return create(trace, newScope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, partialBodyResolveProvider,
+                      collectAllCandidates);
     }
 
     @NotNull
     public Context replaceContextDependency(@NotNull ContextDependency newContextDependency) {
         if (newContextDependency == contextDependency) return self();
-        return create(trace, scope, dataFlowInfo, expectedType, newContextDependency, resolutionResultsCache, collectAllCandidates);
+        return create(trace, scope, dataFlowInfo, expectedType, newContextDependency, resolutionResultsCache, partialBodyResolveProvider,
+                      collectAllCandidates);
     }
 
     @NotNull
     public Context replaceResolutionResultsCache(@NotNull ResolutionResultsCache newResolutionResultsCache) {
         if (newResolutionResultsCache == resolutionResultsCache) return self();
-        return create(trace, scope, dataFlowInfo, expectedType, contextDependency, newResolutionResultsCache, collectAllCandidates);
+        return create(trace, scope, dataFlowInfo, expectedType, contextDependency, newResolutionResultsCache, partialBodyResolveProvider,
+                      collectAllCandidates);
     }
 
     @NotNull
@@ -127,6 +139,13 @@ public abstract class ResolutionContext<Context extends ResolutionContext<Contex
 
     @NotNull
     public Context replaceCollectAllCandidates(boolean newCollectAllCandidates) {
-        return create(trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, newCollectAllCandidates);
+        return create(trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, partialBodyResolveProvider,
+                      newCollectAllCandidates);
+    }
+
+    @NotNull
+    public Context replacePartialBodyResolveProvider(@NotNull PartialBodyResolveProvider partialBodyResolveProvider) {
+        return create(trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, partialBodyResolveProvider,
+                      collectAllCandidates);
     }
 }

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

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.resolve.calls.context;
 
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.resolve.BindingTrace;
+import org.jetbrains.kotlin.resolve.PartialBodyResolveProvider;
 import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
 import org.jetbrains.kotlin.resolve.scopes.JetScope;
@@ -32,10 +33,11 @@ public class SimpleResolutionContext extends ResolutionContext<SimpleResolutionC
             @NotNull ContextDependency contextDependency,
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @NotNull CallChecker callChecker,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean isAnnotationContext,
             boolean collectAllCandidates
     ) {
-        super(trace, scope, expectedType, dataFlowInfo, contextDependency, resolutionResultsCache, callChecker,
+        super(trace, scope, expectedType, dataFlowInfo, contextDependency, resolutionResultsCache, callChecker, partialBodyResolveProvider,
               isAnnotationContext, collectAllCandidates);
     }
 
@@ -45,10 +47,11 @@ public class SimpleResolutionContext extends ResolutionContext<SimpleResolutionC
             @NotNull JetType expectedType,
             @NotNull DataFlowInfo dataFlowInfo,
             @NotNull ContextDependency contextDependency,
-            @NotNull CallChecker callChecker
+            @NotNull CallChecker callChecker,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider
     ) {
         this(trace, scope, expectedType, dataFlowInfo, contextDependency, new ResolutionResultsCacheImpl(),
-             callChecker, false, false);
+             callChecker, partialBodyResolveProvider, false, false);
     }
 
     @Override
@@ -59,10 +62,11 @@ public class SimpleResolutionContext extends ResolutionContext<SimpleResolutionC
             @NotNull JetType expectedType,
             @NotNull ContextDependency contextDependency,
             @NotNull ResolutionResultsCache resolutionResultsCache,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean collectAllCandidates
     ) {
         return new SimpleResolutionContext(
-                trace, scope, expectedType, dataFlowInfo, contextDependency, resolutionResultsCache, callChecker,
+                trace, scope, expectedType, dataFlowInfo, contextDependency, resolutionResultsCache, callChecker, partialBodyResolveProvider,
                 isAnnotationContext, collectAllCandidates);
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/results/ResolutionResultsHandler.java

@@ -102,7 +102,7 @@ public class ResolutionResultsHandler {
             // This check is needed for the following case:
             //    x.foo(unresolved) -- if there are multiple foo's, we'd report an ambiguity, and it does not make sense here
             if (task.checkArguments == CheckValueArgumentsMode.DISABLED ||
-                    !CallUtilPackage.hasUnresolvedArguments(task.call, task.trace.getBindingContext())) {
+                    !CallUtilPackage.hasUnresolvedArguments(task.call, task)) {
                 if (allCandidatesIncomplete) {
                     task.tracing.cannotCompleteResolve(task.trace, results.getResultingCalls());
                 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tasks/ResolutionTask.java

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor;
 import org.jetbrains.kotlin.psi.Call;
 import org.jetbrains.kotlin.psi.JetReferenceExpression;
 import org.jetbrains.kotlin.resolve.BindingTrace;
+import org.jetbrains.kotlin.resolve.PartialBodyResolveProvider;
 import org.jetbrains.kotlin.resolve.calls.context.*;
 import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker;
 import org.jetbrains.kotlin.resolve.calls.model.MutableDataFlowInfoForArguments;
@@ -55,12 +56,13 @@ public class ResolutionTask<D extends CallableDescriptor, F extends D> extends C
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @Nullable MutableDataFlowInfoForArguments dataFlowInfoForArguments,
             @NotNull CallChecker callChecker,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             @NotNull Collection<MutableResolvedCall<F>> resolvedCalls,
             boolean isAnnotationContext,
             boolean collectAllCandidates
     ) {
         super(trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments, resolutionResultsCache,
-              dataFlowInfoForArguments, callChecker, isAnnotationContext, collectAllCandidates);
+              dataFlowInfoForArguments, callChecker, partialBodyResolveProvider, isAnnotationContext, collectAllCandidates);
         this.lazyCandidates = lazyCandidates;
         this.resolvedCalls = resolvedCalls;
         this.tracing = tracing;
@@ -75,7 +77,8 @@ public class ResolutionTask<D extends CallableDescriptor, F extends D> extends C
              context.trace, context.scope, context.call,
              context.expectedType, context.dataFlowInfo, context.contextDependency, context.checkArguments,
              context.resolutionResultsCache, context.dataFlowInfoForArguments,
-             context.callChecker, Lists.<MutableResolvedCall<F>>newArrayList(), context.isAnnotationContext, context.collectAllCandidates);
+             context.callChecker, context.partialBodyResolveProvider, Lists.<MutableResolvedCall<F>>newArrayList(),
+             context.isAnnotationContext, context.collectAllCandidates);
     }
 
     public ResolutionTask(
@@ -113,12 +116,13 @@ public class ResolutionTask<D extends CallableDescriptor, F extends D> extends C
             @NotNull JetType expectedType,
             @NotNull ContextDependency contextDependency,
             @NotNull ResolutionResultsCache resolutionResultsCache,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean collectAllCandidates
     ) {
         return new ResolutionTask<D, F>(
                 lazyCandidates, tracing, trace, scope, call, expectedType, dataFlowInfo, contextDependency, checkArguments,
-                resolutionResultsCache, dataFlowInfoForArguments, callChecker, resolvedCalls, isAnnotationContext,
+                resolutionResultsCache, dataFlowInfoForArguments, callChecker, partialBodyResolveProvider, resolvedCalls,
-                collectAllCandidates);
+                isAnnotationContext, collectAllCandidates);
     }
 
     public ResolutionTask<D, F> replaceContext(@NotNull BasicCallResolutionContext newContext) {
@@ -128,7 +132,7 @@ public class ResolutionTask<D extends CallableDescriptor, F extends D> extends C
     public ResolutionTask<D, F> replaceCall(@NotNull Call newCall) {
         return new ResolutionTask<D, F>(
                 lazyCandidates, tracing, trace, scope, newCall, expectedType, dataFlowInfo, contextDependency, checkArguments,
-                resolutionResultsCache, dataFlowInfoForArguments, callChecker, resolvedCalls,
+                resolutionResultsCache, dataFlowInfoForArguments, callChecker, partialBodyResolveProvider, resolvedCalls,
                 isAnnotationContext, collectAllCandidates);
     }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/util/callUtil.kt

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.resolve.BindingContext.RESOLVED_CALL
 import org.jetbrains.kotlin.utils.sure
 import org.jetbrains.kotlin.psi.psiUtil.getTextWithLocation
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
 
 // resolved call
 
@@ -63,12 +64,12 @@ public fun <D : CallableDescriptor> ResolvedCall<D>.getParameterForArgument(valu
 
 // call
 
-public fun Call.hasUnresolvedArguments(context: BindingContext): Boolean {
+public fun <C: ResolutionContext<C>> Call.hasUnresolvedArguments(context: ResolutionContext<C>): Boolean {
     val arguments = getValueArguments().map { it?.getArgumentExpression() }
     return arguments.any {
         argument ->
-        val expressionType = context[BindingContext.EXPRESSION_TYPE, argument]
+        val expressionType = context.trace.getBindingContext()[BindingContext.EXPRESSION_TYPE, argument]
-        argument != null && !ArgumentTypeResolver.isFunctionLiteralArgument(argument)
+        argument != null && !ArgumentTypeResolver.isFunctionLiteralArgument(argument, context as ResolutionContext<*>)
                 && (expressionType == null || expressionType.isError())
     }
 }

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

@@ -763,7 +763,7 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
         JetType expressionType = value.getType(KotlinBuiltIns.getInstance());
         if (value instanceof IntegerValueTypeConstant && context.contextDependency == INDEPENDENT) {
             expressionType = ((IntegerValueTypeConstant) value).getType(context.expectedType);
-            ArgumentTypeResolver.updateNumberType(expressionType, expression, context.trace);
+            ArgumentTypeResolver.updateNumberType(expressionType, expression, context);
         }
 
         return DataFlowUtils.checkType(expressionType, expression, context, context.dataFlowInfo);

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

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.types.expressions;
 
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.resolve.BindingTrace;
+import org.jetbrains.kotlin.resolve.PartialBodyResolveProvider;
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency;
 import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext;
 import org.jetbrains.kotlin.resolve.calls.context.ResolutionResultsCache;
@@ -40,15 +41,15 @@ public class ExpressionTypingContext extends ResolutionContext<ExpressionTypingC
     ) {
         return newContext(trace, scope, dataFlowInfo, expectedType,
                           ContextDependency.INDEPENDENT, new ResolutionResultsCacheImpl(),
-                          expressionTypingServices.getCallChecker(), false);
+                          expressionTypingServices.getCallChecker(), PartialBodyResolveProvider.NONE, false);
     }
 
     @NotNull
     public static ExpressionTypingContext newContext(@NotNull ResolutionContext context) {
         return new ExpressionTypingContext(
                 context.trace, context.scope, context.dataFlowInfo, context.expectedType,
-                context.contextDependency, context.resolutionResultsCache, context.callChecker, context.isAnnotationContext,
+                context.contextDependency, context.resolutionResultsCache, context.callChecker, context.partialBodyResolveProvider,
-                context.collectAllCandidates
+                context.isAnnotationContext, context.collectAllCandidates
         );
     }
 
@@ -61,11 +62,12 @@ public class ExpressionTypingContext extends ResolutionContext<ExpressionTypingC
             @NotNull ContextDependency contextDependency,
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @NotNull CallChecker callChecker,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean isAnnotationContext
     ) {
         return new ExpressionTypingContext(
-                trace, scope, dataFlowInfo, expectedType, contextDependency,
+                trace, scope, dataFlowInfo, expectedType, contextDependency, resolutionResultsCache, callChecker,
-                resolutionResultsCache, callChecker, isAnnotationContext, false);
+                partialBodyResolveProvider, isAnnotationContext, false);
     }
 
     private CompileTimeConstantChecker compileTimeConstantChecker;
@@ -78,11 +80,12 @@ public class ExpressionTypingContext extends ResolutionContext<ExpressionTypingC
             @NotNull ContextDependency contextDependency,
             @NotNull ResolutionResultsCache resolutionResultsCache,
             @NotNull CallChecker callChecker,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean isAnnotationContext,
             boolean collectAllCandidates
     ) {
         super(trace, scope, expectedType, dataFlowInfo, contextDependency, resolutionResultsCache, callChecker,
-              isAnnotationContext, collectAllCandidates);
+              partialBodyResolveProvider, isAnnotationContext, collectAllCandidates);
     }
 
     @Override
@@ -93,11 +96,12 @@ public class ExpressionTypingContext extends ResolutionContext<ExpressionTypingC
             @NotNull JetType expectedType,
             @NotNull ContextDependency contextDependency,
             @NotNull ResolutionResultsCache resolutionResultsCache,
+            @NotNull PartialBodyResolveProvider partialBodyResolveProvider,
             boolean collectAllCandidates
     ) {
         return new ExpressionTypingContext(trace, scope, dataFlowInfo,
                                            expectedType, contextDependency, resolutionResultsCache, callChecker,
-                                           isAnnotationContext, collectAllCandidates);
+                                           partialBodyResolveProvider, isAnnotationContext, collectAllCandidates);
     }
 
 ///////////// LAZY ACCESSORS

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

@@ -19,8 +19,6 @@ package org.jetbrains.kotlin.types.expressions;
 import com.google.common.base.Function;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.tree.IElementType;
-import kotlin.Function1;
-import kotlin.KotlinPackage;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.analyzer.AnalyzerPackage;
@@ -210,12 +208,7 @@ public class ExpressionTypingServices {
             @NotNull CoercionStrategy coercionStrategyForLastExpression,
             @NotNull ExpressionTypingContext context
     ) {
-        List<JetElement> block = expression.getStatements();
+        List<JetElement> block = ResolvePackage.filterStatements(getPartialBodyResolveProvider(), expression);
-
-        Function1<JetElement, Boolean> filter = getPartialBodyResolveProvider().getFilter();
-        if (filter != null && !(expression instanceof JetPsiUtil.JetExpressionWrapper)) {
-            block = KotlinPackage.filter(block, filter);
-        }
 
         // SCRIPT: get code descriptor for script declaration
         DeclarationDescriptor containingDescriptor = context.scope.getContainingDeclaration();
@@ -235,7 +228,8 @@ public class ExpressionTypingServices {
             r = DataFlowUtils.checkType(builtIns.getUnitType(), expression, context, context.dataFlowInfo);
         }
         else {
-            r = getBlockReturnedTypeWithWritableScope(scope, block, coercionStrategyForLastExpression, context);
+            r = getBlockReturnedTypeWithWritableScope(scope, block, coercionStrategyForLastExpression,
+                                                      context.replacePartialBodyResolveProvider(getPartialBodyResolveProvider()));
         }
         scope.changeLockLevel(WritableScope.LockLevel.READING);