From 1a1a3a2431b8d973e91cf347692c6697435a65cd Mon Sep 17 00:00:00 2001
From: Joshua Stiefer <facedelajunk@gmail.com>
Date: Wed, 12 Dec 2018 18:23:26 -0700
Subject: [PATCH] Handle functional setState in no-unused-state

---
 lib/rules/no-unused-state.js       | 22 +++++++--
 tests/lib/rules/no-unused-state.js | 79 ++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 3 deletions(-)

diff --git a/lib/rules/no-unused-state.js b/lib/rules/no-unused-state.js
index d3c307db4e..0804d5650b 100644
--- a/lib/rules/no-unused-state.js
+++ b/lib/rules/no-unused-state.js
@@ -59,6 +59,14 @@ function getInitialClassInfo() {
   };
 }
 
+function isSetStateCall(node) {
+  return (
+    node.callee.type === 'MemberExpression' &&
+    isThisExpression(node.callee.object) &&
+    getName(node.callee.property) === 'setState'
+  );
+}
+
 module.exports = {
   meta: {
     docs: {
@@ -246,13 +254,21 @@ module.exports = {
         // If we're looking at a `this.setState({})` invocation, record all the
         // properties as state fields.
         if (
-          node.callee.type === 'MemberExpression' &&
-          isThisExpression(node.callee.object) &&
-          getName(node.callee.property) === 'setState' &&
+          isSetStateCall(node) &&
           node.arguments.length > 0 &&
           node.arguments[0].type === 'ObjectExpression'
         ) {
           addStateFields(node.arguments[0]);
+        } else if (
+          isSetStateCall(node) &&
+          node.arguments.length > 0 &&
+          node.arguments[0].type === 'ArrowFunctionExpression' &&
+          node.arguments[0].body.type === 'ObjectExpression'
+        ) {
+          if (node.arguments[0].params.length > 0) {
+            classInfo.aliases.add(getName(node.arguments[0].params[0]));
+          }
+          addStateFields(node.arguments[0].body);
         }
       },
 
diff --git a/tests/lib/rules/no-unused-state.js b/tests/lib/rules/no-unused-state.js
index 8261ae708c..2742500589 100644
--- a/tests/lib/rules/no-unused-state.js
+++ b/tests/lib/rules/no-unused-state.js
@@ -576,6 +576,66 @@ eslintTester.run('no-unused-state', rule, {
         }
       }`,
       parser: 'babel-eslint'
+    }, {
+      code: `
+      class Foo extends Component {
+        state = {
+          initial: 'foo',
+        }
+        handleChange = () => {
+          this.setState(state => ({
+            current: state.initial
+          }));
+        }
+        render() {
+          const { current } = this.state;
+          return <div>{current}</div>
+        }
+      }
+      `,
+      parser: 'babel-eslint'
+    }, {
+      code: `
+      class Foo extends Component {
+        constructor(props) {
+          super(props);
+          this.state = {
+            initial: 'foo',
+          }
+        }
+        handleChange = () => {
+          this.setState(state => ({
+            current: state.initial
+          }));
+        }
+        render() {
+          const { current } = this.state;
+          return <div>{current}</div>
+        }
+      }
+      `,
+      parser: 'babel-eslint'
+    }, {
+      code: `
+      class Foo extends Component {
+        constructor(props) {
+          super(props);
+          this.state = {
+            initial: 'foo',
+          }
+        }
+        handleChange = () => {
+          this.setState((state, props) => ({
+            current: state.initial
+          }));
+        }
+        render() {
+          const { current } = this.state;
+          return <div>{current}</div>
+        }
+      }
+      `,
+      parser: 'babel-eslint'
     }
   ],
 
@@ -914,6 +974,25 @@ eslintTester.run('no-unused-state', rule, {
       }`,
       errors: getErrorMessages(['id']),
       parser: 'babel-eslint'
+    }, {
+      code: `
+      class Foo extends Component {
+        state = {
+          initial: 'foo',
+        }
+        handleChange = () => {
+          this.setState(() => ({
+            current: 'hi'
+          }));
+        }
+        render() {
+          const { current } = this.state;
+          return <div>{current}</div>
+        }
+      }
+      `,
+      parser: 'babel-eslint',
+      errors: getErrorMessages(['initial'])
     }
   ]
 });