From c98e5b206cb5b8f249895722d6af3f129eac9c82 Mon Sep 17 00:00:00 2001
From: rikitariko <leoprogbars@gmail.com>
Date: Thu, 16 Jul 2020 15:36:19 +0300
Subject: [PATCH] Added the ability to check unused DI in the controller which
 is located in component and directive.

---
 rules/di-unused.js          |  5 ++-
 rules/utils/angular-rule.js | 62 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/rules/di-unused.js b/rules/di-unused.js
index 5176c48c..62828b94 100644
--- a/rules/di-unused.js
+++ b/rules/di-unused.js
@@ -65,7 +65,6 @@ module.exports = {
             'angular?animation': reportUnusedVariables,
             'angular?config': reportUnusedVariables,
             'angular?controller': reportUnusedVariables,
-            'angular?directive': reportUnusedVariables,
             'angular?factory': reportUnusedVariables,
             'angular?filter': reportUnusedVariables,
             'angular?inject': reportUnusedVariables,
@@ -75,6 +74,10 @@ module.exports = {
                 reportUnusedVariables(null, providerFn);
                 reportUnusedVariables(null, $get);
             },
+            'angular?directive': function(callee, directiveFn, controllerFn) {
+                reportUnusedVariables(callee, directiveFn);
+                reportUnusedVariables(callee, controllerFn);
+            },
 
             'Program:exit': function() {
                 var globalScope = context.getScope();
diff --git a/rules/utils/angular-rule.js b/rules/utils/angular-rule.js
index 5e731cb0..93e94664 100644
--- a/rules/utils/angular-rule.js
+++ b/rules/utils/angular-rule.js
@@ -202,6 +202,22 @@ function angularRule(ruleDefinition) {
         if (node.type === 'FunctionExpression' || node.type === 'ArrowFunctionExpression' || node.type === 'FunctionDeclaration') {
             return node;
         }
+
+        if (node.type === 'ObjectExpression') {
+            var controller;
+            node.properties.forEach(function(property) {
+                if (property.key.name === 'controller') {
+                    if (property.value.type === 'FunctionExpression' || property.value.type === 'ArrowFunctionExpression') {
+                        controller = property.value;
+                    }
+                    if (property.value.type === 'ArrayExpression') {
+                        controller = property.value.elements[property.value.elements.length - 1];
+                    }
+                }
+            });
+            return controller;
+        }
+
         if (node.type !== 'Identifier') {
             return;
         }
@@ -250,7 +266,6 @@ function angularRule(ruleDefinition) {
             case 'component':
             case 'config':
             case 'controller':
-            case 'directive':
             case 'factory':
             case 'filter':
             case 'run':
@@ -258,6 +273,8 @@ function angularRule(ruleDefinition) {
                 return [node.callExpression.callee, node.fn];
             case 'provider':
                 return assembleProviderArguments(node);
+            case 'directive':
+                return assembleDirectiveArguments(node);
         }
     }
 
@@ -270,6 +287,49 @@ function angularRule(ruleDefinition) {
         return [node.callExpression, node.fn, findProviderGet(node.fn)];
     }
 
+    function assembleDirectiveArguments(node) {
+        return [node.callExpression, node.fn, findDirectiveFunction(node.fn)];
+    }
+
+    function findDirectiveFunction(directiveFn) {
+        if (!directiveFn) {
+            return;
+        }
+
+        var controllerFn;
+        var returnFn;
+        if (directiveFn.body.type !== 'BlockStatement') {
+            returnFn = directiveFn.body.properties;
+            returnFn.forEach(function(arg) {
+                if (arg.key.name === 'controller' && (arg.value.type === 'FunctionExpression' || arg.value.type === 'ArrowFunctionExpression')) {
+                    controllerFn = arg.value;
+                }
+                if (arg.value.type === 'ArrayExpression') {
+                    controllerFn = arg.value.elements[arg.value.elements.length - 1];
+                }
+            });
+            return controllerFn;
+        }
+
+        directiveFn.body.body.forEach(function(statement) {
+            if (statement.type === 'ReturnStatement' && statement.argument.type === 'ObjectExpression') {
+                returnFn = statement.argument.properties;
+                returnFn.forEach(function(arg) {
+                    if (arg.key.name === 'controller') {
+                        if (arg.value.type === 'FunctionExpression' || arg.value.type === 'ArrowFunctionExpression') {
+                            controllerFn = arg.value;
+                        }
+                        if (arg.value.type === 'ArrayExpression') {
+                            controllerFn = arg.value.elements[arg.value.elements.length - 1];
+                        }
+                    }
+                });
+            }
+        });
+
+        return controllerFn;
+    }
+
     /**
      * Find the $get function of a provider based on the provider function body.
      */