wrap if branches in blocks

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiFactory.java

@@ -31,6 +31,7 @@ import org.jetbrains.jet.lexer.JetKeywordToken;
 import org.jetbrains.jet.lexer.JetTokens;
 import org.jetbrains.jet.plugin.JetFileType;
 
+import java.util.Collections;
 import java.util.List;
 
 public class JetPsiFactory {
@@ -484,4 +485,20 @@ public class JetPsiFactory {
         JetClass klass = createClass(project, "class foo { class object { } }");
         return klass.getClassObject();
     }
+
+    public static JetBlockExpression wrapInABlock(@NotNull final JetExpression expression) {
+        if (expression instanceof JetBlockExpression) {
+            return (JetBlockExpression) expression;
+        }
+        JetNamedFunction function = createFunction(expression.getProject(), "fun f() { " + expression.getText() + "}");
+        JetBlockExpression block = (JetBlockExpression) function.getBodyExpression();
+        assert block != null;
+        return new JetBlockExpression(block.getNode()) {
+            @NotNull
+            @Override
+            public List<JetElement> getStatements() {
+                return Collections.<JetElement>singletonList(expression);
+            }
+        };
+    }
 }

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

@@ -478,7 +478,16 @@ public class CandidateResolver {
         BindingContextUtils.updateRecordedType(numberType, expression, context.trace, false);
 
         if (!(expression instanceof JetConstantExpression)) {
-            updateNumberType(numberType, JetPsiUtil.deparenthesizeWithNoTypeResolution(expression, false), context);
+            JetExpression deparenthesized = JetPsiUtil.deparenthesizeWithNoTypeResolution(expression, false);
+            if (deparenthesized != expression) {
+                updateNumberType(numberType, deparenthesized, context);
+            }
+            if (deparenthesized instanceof JetBlockExpression) {
+                JetElement lastStatement = JetPsiUtil.getLastStatementInABlock((JetBlockExpression) deparenthesized);
+                if (lastStatement instanceof JetExpression) {
+                    updateNumberType(numberType, (JetExpression) lastStatement, context);
+                }
+            }
             return;
         }
         CompileTimeConstant<?> constant =

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

@@ -30,6 +30,8 @@ import org.jetbrains.jet.lang.types.JetType;
 import org.jetbrains.jet.lang.types.TypeUtils;
 import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns;
 
+import java.util.List;
+
 import static org.jetbrains.jet.lang.resolve.BindingContext.REFERENCE_TARGET;
 import static org.jetbrains.jet.lang.resolve.BindingContext.RESOLVED_CALL;
 
@@ -158,6 +160,15 @@ public class DataFlowValueFactory {
 
             return getIdForStableIdentifier(innerExpression, bindingContext);
         }
+        else if (expression instanceof JetBlockExpression) {
+            List<JetElement> statements = ((JetBlockExpression) expression).getStatements();
+            if (statements.size() == 1) {
+                JetElement lastStatement = JetPsiUtil.getLastStatementInABlock((JetBlockExpression) expression);
+                if (lastStatement instanceof JetExpression) {
+                    return getIdForStableIdentifier((JetExpression) lastStatement, bindingContext);
+                }
+            }
+        }
         else if (expression instanceof JetQualifiedExpression) {
             JetQualifiedExpression qualifiedExpression = (JetQualifiedExpression) expression;
             JetExpression receiverExpression = qualifiedExpression.getReceiverExpression();

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

@@ -95,7 +95,8 @@ public class ControlStructureTypingUtils {
                 noAnnotations, specialFunctionName, CallableMemberDescriptor.Kind.DECLARATION);
 
         TypeParameterDescriptor typeParameter = TypeParameterDescriptorImpl.createWithDefaultBound(
-                function, noAnnotations, false, Variance.INVARIANT, Name.identifierNoValidate("<TYPE-PARAMETER-FOR-" + constructionName + "-RESOLVE>"), 0);
+                function, noAnnotations, false, Variance.INVARIANT,
+                Name.identifierNoValidate("<TYPE-PARAMETER-FOR-" + constructionName + "-RESOLVE>"), 0);
 
         JetType type = new JetTypeImpl(typeParameter.getTypeConstructor(), JetScope.EMPTY);
         JetType nullableType = new JetTypeImpl(
@@ -163,7 +164,7 @@ public class ControlStructureTypingUtils {
 
     /*package*/ static Call createCallForSpecialConstruction(
             @NotNull final JetExpression expression,
-            @NotNull List<JetExpression> arguments
+            @NotNull List<? extends JetExpression> arguments
     ) {
         final List<ValueArgument> valueArguments = Lists.newArrayList();
         for (JetExpression argument : arguments) {
@@ -272,12 +273,13 @@ public class ControlStructureTypingUtils {
 
             @Override
             public Void visitBlockExpression(JetBlockExpression expression, CheckTypeContext c) {
-                List<JetElement> statements = expression.getStatements();
-                if (!statements.isEmpty()) {
-                    JetElement lastElement = statements.get(statements.size() - 1);
-                    if (lastElement instanceof JetExpression) {
-                        checkExpressionType((JetExpression) lastElement, c);
-                    }
+                if (expression.getStatements().isEmpty()) {
+                    visitExpression(expression, c);
+                    return null;
+                }
+                JetElement lastStatement = JetPsiUtil.getLastStatementInABlock(expression);
+                if (lastStatement instanceof JetExpression) {
+                    checkExpressionType((JetExpression) lastStatement, c);
                 }
                 return null;
             }

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

@@ -114,7 +114,9 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
             return getTypeInfoWhenOnlyOneBranchIsPresent(
                     elseBranch, elseScope, elseInfo, thenInfo, contextWithExpectedType, ifExpression, isStatement);
         }
-        Call callForIf = createCallForSpecialConstruction(ifExpression, Lists.newArrayList(thenBranch, elseBranch));
+        JetBlockExpression thenBlock = JetPsiFactory.wrapInABlock(thenBranch);
+        JetBlockExpression elseBlock = JetPsiFactory.wrapInABlock(elseBranch);
+        Call callForIf = createCallForSpecialConstruction(ifExpression, Lists.newArrayList(thenBlock, elseBlock));
         MutableDataFlowInfoForArguments dataFlowInfoForArguments =
                     createDataFlowInfoForArgumentsForIfCall(callForIf, thenInfo, elseInfo);
         ResolvedCall<FunctionDescriptor> resolvedCall = resolveSpecialConstructionAsCall(