fix: update react-fresh babel plugin

examples/styled-components/src/Spring.js

@@ -0,0 +1,18 @@
+import { animated, useSpring } from 'react-spring';
+import React, { useCallback, useState } from 'react';
+
+export function SpringTest() {
+  const [thingDone, toggleThingDone] = useState(false);
+  const doTheThing = useCallback(() => toggleThingDone(!thingDone), [thingDone]);
+
+  const fader = useSpring({ opacity: thingDone ? 1 : 0 });
+
+  return (
+    <React.Fragment>
+      <animated.h1 style={fader}>You did the thing!</animated.h1>
+      <button type="button" onClick={doTheThing}>
+        {thingDone ? 'Undo The Thing' : 'Do The Thing'}
+      </button>
+    </React.Fragment>
+  );
+}

src/babel.dev.js

@@ -257,12 +257,13 @@ module.exports = composePlugins([
     // removing everything we dont want right now
 
     // registration
-    delete p.visitor.Program;
+    // delete p.visitor.Program;
+    // delete p.visitor.Program.exit;
 
     // registrations
-    delete p.visitor.FunctionDeclaration.enter;
-    delete p.visitor.FunctionDeclaration.leave;
-    delete p.visitor.VariableDeclaration;
+    // delete p.visitor.FunctionDeclaration.enter;
+    // delete p.visitor.FunctionDeclaration.leave;
+    // delete p.visitor.VariableDeclaration;
 
     return p;
   },

src/fresh/babel.js

@@ -165,25 +165,6 @@ export default function(babel) {
     return false;
   }
 
-  let hookCalls = new WeakMap();
-
-  function recordHookCall(functionNode, hookCallPath, hookName) {
-    if (!hookCalls.has(functionNode)) {
-      hookCalls.set(functionNode, []);
-    }
-    let hookCallsForFn = hookCalls.get(functionNode);
-    let key = '';
-    if (hookCallPath.parent.type === 'VariableDeclarator') {
-      // TODO: if there is no LHS, consider some other heuristic.
-      key = hookCallPath.parentPath.get('id').getSource();
-    }
-    hookCallsForFn.push({
-      name: hookName,
-      callee: hookCallPath.node.callee,
-      key,
-    });
-  }
-
   function isBuiltinHook(hookName) {
     switch (hookName) {
       case 'useState':
@@ -236,9 +217,64 @@ export default function(babel) {
 
   let seenForRegistration = new WeakSet();
   let seenForSignature = new WeakSet();
-  let seenForHookCalls = new WeakSet();
   let seenForOutro = new WeakSet();
 
+  let hookCalls = new WeakMap();
+  const HookCallsVisitor = {
+    CallExpression(path) {
+      const node = path.node;
+      const callee = node.callee;
+
+      // Note: this visitor MUST NOT mutate the tree in any way.
+      // It runs early in a separate traversal and should be very fast.
+
+      let name = null;
+      switch (callee.type) {
+        case 'Identifier':
+          name = callee.name;
+          break;
+        case 'MemberExpression':
+          name = callee.property.name;
+          break;
+      }
+      if (name === null || !/^use[A-Z]/.test(name)) {
+        return;
+      }
+      const fnScope = path.scope.getFunctionParent();
+      if (fnScope === null) {
+        return;
+      }
+
+      // This is a Hook call. Record it.
+      const fnNode = fnScope.block;
+      if (!hookCalls.has(fnNode)) {
+        hookCalls.set(fnNode, []);
+      }
+      let hookCallsForFn = hookCalls.get(fnNode);
+      let key = '';
+      if (path.parent.type === 'VariableDeclarator') {
+        // TODO: if there is no LHS, consider some other heuristic.
+        key = path.parentPath.get('id').getSource();
+      }
+
+      // Some built-in Hooks reset on edits to arguments.
+      const args = path.get('arguments');
+      if (name === 'useState' && args.length > 0) {
+        // useState second argument is initial state.
+        key += '(' + args[0].getSource() + ')';
+      } else if (name === 'useReducer' && args.length > 1) {
+        // useReducer second argument is initial state.
+        key += '(' + args[1].getSource() + ')';
+      }
+
+      hookCallsForFn.push({
+        callee: path.node.callee,
+        name,
+        key,
+      });
+    },
+  };
+
   return {
     visitor: {
       ExportDefaultDeclaration(path) {
@@ -289,6 +325,7 @@ export default function(babel) {
       },
       FunctionDeclaration: {
         enter(path) {
+          return;
           const node = path.node;
           let programPath;
           let insertAfterPath;
@@ -342,6 +379,7 @@ export default function(babel) {
           );
         },
         exit(path) {
+          //return;
           const node = path.node;
           const id = node.id;
           if (id === null) {
@@ -358,7 +396,7 @@ export default function(babel) {
             return;
           }
           seenForSignature.add(node);
-          // Don't muatte the tree above this point.
+
 
           // Unlike with __register__, this needs to work for nested
           // declarations too. So we need to search for a path where
@@ -438,6 +476,7 @@ export default function(babel) {
         },
       },
       VariableDeclaration(path) {
+        return;
         const node = path.node;
         let programPath;
         let insertAfterPath;
@@ -510,39 +549,16 @@ export default function(babel) {
           },
         );
       },
-      CallExpression(path) {
-        const node = path.node;
-        const callee = node.callee;
-
-        let name = null;
-        switch (callee.type) {
-          case 'Identifier':
-            name = callee.name;
-            break;
-          case 'MemberExpression':
-            name = callee.property.name;
-            break;
-        }
-        if (name === null || !/^use[A-Z]/.test(name)) {
-          return;
-        }
-
-        // Make sure we're not recording the same calls twice.
-        // This can happen if another Babel plugin replaces parents.
-        if (seenForHookCalls.has(node)) {
-          return;
-        }
-        seenForHookCalls.add(node);
-        // Don't mutate the tree above this point.
-
-        const fn = path.scope.getFunctionParent();
-        if (fn === null) {
-          return;
-        }
-        recordHookCall(fn.block, path, name);
-      },
       Program: {
+        enter(path) {
+          // This is a separate early visitor because we need to collect Hook calls
+          // and "const [foo, setFoo] = ..." signatures before the destructuring
+          // transform mangles them. This extra traversal is not ideal for perf,
+          // but it's the best we can do until we stop transpiling destructuring.
+          path.traverse(HookCallsVisitor);
+        },
         exit(path) {
+          return;
           const registrations = registrationsByProgramPath.get(path);
           if (registrations === undefined) {
             return;

test/__babel_fixtures__/hooks.js

@@ -1,4 +1,5 @@
 import React, {useState} from 'react';
+import {useExternalHook} from 'external-hook'
 
 const NoHooks = () => <div>no hooks</div>;
 
@@ -32,6 +33,7 @@ const useForwardRefFunctionHook = React.forwardRef(function () {
 const useCustomHook = () => {
   useState(42);
   useEffectHook();
+  useExternalHook();
 };
 
 function useFunc () {