Support updating multiple states

Author: qnighy

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 - Added
   - Support for method-binding patterns e.g. `this.foo = this.foo.bind(this);`
+  - Support for multiple states in one `setState()` call
 
 ## 0.1.8
 

README.md

@@ -247,7 +247,7 @@ If you need to enforce specific styles, use Prettier or ESLint or whatever is yo
 - [ ] Support for `this.state`
   - [x] Decompose `this.state` into `useState` variables
   - [x] Rename states if necessary
-  - [ ] Support updating multiple states at once
+  - [x] Support updating multiple states at once
   - [ ] Support functional updates
   - [ ] Support lazy initialization
 - [ ] Support for refs

src/analysis.ts

@@ -20,7 +20,7 @@ export type {
 } from "./analysis/pre.js";
 export { preanalyzeClass } from "./analysis/pre.js";
 export type { LocalManager } from "./analysis/local.js";
-export type { StateObjAnalysis } from "./analysis/state.js";
+export type { StateObjAnalysis, SetStateSite, SetStateFieldSite } from "./analysis/state.js";
 export { needAlias } from "./analysis/prop.js";
 export type { PropsObjAnalysis } from "./analysis/prop.js";
 

src/analysis/state.ts

@@ -9,6 +9,7 @@ import { trackMember } from "./track_member.js";
 
 export type StateObjAnalysis = {
   states: Map<string, StateAnalysis>;
+  setStateSites: SetStateSite[];
 };
 
 export type StateAnalysis = {
@@ -19,7 +20,7 @@ export type StateAnalysis = {
   sites: StateSite[];
 };
 
-export type StateSite = StateInitSite | StateExprSite | SetStateSite;
+export type StateSite = StateInitSite | StateExprSite;
 
 export type StateInitSite = {
   type: "state_init";
@@ -31,9 +32,14 @@ export type StateExprSite = {
   path: NodePath<Expression>;
   owner: string | undefined;
 };
+
 export type SetStateSite = {
-  type: "setState";
   path: NodePath<CallExpression>;
+  fields: SetStateFieldSite[];
+}
+
+export type SetStateFieldSite = {
+  name: string;
   valuePath: NodePath<Expression>;
 }
 
@@ -120,6 +126,7 @@ export function analyzeState(
       throw new AnalysisError(`Non-analyzable this.state`);
     }
   }
+  const setStateSites: SetStateSite[] = [];
   for (const site of setStateAnalysis.sites) {
     if (site.type !== "expr" || site.hasWrite) {
       throw new AnalysisError(`Invalid use of this.setState`);
@@ -135,22 +142,25 @@ export function analyzeState(
     const arg0 = args[0]!;
     if (arg0.isObjectExpression()) {
       const props = arg0.get("properties");
-      if (props.length !== 1) {
-        throw new AnalysisError(`Multiple assignments in setState is not yet supported`);
-      }
-      const prop0 = props[0]!;
-      if (!prop0.isObjectProperty()) {
-        throw new AnalysisError(`Non-analyzable setState`);
-      }
-      const setStateName = memberName(prop0.node);
-      if (setStateName == null) {
-        throw new AnalysisError(`Non-analyzable setState name`);
+      const fields: SetStateFieldSite[] = [];
+      for (const prop of props) {
+        if (!prop.isObjectProperty()) {
+          throw new AnalysisError(`Non-analyzable setState`);
+        }
+        const setStateName = memberName(prop.node);
+        if (setStateName == null) {
+          throw new AnalysisError(`Non-analyzable setState name`);
+        }
+        // Ensure the state exists
+        getState(setStateName);
+        fields.push({
+          name: setStateName,
+          valuePath: prop.get("value") as NodePath<Expression>,
+        });
       }
-      const state = getState(setStateName);
-      state.sites.push({
-        type: "setState",
+      setStateSites.push({
         path: gpPath,
-        valuePath: prop0.get("value") as NodePath<Expression>,
+        fields,
       });
     } else {
       throw new AnalysisError(`Non-analyzable setState`);
@@ -185,5 +195,5 @@ export function analyzeState(
     }
     state.init = state.sites.find((site): site is StateInitSite => site.type === "state_init");
   }
-  return { states };
+  return { states, setStateSites };
 }

src/index.test.ts

@@ -832,6 +832,36 @@ describe("react-declassify", () => {
       expect(transform(input)).toBe(output);
     });
 
+    it("transforms multi-value assignments", () => {
+      const input = dedent`\
+        class C extends React.Component {
+          render() {
+            return <button onClick={() => this.setState({ foo: 3, bar: 4 })} />
+          }
+          frob() {
+            this.setState({
+              foo: 1,
+              bar: 2,
+            });
+          }
+        }
+      `;
+      const output = dedent`\
+        const C = () => {
+          const [foo, setFoo] = React.useState();
+          const [bar, setBar] = React.useState();
+
+          function frob() {
+            setFoo(1);
+            setBar(2);
+          }
+
+          return <button onClick={() => (setFoo(3), setBar(4))} />;
+        };
+      `;
+      expect(transform(input)).toBe(output);
+    });
+
     it("transforms state types (type alias)", () => {
       const input = dedent`\
         type Props = {};

src/index.ts

@@ -1,7 +1,7 @@
 import type { ArrowFunctionExpression, ClassMethod, ClassPrivateMethod, Expression, FunctionDeclaration, FunctionExpression, Identifier, ImportDeclaration, MemberExpression, ObjectMethod, Pattern, RestElement, Statement, TSEntityName, TSType, TSTypeAnnotation, TSTypeParameterDeclaration, VariableDeclaration } from "@babel/types";
 import type { NodePath, PluginObj, PluginPass } from "@babel/core";
 import { assignReturnType, assignTypeAnnotation, assignTypeArguments, assignTypeParameters, importName, isTS, nonNullPath } from "./utils.js";
-import { AnalysisError, analyzeClass, preanalyzeClass, AnalysisResult, PreAnalysisResult, needsProps, LibRef, needAlias } from "./analysis.js";
+import { AnalysisError, analyzeClass, preanalyzeClass, AnalysisResult, PreAnalysisResult, needsProps, LibRef, needAlias, SetStateFieldSite } from "./analysis.js";
 
 type Options = {};
 
@@ -147,17 +147,32 @@ function transformClass(analysis: AnalysisResult, options: { ts: boolean }, babe
       if (site.type === "expr") {
         // this.state.foo -> foo
         site.path.replaceWith(t.identifier(stateAnalysis.localName!));
-      } else if (site.type === "setState") {
-        // this.setState({ foo: 1 }) -> setFoo(1)
-        site.path.replaceWith(
-          t.callExpression(
-            t.identifier(stateAnalysis.localSetterName!),
-            [site.valuePath.node]
-          )
-        );
       }
     }
   }
+  for (const site of analysis.state.setStateSites) {
+    function setter(field: SetStateFieldSite) {
+      // this.setState({ foo: 1 }) -> setFoo(1)
+      return t.callExpression(
+        t.identifier(analysis.state.states.get(field.name)!.localSetterName!),
+        [field.valuePath.node]
+      );
+    }
+    if (site.fields.length === 1) {
+      const field = site.fields[0]!;
+      site.path.replaceWith(setter(field));
+    } else if (site.path.parentPath.isExpressionStatement()) {
+      site.path.parentPath.replaceWithMultiple(
+        site.fields.map((field) =>
+          t.expressionStatement(setter(field))
+        )
+      );
+    } else if (site.fields.length === 0) {
+      site.path.replaceWith(t.nullLiteral());
+    } else {
+      site.path.replaceWith(t.sequenceExpression(site.fields.map(setter)));
+    }
+  }
   for (const [, field] of analysis.userDefined.fields) {
     if (field.type === "user_defined_function" || field.type === "user_defined_ref") {
       for (const site of field.sites) {