change state using arrow functions instead

Author: barelyhuman

source/index.ts

@@ -5,7 +5,7 @@ interface ToModifyVariableI {
   simplified: string
 }
 
-export default function () {
+export default function toMutableTransformer() {
   const toMod: ToModifyVariableI[] = []
   return {
     visitor: {
@@ -46,9 +46,21 @@ function transformToStateByScope(path: any, toMod: ToModifyVariableI[]) {
     VariableDeclaration({node}: {node: t.VariableDeclaration}) {
       transformReactiveDeclarations(node, toMod, path)
     },
-    ExpressionStatement({node}: {node: t.ExpressionStatement}) {
-      transformAssignmentExpression(node, toMod)
+    ExpressionStatement(expressionPath: any) {
+      const node: t.ExpressionStatement = expressionPath.node
+      const replacedStateExpression = transformAssignmentExpression(
+        node,
+        toMod,
+        expressionPath
+      )
       transformUpdateExpression(node, toMod)
+
+      // HACK: any identifier inside the replace expression that is reactive to change itself to
+      // the substate name so $a => a => _prevState
+      if (replacedStateExpression) {
+        const {subStateIdentifier, normName} = replacedStateExpression
+        replaceIdentifiers(expressionPath, normName, subStateIdentifier.name)
+      }
     },
   })
 }
@@ -92,7 +104,8 @@ function transformReactiveDeclarations(
 
 function transformAssignmentExpression(
   node: t.ExpressionStatement,
-  toMod: ToModifyVariableI[]
+  toMod: ToModifyVariableI[],
+  path: any
 ) {
   if (!t.isAssignmentExpression(node.expression)) {
     return
@@ -112,47 +125,52 @@ function transformAssignmentExpression(
   const normName = normalizeName(expression.left.name)
   const setterName = getSetterName(normName)
 
-  let callArgs: t.Expression[]
+  const subStateIdentifier: t.Identifier =
+    path.scope.generateUidIdentifier('prevState')
+
+  let params: any = {}
 
   switch (expression.operator) {
     case '=': {
-      callArgs = [{...expression.right}]
+      params = expression.right
       break
     }
 
     case '+=': {
-      callArgs = [
-        t.binaryExpression('+', t.identifier(normName), expression.right),
-      ]
+      params = t.binaryExpression('+', subStateIdentifier, expression.right)
       break
     }
 
     case '-=': {
-      callArgs = [
-        t.binaryExpression('-', t.identifier(normName), expression.right),
-      ]
+      params = t.binaryExpression('-', subStateIdentifier, expression.right)
       break
     }
 
     case '/=': {
-      callArgs = [
-        t.binaryExpression('/', t.identifier(normName), expression.right),
-      ]
+      params = t.binaryExpression('/', subStateIdentifier, expression.right)
+
       break
     }
 
     case '*=': {
-      callArgs = [
-        t.binaryExpression('*', t.identifier(normName), expression.right),
-      ]
+      params = t.binaryExpression('*', subStateIdentifier, expression.right)
       break
     }
     default: {
-      callArgs = []
+      params = []
     }
   }
 
-  node.expression = t.callExpression(t.identifier(setterName), callArgs)
+  let arrowExpression = params
+  if (!t.isArrowFunctionExpression(params)) {
+    arrowExpression = t.arrowFunctionExpression([subStateIdentifier], params)
+  }
+
+  node.expression = t.callExpression(t.identifier(setterName), [
+    arrowExpression,
+  ])
+
+  return {subStateIdentifier, normName, setterName}
 }
 
 function transformUpdateExpression(
@@ -211,3 +229,13 @@ function getSetterName(normalizedName: string) {
 function normalizeName(n: string) {
   return n.replace(/\$/, '')
 }
+
+function replaceIdentifiers(path: any, toCompare: string, replaceWith: string) {
+  return path.traverse({
+    Identifier({node}: {node: t.Identifier}) {
+      if (node.name === toCompare) {
+        node.name = replaceWith
+      }
+    },
+  })
+}

test/snapshots/test.ts.md

@@ -4,7 +4,7 @@ The actual snapshot is saved in `test.ts.snap`.
 
 Generated by [AVA](https://avajs.dev).
 
-## works
+## Simple Transform
 
 > Snapshot 1
 
@@ -14,15 +14,15 @@ Generated by [AVA](https://avajs.dev).
       const [a, setA] = React.useState(1);␊
     ␊
       const onPress = () => {␊
-        setA(a + 1);␊
+        setA(_prevState => _prevState + 1);␊
       };␊
     ␊
       return /*#__PURE__*/React.createElement("div", null, /*#__PURE__*/React.createElement("p", null, a), /*#__PURE__*/React.createElement("button", {␊
         onClick: onPress␊
       }, "Press"));␊
     }`
 
-## Simple Transform
+## Simple Transform 2
 
 > Snapshot 1
 
@@ -32,14 +32,37 @@ Generated by [AVA](https://avajs.dev).
       const [a, setA] = React.useState(1);␊
     ␊
       const onPress = () => {␊
-        setA(a + 1);␊
+        setA(x => {␊
+          console.log(x);␊
+          return x + 1;␊
+        });␊
       };␊
     ␊
       return /*#__PURE__*/React.createElement("div", null, /*#__PURE__*/React.createElement("p", null, a), /*#__PURE__*/React.createElement("button", {␊
         onClick: onPress␊
       }, "Press"));␊
     }`
 
+## Simple Transform 3
+
+> Snapshot 1
+
+    `import * as React from "react";␊
+    ␊
+    function Component() {␊
+      const [a, setA] = React.useState("");␊
+    ␊
+      const onChange = e => {␊
+        setA(_prevState => e.target.value);␊
+        setA(_prevState2 => _prevState2.toUpperCase());␊
+      };␊
+    ␊
+      return /*#__PURE__*/React.createElement("div", null, /*#__PURE__*/React.createElement("input", {␊
+        value: a,␊
+        onChange: onChange␊
+      }));␊
+    }`
+
 ## Check Functional Scope
 
 > Snapshot 1
@@ -51,7 +74,7 @@ Generated by [AVA](https://avajs.dev).
       const [a, setA] = React.useState(1);␊
     ␊
       const onPress = () => {␊
-        setA(a + 1);␊
+        setA(_prevState => _prevState + 1);␊
         $b = 3;␊
       };␊
     ␊
@@ -71,7 +94,7 @@ Generated by [AVA](https://avajs.dev).
       const [a, setA] = React.useState(1);␊
     ␊
       const onPress = () => {␊
-        setA(a + 1);␊
+        setA(_prevState => _prevState + 1);␊
         $b = 3;␊
       };␊
     ␊
@@ -91,7 +114,7 @@ Generated by [AVA](https://avajs.dev).
       const [a, setA] = React.useState(1);␊
     ␊
       const onPress = () => {␊
-        setA(a + 1);␊
+        setA(_prevState => _prevState + 1);␊
         $b = 3;␊
       };␊
     ␊
@@ -104,7 +127,7 @@ Generated by [AVA](https://avajs.dev).
       const [a, setA] = React.useState(3);␊
     ␊
       const onPress = () => {␊
-        setA(5);␊
+        setA(_prevState2 => 5);␊
         $b = 3;␊
       };␊
     ␊
@@ -126,7 +149,7 @@ Generated by [AVA](https://avajs.dev).
       }, [a]);␊
     ␊
       const onPress = () => {␊
-        setA(a + 1);␊
+        setA(_prevState => _prevState + 1);␊
       };␊
     ␊
       return {␊
@@ -144,8 +167,8 @@ Generated by [AVA](https://avajs.dev).
       const [count, setCount] = React.useState(1);␊
     ␊
       const handleClick = () => {␊
-        setCount(count + 1);␊
-        setCount(count * 2);␊
+        setCount(_prevState => _prevState + 1);␊
+        setCount(_prevState2 => _prevState2 * 2);␊
       };␊
     ␊
       return /*#__PURE__*/React.createElement("div", null, /*#__PURE__*/React.createElement("h1", null, count), /*#__PURE__*/React.createElement("button", {␊
@@ -168,7 +191,31 @@ Generated by [AVA](https://avajs.dev).
         const x = { ...user␊
         };␊
         x.name = "barelyhuman";␊
-        setUser(x);␊
+        setUser(_prevState => x);␊
+      };␊
+    ␊
+      return /*#__PURE__*/React.createElement(React.Fragment, null, /*#__PURE__*/React.createElement("p", null, user.name), /*#__PURE__*/React.createElement("button", {␊
+        onClick: updateUser␊
+      }, "Click Me"));␊
+    }`
+
+## Object Update Arrow Func
+
+> Snapshot 1
+
+    `import * as React from "react";␊
+    ␊
+    function App() {␊
+      const [user, setUser] = React.useState({␊
+        name: "reaper"␊
+      });␊
+    ␊
+      const updateUser = () => {␊
+        setUser(x => Object.assign({}, {␊
+          x␊
+        }, {␊
+          name: "barelyhuman"␊
+        }));␊
       };␊
     ␊
       return /*#__PURE__*/React.createElement(React.Fragment, null, /*#__PURE__*/React.createElement("p", null, user.name), /*#__PURE__*/React.createElement("button", {␊
@@ -191,7 +238,7 @@ Generated by [AVA](https://avajs.dev).
         const _nextUsers = users.slice();␊
     ␊
         _nextUsers[0].name = "barelyhuman";␊
-        setUser(_nextUsers);␊
+        setUsers(_prevState => _nextUsers);␊
       };␊
     ␊
       return /*#__PURE__*/React.createElement(React.Fragment, null, users.map(user => {␊