[Fresh] Initial Babel plugin implementation

packages/react-fresh/src/ReactFreshBabelPlugin.js

@@ -7,9 +7,128 @@
 
 'use strict';
 
-// TODO
 export default function(babel) {
+  const {types: t, template} = babel;
+
+  const registrationsByProgramPath = new Map();
+  function createRegistration(programPath, persistentID) {
+    const handle = programPath.scope.generateUidIdentifier('c');
+    if (!registrationsByProgramPath.has(programPath)) {
+      registrationsByProgramPath.set(programPath, []);
+    }
+    const registrations = registrationsByProgramPath.get(programPath);
+    registrations.push({
+      handle,
+      persistentID,
+    });
+    return handle;
+  }
+
+  const buildRegistrationCall = template(`
+    __register__(HANDLE, PERSISTENT_ID);
+  `);
+
+  function isComponentishName(name) {
+    return typeof name === 'string' && name[0] >= 'A' && name[0] <= 'Z';
+  }
+
+  function isComponentish(node) {
+    switch (node.type) {
+      case 'FunctionDeclaration':
+        return node.id !== null && isComponentishName(node.id.name);
+      case 'VariableDeclarator':
+        return (
+          isComponentishName(node.id.name) &&
+          node.init !== null &&
+          (node.init.type === 'FunctionExpression' ||
+            (node.init.type === 'ArrowFunctionExpression' &&
+              node.init.body.type !== 'ArrowFunctionExpression'))
+        );
+      default:
+        return false;
+    }
+  }
+
   return {
-    visitor: {},
+    visitor: {
+      FunctionDeclaration(path) {
+        let programPath;
+        let insertAfterPath;
+        switch (path.parent.type) {
+          case 'Program':
+            insertAfterPath = path;
+            programPath = path.parentPath;
+            break;
+          case 'ExportNamedDeclaration':
+          case 'ExportDefaultDeclaration':
+            insertAfterPath = path.parentPath;
+            programPath = insertAfterPath.parentPath;
+            break;
+          default:
+            return;
+        }
+        const maybeComponent = path.node;
+        if (!isComponentish(maybeComponent)) {
+          return;
+        }
+        const functionName = path.node.id.name;
+        const handle = createRegistration(programPath, functionName);
+        insertAfterPath.insertAfter(
+          t.expressionStatement(
+            t.assignmentExpression('=', handle, path.node.id),
+          ),
+        );
+      },
+      VariableDeclaration(path) {
+        let programPath;
+        switch (path.parent.type) {
+          case 'Program':
+            programPath = path.parentPath;
+            break;
+          case 'ExportNamedDeclaration':
+          case 'ExportDefaultDeclaration':
+            programPath = path.parentPath.parentPath;
+            break;
+          default:
+            return;
+        }
+        const declPath = path.get('declarations');
+        if (declPath.length !== 1) {
+          return;
+        }
+        const firstDeclPath = declPath[0];
+        const maybeComponent = firstDeclPath.node;
+        if (!isComponentish(maybeComponent)) {
+          return;
+        }
+        const functionName = maybeComponent.id.name;
+        const initPath = firstDeclPath.get('init');
+        const handle = createRegistration(programPath, functionName);
+        initPath.replaceWith(
+          t.assignmentExpression('=', handle, initPath.node),
+        );
+      },
+      Program: {
+        exit(path) {
+          const registrations = registrationsByProgramPath.get(path);
+          if (registrations === undefined) {
+            return;
+          }
+          registrationsByProgramPath.delete(path);
+          const declarators = [];
+          path.pushContainer('body', t.variableDeclaration('var', declarators));
+          registrations.forEach(({handle, persistentID}) => {
+            path.pushContainer(
+              'body',
+              buildRegistrationCall({
+                HANDLE: handle,
+                PERSISTENT_ID: t.stringLiteral(persistentID),
+              }),
+            );
+            declarators.push(t.variableDeclarator(handle));
+          });
+        },
+      },
+    },
   };
 }

packages/react-fresh/src/__tests__/ReactFreshBabelPlugin-test.js

@@ -12,12 +12,178 @@ let freshPlugin = require('react-fresh/babel');
 
 function transform(input, options = {}) {
   return babel.transform(input, {
-    plugins: [[freshPlugin]],
+    babelrc: false,
+    plugins: ['syntax-jsx', freshPlugin],
   }).code;
 }
 
 describe('ReactFreshBabelPlugin', () => {
-  it('hello world', () => {
-    expect(transform(`hello()`)).toMatchSnapshot();
+  it('registers top-level function declarations', () => {
+    // Hello and Bar should be registered, handleClick shouldn't.
+    expect(
+      transform(`
+        function Hello() {
+          function handleClick() {}
+          return <h1 onClick={handleClick}>Hi</h1>;
+        }
+
+        function Bar() {
+          return <Hello />;
+        }
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('registers top-level exported function declarations', () => {
+    expect(
+      transform(`
+        export function Hello() {
+          function handleClick() {}
+          return <h1 onClick={handleClick}>Hi</h1>;
+        }
+
+        export default function Bar() {
+          return <Hello />;
+        }
+
+        function Baz() {
+          return <h1>OK</h1>;
+        }
+
+        const NotAComp = 'hi';
+        export { Baz, NotAComp };
+
+        export function sum() {}
+        export const Bad = 42;
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('registers top-level exported named arrow functions', () => {
+    expect(
+      transform(`
+        export const Hello = () => {
+          function handleClick() {}
+          return <h1 onClick={handleClick}>Hi</h1>;
+        };
+
+        export let Bar = (props) => <Hello />;
+
+        export default () => {
+          // This one should be ignored.
+          // You should name your components.
+          return <Hello />;
+        };
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('uses original function declaration if it get reassigned', () => {
+    // This should register the original version.
+    // TODO: in the future, we may *also* register the wrapped one.
+    expect(
+      transform(`
+        function Hello() {
+          return <h1>Hi</h1>;
+        }
+        Hello = connect(Hello);
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('only registers pascal case functions', () => {
+    // Should not get registered.
+    expect(
+      transform(`
+        function hello() {
+          return 2 * 2;
+        }
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('registers top-level variable declarations with function expressions', () => {
+    // Hello and Bar should be registered; handleClick, sum, Baz, and Qux shouldn't.
+    expect(
+      transform(`
+        let Hello = function() {
+          function handleClick() {}
+          return <h1 onClick={handleClick}>Hi</h1>;
+        };
+        const Bar = function Baz() {
+          return <Hello />;
+        };
+        function sum() {}
+        let Baz = 10;
+        var Qux;
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('registers top-level variable declarations with arrow functions', () => {
+    // Hello, Bar, and Baz should be registered; handleClick and sum shouldn't.
+    expect(
+      transform(`
+        let Hello = () => {
+          const handleClick = () => {};
+          return <h1 onClick={handleClick}>Hi</h1>;
+        }
+        const Bar = () => {
+          return <Hello />;
+        };
+        var Baz = () => <div />;
+        var sum = () => {};
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('ignores HOC definitions', () => {
+    // TODO: we might want to handle HOCs at usage site, however.
+    // TODO: it would be nice if we could always avoid registering
+    // a function that is known to return a function or other non-node.
+    expect(
+      transform(`
+        let connect = () => {
+          function Comp() {
+            const handleClick = () => {};
+            return <h1 onClick={handleClick}>Hi</h1>;
+          }
+          return Comp;
+        };
+        function withRouter() {
+          return function Child() {
+            const handleClick = () => {};
+            return <h1 onClick={handleClick}>Hi</h1>;
+          }
+        };
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('ignores complex definitions', () => {
+    expect(
+      transform(`
+        let A = foo ? () => {
+          return <h1>Hi</h1>;
+        } : null
+        const B = (function Foo() {
+          return <h1>Hi</h1>;
+        })();
+        let C = () => () => {
+          return <h1>Hi</h1>;
+        };
+        let D = bar && (() => {
+          return <h1>Hi</h1>;
+        });
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('ignores unnamed function declarations', () => {
+    expect(
+      transform(`
+        export default function() {}
+    `),
+    ).toMatchSnapshot();
   });
 });