working for function expressions

Author: alexreardon

community/memoize-one/src/5.0.0/transform.spec.ts

@@ -3,7 +3,17 @@ import * as transformer from './transform';
 import prettier from 'prettier';
 
 function format(source: string): string {
-  return prettier.format(source, { parser: 'typescript' }).trim();
+  return prettier
+    .format(source, {
+      parser: 'typescript',
+      trailingComma: 'all',
+      semi: true,
+      tabWidth: 2,
+      useTabs: false,
+      singleQuote: true,
+      printWidth: 100,
+    })
+    .trim();
 }
 
 describe('memoize-one@5.0.0 transform', () => {
@@ -35,7 +45,47 @@ describe('memoize-one@5.0.0 transform', () => {
     );
   });
 
-  it.only('should wrap inline equality functions', () => {
+  it.only('should wrap inline equality arrow functions', () => {
+    const result = applyTransform(
+      transformer,
+      format(`
+        import memoize from 'memoize-one';
+
+        function add(a: number, b: number) {
+          return a + b;
+        }
+
+        const memoized = memoize(add, (a, b) => {
+          return a === b;
+        });
+      `),
+      { parser: 'tsx' },
+    );
+
+    expect(result).toEqual(
+      format(`
+      import memoize from 'memoize-one';
+
+      function add(a: number, b: number) {
+        return a + b;
+      }
+
+      const memoized = memoize(add, (newArgs, lastArgs) => {
+        if (newArgs.length !== lastArgs.length) {
+          return false;
+        }
+
+        const __equalityFn = (a, b) => {
+          return a === b;
+        };
+
+        return newArgs.every((newArg, index) => __equalityFn(newArg, lastArgs[index]));
+      });
+      `),
+    );
+  });
+
+  it('should wrap inline equality function declarations', () => {
     const result = applyTransform(
       transformer,
       format(`
@@ -65,11 +115,11 @@ describe('memoize-one@5.0.0 transform', () => {
           return false;
         }
 
-        function isEqual(a, b) {
+        const __equalityFn = function isEqual(a, b) {
           return a === b;
         }
 
-        return newArgs.every((newArg, index) => isEqual(newArg, lastArgs[index]));
+        return newArgs.every((newArg, index) => __equalityFn(newArg, lastArgs[index]));
       });
       `),
     );

community/memoize-one/src/5.0.0/transform.ts

@@ -40,7 +40,10 @@ export default function transformer(
       if (second == null) {
         return;
       }
-      if (second.type === 'FunctionExpression') {
+      if (
+        second.type === 'FunctionExpression' ||
+        second.type === 'ArrowFunctionExpression'
+      ) {
         const customEqualityFn = j.arrowFunctionExpression(
           [j.identifier('newArgs'), j.identifier('lastArgs')],
           j.blockStatement([
@@ -56,25 +59,37 @@ export default function transformer(
                   j.identifier('length'),
                 ),
               ),
-              j.returnStatement(j.booleanLiteral(false)),
+              j.blockStatement([j.returnStatement(j.booleanLiteral(false))]),
+            ),
+            j.variableDeclaration('const', [
+              j.variableDeclarator(j.identifier('__equalityFn'), second),
+            ]),
+            j.returnStatement(
+              j.callExpression(
+                j.memberExpression(
+                  j.identifier('newArgs'),
+                  j.identifier('every'),
+                ),
+                [
+                  j.arrowFunctionExpression(
+                    [j.identifier('newArg'), j.identifier('index')],
+                    j.callExpression(j.identifier('__equalityFn'), [
+                      j.identifier('newArg'),
+                      j.memberExpression(
+                        j.identifier('lastArgs'),
+                        j.identifier('index'),
+                        // computed lastArgs[index]
+                        true,
+                      ),
+                    ]),
+                  ),
+                ],
+              ),
             ),
-            j.returnStatement(j.booleanLiteral(true)),
           ]),
         );
 
-        // (newArgs, lastArgs) => {
-        //   if (newArgs.length !== lastArgs.length) {
-        //     return false;
-        //   }
-
-        //   return newArgs.every((newArg, index) =>
-        //     isEqual(newArg, lastArgs[index]),
-        //   );
-        // };
-
         call.value.arguments = [first, customEqualityFn];
-        // console.log('FunctionExpression', call);
-        // call.replace(j.functionExpression([first, 'hi']));
         return;
       }