More fixes for possible coercion to unit of type of last statement in block

compiler/frontend/src/org/jetbrains/jet/lang/diagnostics/Errors.java

@@ -192,6 +192,7 @@ public interface Errors {
     SimpleDiagnosticFactory RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY = SimpleDiagnosticFactory.create(ERROR, "Returns are not allowed for functions with expression body. Use block body in '{...}'");
     SimpleDiagnosticFactory NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY = SimpleDiagnosticFactory.create(ERROR, "A 'return' expression required in a function with a block body ('{...}')");
     ParameterizedDiagnosticFactory1<JetType> RETURN_TYPE_MISMATCH = ParameterizedDiagnosticFactory1.create(ERROR, "This function must return a value of type {0}");
+    ParameterizedDiagnosticFactory1<JetType> EXPECTED_TYPE_MISMATCH = ParameterizedDiagnosticFactory1.create(ERROR, "Expected a value of type {0}");
 
     ParameterizedDiagnosticFactory1<JetType> UPPER_BOUND_VIOLATED = ParameterizedDiagnosticFactory1.create(ERROR, "An upper bound {0} is violated"); // TODO : Message
     ParameterizedDiagnosticFactory1<JetType> FINAL_CLASS_OBJECT_UPPER_BOUND = ParameterizedDiagnosticFactory1.create(ERROR, "{0} is a final type, and thus a class object cannot extend it");

compiler/frontend/src/org/jetbrains/jet/lang/types/JetTypeInferrer.java

@@ -423,42 +423,61 @@ public class JetTypeInferrer {
                 trace.record(STATEMENT, statement);
                 final JetExpression statementExpression = (JetExpression) statement;
                 //TODO constructor assert context.expectedType != FORBIDDEN : ""
-                if (!iterator.hasNext() && context.expectedType != NO_EXPECTED_TYPE) {
+                if (!iterator.hasNext()) {
-                    if (coercionStrategyForLastExpression == CoercionStrategy.COERCION_TO_UNIT && JetStandardClasses.isUnit(context.expectedType)) {
+                    if (context.expectedType != NO_EXPECTED_TYPE) {
-                        // This implements coercion to Unit
+                        if (coercionStrategyForLastExpression == CoercionStrategy.COERCION_TO_UNIT && JetStandardClasses.isUnit(context.expectedType)) {
-                        TemporaryBindingTrace temporaryTraceExpectingUnit = TemporaryBindingTrace.create(trace);
+                            // This implements coercion to Unit
-                        final boolean[] mismatch = new boolean[1];
+                            TemporaryBindingTrace temporaryTraceExpectingUnit = TemporaryBindingTrace.create(trace);
-                        ObservableBindingTrace errorInterceptingTrace = makeTraceInterceptingTypeMismatch(temporaryTraceExpectingUnit, statementExpression, mismatch);
+                            final boolean[] mismatch = new boolean[1];
-                        newContext = newContext(errorInterceptingTrace, scope, newContext.dataFlowInfo, context.expectedType, context.expectedReturnType);
+                            ObservableBindingTrace errorInterceptingTrace = makeTraceInterceptingTypeMismatch(temporaryTraceExpectingUnit, statementExpression, mismatch);
-                        result = blockLevelVisitor.getType(statementExpression, newContext);
+                            newContext = newContext(errorInterceptingTrace, scope, newContext.dataFlowInfo, context.expectedType, context.expectedReturnType);
-                        if (mismatch[0]) {
-                            TemporaryBindingTrace temporaryTraceNoExpectedType = TemporaryBindingTrace.create(trace);
-                            mismatch[0] = false;
-                            ObservableBindingTrace interceptingTrace = makeTraceInterceptingTypeMismatch(temporaryTraceNoExpectedType, statementExpression, mismatch);
-                            newContext = newContext(interceptingTrace, scope, newContext.dataFlowInfo, NO_EXPECTED_TYPE, context.expectedReturnType);
                             result = blockLevelVisitor.getType(statementExpression, newContext);
                             if (mismatch[0]) {
-                                temporaryTraceExpectingUnit.commit();
+                                TemporaryBindingTrace temporaryTraceNoExpectedType = TemporaryBindingTrace.create(trace);
+                                mismatch[0] = false;
+                                ObservableBindingTrace interceptingTrace = makeTraceInterceptingTypeMismatch(temporaryTraceNoExpectedType, statementExpression, mismatch);
+                                newContext = newContext(interceptingTrace, scope, newContext.dataFlowInfo, NO_EXPECTED_TYPE, context.expectedReturnType);
+                                result = blockLevelVisitor.getType(statementExpression, newContext);
+                                if (mismatch[0]) {
+                                    temporaryTraceExpectingUnit.commit();
+                                }
+                                else {
+                                    temporaryTraceNoExpectedType.commit();
+                                }
                             }
                             else {
-                                temporaryTraceNoExpectedType.commit();
+                                temporaryTraceExpectingUnit.commit();
                             }
                         }
                         else {
-                            temporaryTraceExpectingUnit.commit();
+                            newContext = newContext(trace, scope, newContext.dataFlowInfo, context.expectedType, context.expectedReturnType);
+                            result = blockLevelVisitor.getType(statementExpression, newContext);
                         }
-
                     }
                     else {
-                        newContext = newContext(trace, scope, newContext.dataFlowInfo, context.expectedType, context.expectedReturnType);
                         result = blockLevelVisitor.getType(statementExpression, newContext);
+                        if (coercionStrategyForLastExpression == CoercionStrategy.COERCION_TO_UNIT) {
+                            boolean mightBeUnit = false;
+                            if (statementExpression instanceof JetDeclaration) {
+                                mightBeUnit = true;
+                            } 
+                            if (statementExpression instanceof JetBinaryExpression) {
+                                JetBinaryExpression binaryExpression = (JetBinaryExpression) statementExpression;
+                                IElementType operationType = binaryExpression.getOperationToken();
+                                if (operationType == JetTokens.EQ || assignmentOperationNames.containsKey(operationType)) {
+                                    mightBeUnit = true;
+                                }
+                            }
+                            if (mightBeUnit) {
+                                // TypeInferrerVisitorWithWritableScope should return only null or Unit for declarations and assignments
+                                assert result == null || JetStandardClasses.isUnit(result);
+                                result = JetStandardClasses.getUnitType();
+                            }
+                        }
                     }
                 }
                 else {
                     result = blockLevelVisitor.getType(statementExpression, newContext);
-                    if (coercionStrategyForLastExpression == CoercionStrategy.COERCION_TO_UNIT && result == null) {
-                        result = JetStandardClasses.getUnitType();
-                    }
                 }
 
                 DataFlowInfo newDataFlowInfo = blockLevelVisitor.getResultingDataFlowInfo();
@@ -1023,8 +1042,8 @@ public class JetTypeInferrer {
 
             if (functionTypeExpected) {
                 JetType expectedReturnType = JetStandardClasses.getReturnType(expectedType);
-                functionDescriptor.setReturnType(expectedReturnType);
                 if (JetStandardClasses.isUnit(expectedReturnType)) {
+                    functionDescriptor.setReturnType(expectedReturnType);
                     return context.services.checkType(JetStandardClasses.getFunctionType(Collections.<AnnotationDescriptor>emptyList(), effectiveReceiverType, parameterTypes, expectedReturnType), expression, context);
                 }
 
@@ -2335,10 +2354,10 @@ public class JetTypeInferrer {
                 result = getTypeForBinaryCall(context.scope, binaryOperationNames.get(operationType), context, expression);
             }
             else if (operationType == JetTokens.EQ) {
-                result = visitAssignment(expression, context);
+                result = visitAssignment(expression, contextWithExpectedType);
             }
             else if (assignmentOperationNames.containsKey(operationType)) {
-                result = visitAssignmentOperation(expression, context);
+                result = visitAssignmentOperation(expression, contextWithExpectedType);
             }
             else if (comparisonOperations.contains(operationType)) {
                 JetType compareToReturnType = getTypeForBinaryCall(context.scope, "compareTo", context, expression);
@@ -2652,7 +2671,7 @@ public class JetTypeInferrer {
                 PropertyDescriptor propertyDescriptor = context.classDescriptorResolver.resolveObjectDeclarationAsPropertyDescriptor(scope.getContainingDeclaration(), declaration, classDescriptor);
                 scope.addVariableDescriptor(propertyDescriptor);
             }
-            return null;
+            return checkExpectedType(declaration, context);
         }
 
         @Override
@@ -2688,7 +2707,7 @@ public class JetTypeInferrer {
             }
 
             scope.addVariableDescriptor(propertyDescriptor);
-            return null;
+            return checkExpectedType(property, context);
         }
 
         @Override
@@ -2696,7 +2715,7 @@ public class JetTypeInferrer {
             FunctionDescriptorImpl functionDescriptor = context.classDescriptorResolver.resolveFunctionDescriptor(scope.getContainingDeclaration(), scope, function);
             scope.addFunctionDescriptor(functionDescriptor);
             context.services.checkFunctionReturnType(context.scope, function, functionDescriptor, context.dataFlowInfo);
-            return null; 
+            return checkExpectedType(function, context);
         }
 
         @Override
@@ -2711,7 +2730,7 @@ public class JetTypeInferrer {
 
         @Override
         public JetType visitDeclaration(JetDeclaration dcl, TypeInferenceContext context) {
-            return visitJetElement(dcl, context);
+            return checkExpectedType(dcl, context);
         }
 
         @Override
@@ -2733,6 +2752,17 @@ public class JetTypeInferrer {
             else {
                 temporaryBindingTrace.commit();
             }
+            return checkExpectedType(expression, context);
+        }
+
+        @Nullable
+        private JetType checkExpectedType(JetExpression expression, TypeInferenceContext context) {
+            if (context.expectedType != NO_EXPECTED_TYPE) {
+                if (JetStandardClasses.isUnit(context.expectedType)) {
+                    return JetStandardClasses.getUnitType();
+                }
+                context.trace.report(EXPECTED_TYPE_MISMATCH.on(expression, context.expectedType));
+            }
             return null;
         }
 
@@ -2765,7 +2795,7 @@ public class JetTypeInferrer {
                     }
                 }
             }
-            return null;
+            return checkExpectedType(expression, context);
         }
 
         private JetType resolveArrayAccessToLValue(JetArrayAccessExpression arrayAccessExpression, JetExpression rightHandSide, JetSimpleNameExpression operationSign, TypeInferenceContext context) {

idea/testData/checkerWithErrorTypes/full/FunctionReturnTypes.jet

@@ -168,4 +168,42 @@ fun f(): Int = if (1 < 2) 1 else returnNothing()
 public fun <!PUBLIC_MEMBER_SHOULD_SPECIFY_TYPE!>f<!>() = 1
 class B() {
    protected fun <!PUBLIC_MEMBER_SHOULD_SPECIFY_TYPE!>f<!>() = "ss"
+}
+
+fun testFunctionLiterals() {
+    val endsWithVarDeclaration : fun() : Boolean = {
+        <!EXPECTED_TYPE_MISMATCH!>val x = 2<!>
+    }
+
+    val endsWithAssignment = { () : Int =>
+        val x = 1
+        <!EXPECTED_TYPE_MISMATCH!>x = 333<!>
+    }
+
+    val endsWithReAssignment = { () : Int =>
+        val x = 1
+        <!EXPECTED_TYPE_MISMATCH!>x += 333<!>
+    }
+
+    val endsWithFunDeclaration : fun() : String = {
+        val x = 1
+        x = 333
+        <!EXPECTED_TYPE_MISMATCH!>fun meow() : Unit {}<!>
+    }
+
+    val endsWithObjectDeclaration : fun() : Int = {
+        val x = 1
+        x = 333
+        <!EXPECTED_TYPE_MISMATCH!>object A {}<!>
+    }
+
+    val expectedUnitReturnType1 = { () : Unit =>
+        val x = 1
+    }
+
+    val expectedUnitReturnType2 = { () : Unit =>
+        fun meow() : Unit {}
+        object A {}
+    }
+
 }

idea/testData/codegen/regressions/kt343.jet

@@ -1,5 +1,9 @@
 import java.util.ArrayList
 
+fun launch(f : fun() : Unit) {
+    f()
+}
+
 fun box(): String {
     val list = ArrayList<Int>()
     val foo : fun() : Unit = {
@@ -7,10 +11,14 @@ fun box(): String {
     }
     foo()
 
+    launch({
+        list.add(3)
+    })
+
     val bar = {
         val x = 1   //second exception
     }
     bar()
 
-    return if (list.get(0) == 2) "OK" else "fail"
+    return if (list.size() == 2 && list.get(0) == 2 && list.get(1) == 3) "OK" else "fail"
 }