resolve argument of 'if' call

with right data flow info

compiler/frontend/src/org/jetbrains/jet/lang/resolve/calls/CallResolver.java

@@ -29,9 +29,13 @@ import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.*;
 import org.jetbrains.jet.lang.resolve.calls.autocasts.DataFlowInfo;
 import org.jetbrains.jet.lang.resolve.calls.context.*;
+import org.jetbrains.jet.lang.resolve.calls.model.MutableDataFlowInfoForArguments;
 import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCallImpl;
 import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCallWithTrace;
-import org.jetbrains.jet.lang.resolve.calls.results.*;
+import org.jetbrains.jet.lang.resolve.calls.results.OverloadResolutionResults;
+import org.jetbrains.jet.lang.resolve.calls.results.OverloadResolutionResultsImpl;
+import org.jetbrains.jet.lang.resolve.calls.results.ResolutionDebugInfo;
+import org.jetbrains.jet.lang.resolve.calls.results.ResolutionResultsHandler;
 import org.jetbrains.jet.lang.resolve.calls.tasks.*;
 import org.jetbrains.jet.lang.resolve.calls.util.DelegatingCall;
 import org.jetbrains.jet.lang.resolve.calls.util.ExpressionAsFunctionDescriptor;
@@ -271,14 +275,14 @@ public class CallResolver {
             @Nullable TracingStrategy tracing,
             @NotNull JetReferenceExpression reference,
             @NotNull ResolutionContext<?> context,
-            @NotNull CallableDescriptor candidate
+            @NotNull ResolutionCandidate<CallableDescriptor> candidate,
+            @Nullable MutableDataFlowInfoForArguments dataFlowInfoForArguments
     ) {
-        Collection<ResolutionCandidate<CallableDescriptor>> resolutionCandidates = Lists.newArrayList();
+        BasicCallResolutionContext basicCallResolutionContext =
-        resolutionCandidates.add(ResolutionCandidate.create(candidate, null));
+                BasicCallResolutionContext.create(context, call, CheckValueArgumentsMode.ENABLED, dataFlowInfoForArguments);
-        BasicCallResolutionContext basicCallResolutionContext = BasicCallResolutionContext.create(context, call, CheckValueArgumentsMode.ENABLED);
 
         List<ResolutionTask<CallableDescriptor, FunctionDescriptor>> tasks = TaskPrioritizer
-                .computePrioritizedTasksFromCandidates(basicCallResolutionContext, reference, resolutionCandidates, tracing);
+                .computePrioritizedTasksFromCandidates(basicCallResolutionContext, reference, Collections.singleton(candidate), tracing);
         return doResolveCallOrGetCachedResults(ResolutionResultsCache.FUNCTION_MEMBER_TYPE, basicCallResolutionContext, tasks,
                                                CallTransformer.FUNCTION_CALL_TRANSFORMER, reference);
     }

compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/ControlStructureTypingUtils.java

@@ -28,11 +28,14 @@ import org.jetbrains.jet.lang.descriptors.impl.TypeParameterDescriptorImpl;
 import org.jetbrains.jet.lang.descriptors.impl.ValueParameterDescriptorImpl;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.BindingTrace;
+import org.jetbrains.jet.lang.resolve.calls.autocasts.DataFlowInfo;
 import org.jetbrains.jet.lang.resolve.calls.inference.ConstraintSystem;
 import org.jetbrains.jet.lang.resolve.calls.inference.InferenceErrorData;
+import org.jetbrains.jet.lang.resolve.calls.model.MutableDataFlowInfoForArguments;
 import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCall;
 import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCallWithTrace;
 import org.jetbrains.jet.lang.resolve.calls.results.OverloadResolutionResults;
+import org.jetbrains.jet.lang.resolve.calls.tasks.ResolutionCandidate;
 import org.jetbrains.jet.lang.resolve.calls.tasks.TracingStrategy;
 import org.jetbrains.jet.lang.resolve.calls.tasks.TracingStrategyImpl;
 import org.jetbrains.jet.lang.resolve.calls.util.CallMaker;
@@ -58,7 +61,9 @@ public class ControlStructureTypingUtils {
 
     /*package*/ static ResolvedCall<FunctionDescriptor> resolveIfAsCall(
             @NotNull Call callForIf,
-            @NotNull ExpressionTypingContext context
+            @NotNull ExpressionTypingContext context,
+            @NotNull DataFlowInfo thenInfo,
+            @NotNull DataFlowInfo elseInfo
     ) {
         List<AnnotationDescriptor> noAnnotations = Collections.emptyList();
         Name specialFunctionName = Name.identifierNoValidate("<SPECIAL-FUNCTION-FOR-IF-RESOLVE>");
@@ -89,14 +94,49 @@ public class ControlStructureTypingUtils {
 
         JetReferenceExpression ifReference = JetPsiFactory.createSimpleName(context.expressionTypingServices.getProject(), "fakeIfCall");
         TracingStrategy tracingForIf = createTracingForIf(callForIf);
+        MutableDataFlowInfoForArguments dataFlowInfoForArguments = createDataFlowInfoForArgumentsForCall(callForIf, thenInfo, elseInfo);
+        ResolutionCandidate<CallableDescriptor> resolutionCandidate = ResolutionCandidate.<CallableDescriptor>create(ifFunction, null);
         OverloadResolutionResults<FunctionDescriptor>
                 results = context.expressionTypingServices.getCallResolver().resolveCallWithKnownCandidate(
-                callForIf, tracingForIf, ifReference, context, ifFunction);
+                callForIf, tracingForIf, ifReference, context, resolutionCandidate, dataFlowInfoForArguments);
         assert results.isSingleResult() : "Not single result after resolving one known candidate";
         return results.getResultingCall();
     }
 
+    private static MutableDataFlowInfoForArguments createDataFlowInfoForArgumentsForCall(
+            final Call callForIf, final DataFlowInfo thenInfo, final DataFlowInfo elseInfo
+    ) {
+        return new MutableDataFlowInfoForArguments() {
+            private DataFlowInfo initialDataFlowInfo;
 
+            @Override
+            public void setInitialDataFlowInfo(@NotNull DataFlowInfo dataFlowInfo) {
+                this.initialDataFlowInfo = dataFlowInfo;
+            }
+
+            @Override
+            public void updateInfo(@NotNull ValueArgument valueArgument, @NotNull DataFlowInfo dataFlowInfo) {
+            }
+
+            @NotNull
+            @Override
+            public DataFlowInfo getInfo(@NotNull ValueArgument valueArgument) {
+                if (valueArgument == callForIf.getValueArguments().get(0)) {
+                    return thenInfo;
+                }
+                if (valueArgument == callForIf.getValueArguments().get(1)) {
+                    return elseInfo;
+                }
+                throw new IllegalArgumentException();
+            }
+
+            @NotNull
+            @Override
+            public DataFlowInfo getResultInfo() {
+                return initialDataFlowInfo;
+            }
+        };
+    }
 
     /*package*/ static Call createCallForIf(
             final JetIfExpression ifExpression,

compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/ControlStructureTypingVisitor.java

@@ -22,7 +22,10 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
-import org.jetbrains.jet.lang.descriptors.*;
+import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
+import org.jetbrains.jet.lang.descriptors.FunctionDescriptor;
+import org.jetbrains.jet.lang.descriptors.SimpleFunctionDescriptor;
+import org.jetbrains.jet.lang.descriptors.VariableDescriptor;
 import org.jetbrains.jet.lang.diagnostics.DiagnosticFactory1;
 import org.jetbrains.jet.lang.diagnostics.Errors;
 import org.jetbrains.jet.lang.psi.*;
@@ -117,15 +120,12 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
 
         if (contextWithExpectedType.expectedType == UNKNOWN_EXPECTED_TYPE) {
             Call callForIf = ControlStructureTypingUtils.createCallForIf(ifExpression, thenBranch, elseBranch);
-            ResolvedCall<FunctionDescriptor> resultingCall =
+            ControlStructureTypingUtils.resolveIfAsCall(callForIf, contextWithExpectedType, thenInfo, elseInfo);
-                    ControlStructureTypingUtils.resolveIfAsCall(callForIf, contextWithExpectedType);
 
-            List<ValueParameterDescriptor> valueParameters = resultingCall.getResultingDescriptor().getValueParameters();
+            thenTypeInfo = BindingContextUtils.getRecordedTypeInfo(thenBranch, context.trace.getBindingContext());
-            ValueParameterDescriptor thenParameter = valueParameters.get(0);
+            elseTypeInfo = BindingContextUtils.getRecordedTypeInfo(elseBranch, context.trace.getBindingContext());
-            ValueParameterDescriptor elseParameter = valueParameters.get(1);
+            assert thenTypeInfo != null : "'Then' branch of if expression  was not processed: " + ifExpression;
-            //todo correct dataFlowInfo
+            assert elseTypeInfo != null : "'Else' branch of if expression  was not processed: " + ifExpression;
-            thenTypeInfo = JetTypeInfo.create(thenParameter.getType(), context.dataFlowInfo);
-            elseTypeInfo = JetTypeInfo.create(elseParameter.getType(), context.dataFlowInfo);
         }
         else {
             CoercionStrategy coercionStrategy = isStatement ? CoercionStrategy.COERCION_TO_UNIT : CoercionStrategy.NO_COERCION;

compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/DataFlowUtils.java

@@ -22,6 +22,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.resolve.BindingTrace;
 import org.jetbrains.jet.lang.resolve.calls.context.ResolutionContext;
 import org.jetbrains.jet.lang.resolve.calls.autocasts.DataFlowInfo;
 import org.jetbrains.jet.lang.resolve.calls.autocasts.DataFlowValue;
@@ -138,28 +139,35 @@ public class DataFlowUtils {
 
     @Nullable
     public static JetType checkType(@Nullable JetType expressionType, @NotNull JetExpression expression, @NotNull ResolutionContext context) {
-        if (!noExpectedType(context.expectedType)) {
+        return checkType(expressionType, expression, context.expectedType, context.dataFlowInfo, context.trace);
-            context.trace.record(BindingContext.EXPECTED_EXPRESSION_TYPE, expression, context.expectedType);
+    }
+
+    @Nullable
+    public static JetType checkType(@Nullable JetType expressionType, @NotNull JetExpression expression,
+            @NotNull JetType expectedType, @NotNull DataFlowInfo dataFlowInfo, @NotNull BindingTrace trace
+    ) {
+        if (!noExpectedType(expectedType)) {
+            trace.record(BindingContext.EXPECTED_EXPRESSION_TYPE, expression, expectedType);
         }
 
-        if (expressionType == null || noExpectedType(context.expectedType) ||
+        if (expressionType == null || noExpectedType(expectedType) ||
-            JetTypeChecker.INSTANCE.isSubtypeOf(expressionType, context.expectedType)) {
+            JetTypeChecker.INSTANCE.isSubtypeOf(expressionType, expectedType)) {
             return expressionType;
         }
 
-        DataFlowValue dataFlowValue = DataFlowValueFactory.INSTANCE.createDataFlowValue(expression, expressionType, context.trace.getBindingContext());
+        DataFlowValue dataFlowValue = DataFlowValueFactory.INSTANCE.createDataFlowValue(expression, expressionType, trace.getBindingContext());
-        for (JetType possibleType : context.dataFlowInfo.getPossibleTypes(dataFlowValue)) {
+        for (JetType possibleType : dataFlowInfo.getPossibleTypes(dataFlowValue)) {
-            if (JetTypeChecker.INSTANCE.isSubtypeOf(possibleType, context.expectedType)) {
+            if (JetTypeChecker.INSTANCE.isSubtypeOf(possibleType, expectedType)) {
                 if (dataFlowValue.isStableIdentifier()) {
-                    context.trace.record(AUTOCAST, expression, possibleType);
+                    trace.record(AUTOCAST, expression, possibleType);
                 }
                 else {
-                    context.trace.report(AUTOCAST_IMPOSSIBLE.on(expression, possibleType, expression.getText()));
+                    trace.report(AUTOCAST_IMPOSSIBLE.on(expression, possibleType, expression.getText()));
                 }
                 return possibleType;
             }
         }
-        context.trace.report(TYPE_MISMATCH.on(expression, context.expectedType, expressionType));
+        trace.report(TYPE_MISMATCH.on(expression, expectedType, expressionType));
         return expressionType;
     }