Support for arrow function

ng-annotate-main.js

@@ -705,7 +705,7 @@ function followReference(node) {
         // {type: "VariableDeclarator", id: {type: "Identifier", name: "foo"}, init: ..}
         return parent;
     } else if (kind === "fun") {
-        assert(ptype === "FunctionDeclaration" || ptype === "FunctionExpression")
+        assert(ptype === "FunctionDeclaration" || ptype === "FunctionExpression" || ptype === "ArrowFunctionExpression")
         // FunctionDeclaration is the common case, i.e.
         // function foo(a, b) {}
 
@@ -958,11 +958,10 @@ function judgeInjectArraySuspect(node, ctx) {
 }
 
 function jumpOverIife(node) {
-    let outerfn;
-    if (!(node.type === "CallExpression" && (outerfn = node.callee).type === "FunctionExpression")) {
+    const outerfn = node.callee
+    if (!(node.type === "CallExpression" && (outerfn.type === "FunctionExpression" || outerfn.type === "ArrowFunctionExpression"))) {
         return node;
     }
-
     const outerbody = outerfn.body.body;
     for (let i = 0; i < outerbody.length; i++) {
         const statement = outerbody[i];
@@ -988,9 +987,10 @@ function isAnnotatedArray(node) {
         return false;
     }
     const elements = node.elements;
+    const lastElement = last(elements);
 
     // last should be a function expression
-    if (elements.length === 0 || last(elements).type !== "FunctionExpression") {
+    if (elements.length === 0 || (lastElement.type !== "FunctionExpression" && lastElement.type !== "ArrowFunctionExpression")) {
         return false;
     }
 
@@ -1005,7 +1005,7 @@ function isAnnotatedArray(node) {
     return true;
 }
 function isFunctionExpressionWithArgs(node) {
-    return node.type === "FunctionExpression" && node.params.length >= 1;
+    return (node.type === "FunctionExpression" || node.type === "ArrowFunctionExpression") && node.params.length >= 1;
 }
 function isFunctionDeclarationWithArgs(node) {
     return node.type === "FunctionDeclaration" && node.params.length >= 1;

nginject.js

@@ -14,7 +14,7 @@ module.exports = {
 function inspectNode(node, ctx) {
     if (node.type === "CallExpression") {
         inspectCallExpression(node, ctx);
-    } else if (node.type === "FunctionExpression" || node.type === "FunctionDeclaration") {
+    } else if (node.type === "FunctionExpression" || node.type === "ArrowFunctionExpression" || node.type === "FunctionDeclaration") {
         inspectFunction(node, ctx);
     }
 }
@@ -81,6 +81,11 @@ function inspectFunction(node, ctx) {
 function matchPrologueDirectives(prologueDirectives, node) {
     const body = node.body.body;
 
+    // There is no body when not using curly brace, such as x => x
+    if (!body) {
+        return;
+    }
+
     let found = null;
     for (let i = 0; i < body.length; i++) {
         if (body[i].type !== "ExpressionStatement") {
@@ -195,7 +200,7 @@ function nestedObjectValues(node, res) {
 
     node.properties.forEach(function(prop) {
         const v = prop.value;
-        if (is.someof(v.type, ["FunctionExpression", "ArrayExpression"])) {
+        if (is.someof(v.type, ["FunctionExpression", "ArrowFunctionExpression", "ArrayExpression"])) {
             res.push(v);
         } else if (v.type === "ObjectExpression") {
             nestedObjectValues(v, res);

tests/original.js

@@ -1016,3 +1016,12 @@ var MyCtrl = (function() {
 })();
 
 myMod.service("a", MyCtrl);
+
+var ArrowFnCtrl = ($scope) => {
+    "ngInject"
+    return $scope
+}
+
+var ArrowFnCtrl = /*@ngInject*/($scope) => { $scope }
+
+var ArrowFnCtrl = /*@ngInject*/($filter) => $filter

tests/with_annotations.js

@@ -1059,3 +1059,13 @@ var MyCtrl = (function() {
 })();
 
 myMod.service("a", MyCtrl);
+
+var ArrowFnCtrl = ($scope) => {
+    "ngInject"
+    return $scope
+}
+ArrowFnCtrl.$inject = ["$scope"];
+
+var ArrowFnCtrl = /*@ngInject*/["$scope", ($scope) => { $scope }]
+
+var ArrowFnCtrl = /*@ngInject*/["$filter", ($filter) => $filter]