From c857785910c303823b8ae273c2ccb577b728ed97 Mon Sep 17 00:00:00 2001
From: Svetlana Isakova <svetlana.isakova@jetbrains.com>
Date: Fri, 9 Mar 2012 12:20:10 +0400
Subject: [PATCH] KT-1002 If there's no return in function highlight only last
 token but not last statement

---
 .../lang/cfg/JetFlowInformationProvider.java  |  8 +++-
 .../jet/lang/diagnostics/Errors.java          | 13 +++++-
 .../jet/lang/psi/JetBlockExpression.java      | 11 +++++
 .../diagnostics/tests/FunctionReturnTypes.jet | 44 +++++++++----------
 .../diagnostics/tests/TypeInference.jet       |  4 +-
 .../tests/controlStructures/kt786.jet         |  6 +--
 .../diagnostics/tests/regressions/kt402.jet   |  4 +-
 .../diagnostics/tests/regressions/kt442.jet   |  2 +-
 .../diagnostics/tests/regressions/kt456.jet   | 12 ++---
 .../diagnostics/tests/regressions/kt58.jet    | 12 ++---
 idea/testData/checker/FunctionReturnTypes.jet | 44 +++++++++----------
 11 files changed, 93 insertions(+), 67 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetFlowInformationProvider.java b/compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetFlowInformationProvider.java
index a2b37b27648..5ffb96e6c83 100644
--- a/compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetFlowInformationProvider.java
+++ b/compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetFlowInformationProvider.java
@@ -149,7 +149,7 @@ public class JetFlowInformationProvider {
         }
     }
     
-    public void checkDefiniteReturn(@NotNull JetDeclarationWithBody function, final @NotNull JetType expectedReturnType) {
+    public void checkDefiniteReturn(@NotNull final JetDeclarationWithBody function, final @NotNull JetType expectedReturnType) {
         assert function instanceof JetDeclaration;
 
         JetExpression bodyExpression = function.getBodyExpression();
@@ -172,6 +172,7 @@ public class JetFlowInformationProvider {
             trace.report(UNREACHABLE_CODE.on(element));
         }
 
+        final boolean[] noReturnError = new boolean[] { false };
         for (JetElement returnedExpression : returnedExpressions) {
             returnedExpression.accept(new JetVisitorVoid() {
                 @Override
@@ -184,11 +185,14 @@ public class JetFlowInformationProvider {
                 @Override
                 public void visitExpression(JetExpression expression) {
                     if (blockBody && expectedReturnType != NO_EXPECTED_TYPE && !JetStandardClasses.isUnit(expectedReturnType) && !rootUnreachableElements.contains(expression)) {
-                        trace.report(NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY.on(expression));
+                        noReturnError[0] = true;
                     }
                 }
             });
         }
+        if (noReturnError[0]) {
+            trace.report(NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY.on(function));
+        }
     }
 
     private Set<JetElement> collectUnreachableCode(@NotNull JetElement subroutine) {
diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/diagnostics/Errors.java b/compiler/frontend/src/org/jetbrains/jet/lang/diagnostics/Errors.java
index 7afc0d7a16d..c599e0e1d84 100644
--- a/compiler/frontend/src/org/jetbrains/jet/lang/diagnostics/Errors.java
+++ b/compiler/frontend/src/org/jetbrains/jet/lang/diagnostics/Errors.java
@@ -232,7 +232,18 @@ public interface Errors {
     DiagnosticFactory1<JetExpression, JetType> CALLEE_NOT_A_FUNCTION = DiagnosticFactory1.create(ERROR, "Expecting a function type, but found {0}");
 
     DiagnosticFactory<JetReturnExpression> RETURN_IN_FUNCTION_WITH_EXPRESSION_BODY = DiagnosticFactory.create(ERROR, "Returns are not allowed for functions with expression body. Use block body in '{...}'");
-    DiagnosticFactory<JetExpression> NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY = DiagnosticFactory.create(ERROR, "A 'return' expression required in a function with a block body ('{...}')");
+    DiagnosticFactory<JetDeclarationWithBody> NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY = DiagnosticFactory.create(ERROR, "A 'return' expression required in a function with a block body ('{...}')", new PositioningStrategy<JetDeclarationWithBody>() {
+        @NotNull
+        @Override
+        public List<TextRange> mark(@NotNull JetDeclarationWithBody element) {
+            JetExpression bodyExpression = element.getBodyExpression();
+            if (!(bodyExpression instanceof JetBlockExpression)) return Collections.emptyList();
+            JetBlockExpression blockExpression = (JetBlockExpression) bodyExpression;
+            TextRange lastBracketRange = blockExpression.getLastBracketRange();
+            if (lastBracketRange == null) return Collections.emptyList();
+            return markRange(lastBracketRange);
+        }
+    });
     DiagnosticFactory1<JetExpression, JetType> RETURN_TYPE_MISMATCH = DiagnosticFactory1.create(ERROR, "This function must return a value of type {0}");
     DiagnosticFactory1<JetExpression, JetType> EXPECTED_TYPE_MISMATCH = DiagnosticFactory1.create(ERROR, "Expected a value of type {0}");
     DiagnosticFactory1<JetBinaryExpression, JetType> ASSIGNMENT_TYPE_MISMATCH = DiagnosticFactory1.create(ERROR, "Expected a value of type {0}. Assignment operation is not an expression, so it does not return any value");
diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetBlockExpression.java b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetBlockExpression.java
index 789a7a78f77..2fc18d0b772 100644
--- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetBlockExpression.java
+++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetBlockExpression.java
@@ -16,8 +16,13 @@
 
 package org.jetbrains.jet.lang.psi;
 
+import com.google.common.collect.Lists;
 import com.intellij.lang.ASTNode;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lexer.JetTokens;
 
 import java.util.Arrays;
 import java.util.List;
@@ -44,4 +49,10 @@ public class JetBlockExpression extends JetExpression implements JetStatementExp
     public List<JetElement> getStatements() {
         return Arrays.asList(findChildrenByClass(JetElement.class));
     }
+
+    @Nullable
+    public TextRange getLastBracketRange() {
+        PsiElement rBrace = findChildByType(JetTokens.RBRACE);
+        return rBrace != null ? rBrace.getTextRange() : null;
+    }
 }
diff --git a/compiler/testData/diagnostics/tests/FunctionReturnTypes.jet b/compiler/testData/diagnostics/tests/FunctionReturnTypes.jet
index f623565a366..068f6388c9c 100644
--- a/compiler/testData/diagnostics/tests/FunctionReturnTypes.jet
+++ b/compiler/testData/diagnostics/tests/FunctionReturnTypes.jet
@@ -31,12 +31,12 @@ fun unitShort() : Unit = #()
 fun unitShortConv() : Unit = <!TYPE_MISMATCH!>1<!>
 fun unitShortNull() : Unit = <!TYPE_MISMATCH!>null<!>
 
-fun intEmpty() : Int <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>{}<!>
+fun intEmpty() : Int {<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 fun intShortInfer() = 1
 fun intShort() : Int = 1
 //fun intBlockInfer()  {1}
 fun intBlock() : Int {return 1}
-fun intBlock1() : Int {<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY, UNUSED_EXPRESSION!>1<!>}
+fun intBlock1() : Int {<!UNUSED_EXPRESSION!>1<!><!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 
 fun intString(): Int = <!TYPE_MISMATCH!>"s"<!>
 fun intFunctionLiteral(): Int = <!TYPE_MISMATCH!>{ 10 }<!>
@@ -47,8 +47,8 @@ fun blockReturnValueTypeMatch() : Int {return 1}
 fun blockReturnValueTypeMismatchUnit() : Int {return <!TYPE_MISMATCH!>#()<!>}
 
 fun blockAndAndMismatch() : Int {
-  <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>true && false<!>
-}
+  true && false
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 fun blockAndAndMismatch1() : Int {
   return <!TYPE_MISMATCH!>true && false<!>
 }
@@ -57,8 +57,8 @@ fun blockAndAndMismatch2() : Int {
 }
 
 fun blockAndAndMismatch3() : Int {
-  <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>true || false<!>
-}
+  true || false
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 fun blockAndAndMismatch4() : Int {
   return <!TYPE_MISMATCH!>true || false<!>
 }
@@ -91,9 +91,9 @@ fun blockReturnValueTypeMatch6() : Int {
 }
 fun blockReturnValueTypeMatch7() : Int {
   if (1 > 2)
-    <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>1.0<!>
-  else <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>2.0<!>
-}
+    1.0
+  else 2.0
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 fun blockReturnValueTypeMatch8() : Int {
   if (1 > 2)
     1.0
@@ -101,38 +101,38 @@ fun blockReturnValueTypeMatch8() : Int {
   return 1
 }
 fun blockReturnValueTypeMatch9() : Int {
-  <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>if (1 > 2)
-    <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>1.0<!><!>
-}
+  if (1 > 2)
+    1.0
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 fun blockReturnValueTypeMatch10() : Int {
   return <!TYPE_MISMATCH!>if (1 > 2)
     1<!>
 }
 fun blockReturnValueTypeMatch11() : Int {
-  <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>if (1 > 2)
-  else <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>1.0<!><!>
-}
+  if (1 > 2)
+  else 1.0
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 fun blockReturnValueTypeMatch12() : Int {
   if (1 > 2)
     return 1
   else return <!ERROR_COMPILE_TIME_VALUE!>1.0<!>
 }
 fun blockNoReturnIfValDeclaration(): Int {
-  <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>val <!UNUSED_VARIABLE!>x<!> = 1<!>
-}
+  val <!UNUSED_VARIABLE!>x<!> = 1
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 fun blockNoReturnIfEmptyIf(): Int {
-  if (1 < 2) <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>{}<!> else <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>{}<!>
-}
+  if (1 < 2) {} else {}
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 fun blockNoReturnIfUnitInOneBranch(): Int {
   if (1 < 2) {
     return 1
   } else {
-    if (3 < 4) <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>{
-    }<!> else {
+    if (3 < 4) {
+    } else {
       return 2
     }
   }
-}
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 fun nonBlockReturnIfEmptyIf(): Int = if (1 < 2) <!TYPE_MISMATCH!>{}<!> else <!TYPE_MISMATCH!>{}<!>
 fun nonBlockNoReturnIfUnitInOneBranch(): Int = if (1 < 2) <!TYPE_MISMATCH!>{}<!> else 2
 
diff --git a/compiler/testData/diagnostics/tests/TypeInference.jet b/compiler/testData/diagnostics/tests/TypeInference.jet
index 4dd8649fff0..54631bc97c8 100644
--- a/compiler/testData/diagnostics/tests/TypeInference.jet
+++ b/compiler/testData/diagnostics/tests/TypeInference.jet
@@ -1,9 +1,9 @@
 class C<T>() {
-  fun foo() : T <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>{}<!>
+  fun foo() : T {<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 }
 
 fun foo(<!UNUSED_PARAMETER!>c<!>: C<Int>) {}
-fun bar<T>() : C<T> <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>{}<!>
+fun bar<T>() : C<T> {<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 
 fun main(args : Array<String>) {
   val <!UNUSED_VARIABLE!>a<!> : C<Int> = C();
diff --git a/compiler/testData/diagnostics/tests/controlStructures/kt786.jet b/compiler/testData/diagnostics/tests/controlStructures/kt786.jet
index 102b7dbc3cc..5c303d00f4c 100644
--- a/compiler/testData/diagnostics/tests/controlStructures/kt786.jet
+++ b/compiler/testData/diagnostics/tests/controlStructures/kt786.jet
@@ -5,11 +5,11 @@ fun foo() : Int {
     val d = 2
     var z = 0
     when(d) {
-        is 5, is 3 -> <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY, UNUSED_CHANGED_VALUE!>z++<!>
-        <!ELSE_MISPLACED_IN_WHEN!>else<!> -> { <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>z = <!UNUSED_VALUE!>-1000<!><!> }
+        is 5, is 3 -> <!UNUSED_CHANGED_VALUE!>z++<!>
+        <!ELSE_MISPLACED_IN_WHEN!>else<!> -> { z = <!UNUSED_VALUE!>-1000<!> }
         <!UNREACHABLE_CODE!>return z -> 34<!>
     }
-}
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 
 //test unreachable code
 fun fff(): Int {
diff --git a/compiler/testData/diagnostics/tests/regressions/kt402.jet b/compiler/testData/diagnostics/tests/regressions/kt402.jet
index a1a83723e34..2effc4eb86b 100644
--- a/compiler/testData/diagnostics/tests/regressions/kt402.jet
+++ b/compiler/testData/diagnostics/tests/regressions/kt402.jet
@@ -1,8 +1,8 @@
 package kt402
 
 fun getTypeChecker() : (Any)->Boolean {
-  <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY, UNUSED_FUNCTION_LITERAL!>{ (a : Any) -> a is <!UNRESOLVED_REFERENCE!>T<!> }<!> // reports unsupported
-}
+  <!UNUSED_FUNCTION_LITERAL!>{ (a : Any) -> a is <!UNRESOLVED_REFERENCE!>T<!> }<!> // reports unsupported
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 fun f() : (Any) -> Boolean {
   return { (a : Any) -> a is String }
 }
\ No newline at end of file
diff --git a/compiler/testData/diagnostics/tests/regressions/kt442.jet b/compiler/testData/diagnostics/tests/regressions/kt442.jet
index 9b4e1c745d2..a10fd8565d2 100644
--- a/compiler/testData/diagnostics/tests/regressions/kt442.jet
+++ b/compiler/testData/diagnostics/tests/regressions/kt442.jet
@@ -6,7 +6,7 @@ fun testFunny() {
     val <!UNUSED_VARIABLE!>a<!> : Int = funny {1}
 }
 
-fun <T> funny2(<!UNUSED_PARAMETER!>f<!> : (t : T) -> T) : T <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>{}<!>
+fun <T> funny2(<!UNUSED_PARAMETER!>f<!> : (t : T) -> T) : T {<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 
 fun testFunny2() {
     val <!UNUSED_VARIABLE!>a<!> : Int = funny2 {it}
diff --git a/compiler/testData/diagnostics/tests/regressions/kt456.jet b/compiler/testData/diagnostics/tests/regressions/kt456.jet
index 8d67585ff38..5ad3ffc2a6f 100644
--- a/compiler/testData/diagnostics/tests/regressions/kt456.jet
+++ b/compiler/testData/diagnostics/tests/regressions/kt456.jet
@@ -4,15 +4,15 @@ package kt456
 
 class A() {
     val i: Int
-    get() : Int <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>{  //no error
-    }<!>
+    get() : Int {  //no error
+    <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 }
 
 //more tests
 class B() {
     val i: Int
-    get() <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>{  //no error
-    }<!>
+    get() {  //no error
+    <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 }
 
 class C() {
@@ -22,9 +22,9 @@ class C() {
             doSmth()
         }
         finally {
-            <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>doSmth()<!>
+            doSmth()
         }
-    }
+    <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 }
 
 fun doSmth() {}
\ No newline at end of file
diff --git a/compiler/testData/diagnostics/tests/regressions/kt58.jet b/compiler/testData/diagnostics/tests/regressions/kt58.jet
index 7b7a7526755..4d2edb3c7b4 100644
--- a/compiler/testData/diagnostics/tests/regressions/kt58.jet
+++ b/compiler/testData/diagnostics/tests/regressions/kt58.jet
@@ -42,18 +42,18 @@ fun t3() : Int {
         doSmth(2)
     }
     finally {
-        <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>doSmth(3)<!>
+        doSmth(3)
     }
-}
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 
 fun t4() : Int {
     try {
         return 1
     }
     catch (e: UnsupportedOperationException) {
-        <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>doSmth(2)<!>
+        doSmth(2)
     }
-}
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 
 fun t5() : Int {
     try {
@@ -84,9 +84,9 @@ fun t7() : Int {
         return 2
     }
     finally {
-        <!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>doSmth(3)<!>
+        doSmth(3)
     }
-}
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 
 fun doSmth(<!UNUSED_PARAMETER!>i<!>: Int) {
 }
\ No newline at end of file
diff --git a/idea/testData/checker/FunctionReturnTypes.jet b/idea/testData/checker/FunctionReturnTypes.jet
index 88d53b7bdc8..c12170c87d0 100644
--- a/idea/testData/checker/FunctionReturnTypes.jet
+++ b/idea/testData/checker/FunctionReturnTypes.jet
@@ -26,12 +26,12 @@ fun unitShort() : Unit = #()
 fun unitShortConv() : Unit = <error>1</error>
 fun unitShortNull() : Unit = <error>null</error>
 
-fun intEmpty() : Int <error>{}</error>
+fun intEmpty() : Int {<error>}</error>
 fun intShortInfer() = 1
 fun intShort() : Int = 1
 //fun intBlockInfer()  {1}
 fun intBlock() : Int {return 1}
-fun intBlock1() : Int {<error>1</error>}
+fun intBlock1() : Int {<warning>1</warning><error>}</error>
 
 fun intString(): Int = <error>"s"</error>
 fun intFunctionLiteral(): Int = <error>{ 10 }</error>
@@ -42,8 +42,8 @@ fun blockReturnValueTypeMatch() : Int {return 1}
 fun blockReturnValueTypeMismatchUnit() : Int {return <error>#()</error>}
 
 fun blockAndAndMismatch() : Int {
-  <error>true && false</error>
-}
+  true && false
+<error>}</error>
 fun blockAndAndMismatch1() : Int {
   return <error>true && false</error>
 }
@@ -52,8 +52,8 @@ fun blockAndAndMismatch2() : Int {
 }
 
 fun blockAndAndMismatch3() : Int {
-  <error>true || false</error>
-}
+  true || false
+<error>}</error>
 fun blockAndAndMismatch4() : Int {
   return <error>true || false</error>
 }
@@ -86,9 +86,9 @@ fun blockReturnValueTypeMatch6() : Int {
 }
 fun blockReturnValueTypeMatch7() : Int {
   if (1 > 2)
-    <error>1.0</error>
-  else <error>2.0</error>
-}
+    1.0
+  else 2.0
+<error>}</error>
 fun blockReturnValueTypeMatch8() : Int {
   if (1 > 2)
     1.0
@@ -96,38 +96,38 @@ fun blockReturnValueTypeMatch8() : Int {
   return 1
 }
 fun blockReturnValueTypeMatch9() : Int {
-  <error>if (1 > 2)
-    <error>1.0</error></error>
-}
+  if (1 > 2)
+    1.0
+<error>}</error>
 fun blockReturnValueTypeMatch10() : Int {
   return <error>if (1 > 2)
     1</error>
 }
 fun blockReturnValueTypeMatch11() : Int {
-  <error>if (1 > 2)
-  else <error>1.0</error></error>
-}
+  if (1 > 2)
+  else 1.0
+<error>}</error>
 fun blockReturnValueTypeMatch12() : Int {
   if (1 > 2)
     return 1
   else return <error>1.0</error>
 }
 fun blockNoReturnIfValDeclaration(): Int {
-  <error>val x = 1</error>
-}
+  val <warning>x</warning> = 1
+<error>}</error>
 fun blockNoReturnIfEmptyIf(): Int {
-  if (1 < 2) <error>{}</error> else <error>{}</error>
-}
+  if (1 < 2) {} else {}
+<error>}</error>
 fun blockNoReturnIfUnitInOneBranch(): Int {
   if (1 < 2) {
     return 1
   } else {
-    if (3 < 4) <error>{
-    }</error> else {
+    if (3 < 4) {
+    } else {
       return 2
     }
   }
-}
+<error>}</error>
 fun nonBlockReturnIfEmptyIf(): Int = if (1 < 2) <error>{}</error> else <error>{}</error>
 fun nonBlockNoReturnIfUnitInOneBranch(): Int = if (1 < 2) <error>{}</error> else 2