Modified the CheckMissingReturn pass to recognize arrow functions

(which may have expressions rather than block function bodies).

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=157847043

src/com/google/javascript/jscomp/CheckMissingReturn.java

@@ -90,11 +90,22 @@ public boolean apply(DiGraphEdge<Node, ControlFlowGraph.Branch> input) {
 
   @Override
   public void enterScope(NodeTraversal t) {
-    JSType returnType = explicitReturnExpected(t.getScopeRoot());
+    Node n = t.getScopeRoot();
+    JSType returnType = explicitReturnExpected(n);
+
     if (returnType == null) {
+      // No return value is expected, so nothing to check.
       return;
     }
 
+    if (n.isArrowFunction()) {
+      Node functionBody = NodeUtil.getFunctionBody(n);
+      if (!functionBody.isNormalBlock()) {
+        // Body is an expression, which is the implicit return value.
+        return;
+      }
+    }
+
     if (fastAllPathsReturnCheck(t.getControlFlowGraph())) {
       return;
     }
@@ -155,14 +166,14 @@ public void visit(NodeTraversal t, Node n, Node parent) {
    *
    * @return If a return type is expected, returns it. Otherwise, returns null.
    */
-  private JSType explicitReturnExpected(Node scope) {
-    FunctionType scopeType = JSType.toMaybeFunctionType(scope.getJSType());
+  private JSType explicitReturnExpected(Node scopeRoot) {
+    FunctionType scopeType = JSType.toMaybeFunctionType(scopeRoot.getJSType());
 
     if (scopeType == null) {
       return null;
     }
 
-    if (isEmptyFunction(scope)) {
+    if (isEmptyFunction(scopeRoot)) {
       return null;
     }
 

test/com/google/javascript/jscomp/CheckMissingReturnTest.java

@@ -127,31 +127,31 @@ public void testFinallyStatements() {
     // return statements in the three possible configurations: both scopes
     // return; enclosed doesn't return; enclosing doesn't return.
     testNotMissing(
-        "try {" +
-        "   /** @return {number} */ function f() {" +
-        "       try { return 1; }" +
-        "       finally { }" +
-        "   };" +
-        "   return 1;" +
-        "}" +
-        "finally { }");
+        "try {"
+            + "   /** @return {number} */ function f() {"
+            + "       try { return 1; }"
+            + "       finally { }"
+            + "   };"
+            + "   return 1;"
+            + "}"
+            + "finally { }");
     testMissing(
-        "try {" +
-        "   /** @return {number} */ function f() {" +
-        "       try { }" +
-        "       finally { }" +
-        "   };" +
-        "   return 1;" +
-        "}" +
-        "finally { }");
+        "try {"
+            + "   /** @return {number} */ function f() {"
+            + "       try { }"
+            + "       finally { }"
+            + "   };"
+            + "   return 1;"
+            + "}"
+            + "finally { }");
     testMissing(
-        "try {" +
-        "   /** @return {number} */ function f() {" +
-        "       try { return 1; }" +
-        "       finally { }" +
-        "   };" +
-        "}" +
-        "finally { }");
+        "try {"
+            + "   /** @return {number} */ function f() {"
+            + "       try { return 1; }"
+            + "       finally { }"
+            + "   };"
+            + "}"
+            + "finally { }");
   }
 
   public void testKnownConditions() {
@@ -192,8 +192,7 @@ public void testMultiConditions() {
 
   public void testIssue779() {
     testNotMissing(
-        "var a = f(); try { alert(); if (a > 0) return 1; }" +
-        "finally { a = 5; } return 2;");
+        "var a = f(); try { alert(); if (a > 0) return 1; }" + "finally { a = 5; } return 2;");
   }
 
   public void testConstructors() {
@@ -207,9 +206,9 @@ public void testConstructors() {
 
   public void testClosureAsserts() {
     String closureDefs =
-        "/** @const */ var goog = {};\n" +
-        "goog.asserts = {};\n" +
-        "goog.asserts.fail = function(x) {};";
+        "/** @const */ var goog = {};\n"
+            + "goog.asserts = {};\n"
+            + "goog.asserts.fail = function(x) {};";
 
     testNotMissing(closureDefs + "goog.asserts.fail('');");
 
@@ -275,4 +274,28 @@ private void testNotMissing(String body) {
   private void testMissing(String body) {
     testMissing("number", body);
   }
+
+  public void testArrowFunctions() {
+    setAcceptedLanguage(LanguageMode.ECMASCRIPT_2015);
+
+    // undefined return statement
+    testNoWarning(LINE_JOINER.join("/** @return {undefined} */", "() => {}"));
+
+    // arrow function with expression function body
+    testSame(LINE_JOINER.join("/** @return {number} */", "() => 1"));
+
+    testSame(LINE_JOINER.join("/** @return {number} */", "(a) => (a > 3) ? 1 : 0"));
+
+    // arrow function with block function body
+    testSame(
+        LINE_JOINER.join(
+            "/** @return {number} */", "(a) => { if (a > 3) { return 1; } else { return 0; }}"));
+
+    testWarning(
+        LINE_JOINER.join("/** @return {number} */", "(a) => { if (a > 3) { return 1; } else { } }"),
+        CheckMissingReturn.MISSING_RETURN_STATEMENT);
+
+    // arrow function return is object literal
+    testSame("(a) => ({foo: 1});");
+  }
 }