[Fresh] Capture Hook signatures lazily on first render (#15832)

* Split the signature call into two calls

This adds a render-time signature call by making __signature__ curried. We need both calls. The init time tells us which type has which signature. The render time call says when's a good time to capture the lazy Hooks tree. This is necessary for supporting inline requires. I will do that in next commit.

* Lazily compute Hook list on first render

This ensures inline requires don't break comparisons between Hook signatures of previous and next versions by caching Hook list at the time of first render.

* Refactor computing Hook signature keys

Instead of a traversal during the comparison, explicitly compute full keys. This makes it easier to debug mismatches.

packages/react-fresh/src/ReactFreshBabelPlugin.js

@@ -410,7 +410,25 @@ export default function(babel) {
             return;
           }
           seenForSignature.add(node);
-          // Don't muatte the tree above this point.
+          // Don't mutate the tree above this point.
+
+          const sigCallID = path.scope.generateUidIdentifier('_s');
+          path.scope.parent.push({
+            id: sigCallID,
+            init: t.callExpression(t.identifier('__signature__'), []),
+          });
+
+          // The signature call is split in two parts. One part is called inside the function.
+          // This is used to signal when first render happens.
+          path
+            .get('body')
+            .unshiftContainer(
+              'body',
+              t.expressionStatement(t.callExpression(sigCallID, [])),
+            );
+
+          // The second call is around the function itself.
+          // This is used to associate a type with a signature.
 
           // Unlike with __register__, this needs to work for nested
           // declarations too. So we need to search for a path where
@@ -429,7 +447,7 @@ export default function(babel) {
           insertAfterPath.insertAfter(
             t.expressionStatement(
               t.callExpression(
-                t.identifier('__signature__'),
+                sigCallID,
                 createArgumentsForSignature(
                   id,
                   signature,
@@ -456,6 +474,24 @@ export default function(babel) {
           seenForSignature.add(node);
           // Don't mutate the tree above this point.
 
+          const sigCallID = path.scope.generateUidIdentifier('_s');
+          path.scope.parent.push({
+            id: sigCallID,
+            init: t.callExpression(t.identifier('__signature__'), []),
+          });
+
+          // The signature call is split in two parts. One part is called inside the function.
+          // This is used to signal when first render happens.
+          path
+            .get('body')
+            .unshiftContainer(
+              'body',
+              t.expressionStatement(t.callExpression(sigCallID, [])),
+            );
+
+          // The second call is around the function itself.
+          // This is used to associate a type with a signature.
+
           if (path.parent.type === 'VariableDeclarator') {
             let insertAfterPath = null;
             path.find(p => {
@@ -475,7 +511,7 @@ export default function(babel) {
             insertAfterPath.insertAfter(
               t.expressionStatement(
                 t.callExpression(
-                  t.identifier('__signature__'),
+                  sigCallID,
                   createArgumentsForSignature(
                     path.parent.id,
                     signature,
@@ -489,7 +525,7 @@ export default function(babel) {
             // let Foo = hoc(() => {})
             path.replaceWith(
               t.callExpression(
-                t.identifier('__signature__'),
+                sigCallID,
                 createArgumentsForSignature(node, signature, path.scope),
               ),
             );

packages/react-fresh/src/ReactFreshRuntime.js

@@ -15,8 +15,9 @@ import type {
 import {REACT_MEMO_TYPE, REACT_FORWARD_REF_TYPE} from 'shared/ReactSymbols';
 
 type Signature = {|
-  key: string,
+  ownKey: string,
   forceReset: boolean,
+  fullKey: string | null, // Contains keys of nested Hooks. Computed lazily.
   getCustomHooks: () => Array<Function>,
 |};
 
@@ -33,6 +34,49 @@ const familiesByType: WeakMap<any, Family> = new WeakMap();
 // It is an array of [Family, NextType] tuples.
 let pendingUpdates: Array<[Family, any]> = [];
 
+function computeFullKey(signature: Signature): string {
+  if (signature.fullKey !== null) {
+    return signature.fullKey;
+  }
+
+  let fullKey: string = signature.ownKey;
+  let hooks;
+  try {
+    hooks = signature.getCustomHooks();
+  } catch (err) {
+    // This can happen in an edge case, e.g. if expression like Foo.useSomething
+    // depends on Foo which is lazily initialized during rendering.
+    // In that case just assume we'll have to remount.
+    signature.forceReset = true;
+    signature.fullKey = fullKey;
+    return fullKey;
+  }
+
+  for (let i = 0; i < hooks.length; i++) {
+    const hook = hooks[i];
+    if (typeof hook !== 'function') {
+      // Something's wrong. Assume we need to remount.
+      signature.forceReset = true;
+      signature.fullKey = fullKey;
+      return fullKey;
+    }
+    const nestedHookSignature = allSignaturesByType.get(hook);
+    if (nestedHookSignature === undefined) {
+      // No signature means Hook wasn't in the source code, e.g. in a library.
+      // We'll skip it because we can assume it won't change during this session.
+      continue;
+    }
+    const nestedHookKey = computeFullKey(nestedHookSignature);
+    if (nestedHookSignature.forceReset) {
+      signature.forceReset = true;
+    }
+    fullKey += '\n---\n' + nestedHookKey;
+  }
+
+  signature.fullKey = fullKey;
+  return fullKey;
+}
+
 function haveEqualSignatures(prevType, nextType) {
   const prevSignature = allSignaturesByType.get(prevType);
   const nextSignature = allSignaturesByType.get(nextType);
@@ -43,27 +87,13 @@ function haveEqualSignatures(prevType, nextType) {
   if (prevSignature === undefined || nextSignature === undefined) {
     return false;
   }
-  if (prevSignature.key !== nextSignature.key) {
+  if (computeFullKey(prevSignature) !== computeFullKey(nextSignature)) {
     return false;
   }
   if (nextSignature.forceReset) {
     return false;
   }
 
-  // TODO: we might need to calculate previous signature earlier in practice,
-  // such as during the first time a component is resolved. We'll revisit this.
-  const prevCustomHooks = prevSignature.getCustomHooks();
-  const nextCustomHooks = nextSignature.getCustomHooks();
-  if (prevCustomHooks.length !== nextCustomHooks.length) {
-    return false;
-  }
-
-  for (let i = 0; i < nextCustomHooks.length; i++) {
-    if (!haveEqualSignatures(prevCustomHooks[i], nextCustomHooks[i])) {
-      return false;
-    }
-  }
-
   return true;
 }
 
@@ -161,8 +191,18 @@ export function setSignature(
   getCustomHooks?: () => Array<Function>,
 ): void {
   allSignaturesByType.set(type, {
-    key,
     forceReset,
+    ownKey: key,
+    fullKey: null,
     getCustomHooks: getCustomHooks || (() => []),
   });
 }
+
+// This is lazily called during first render for a type.
+// It captures Hook list at that time so inline requires don't break comparisons.
+export function collectCustomHooksForSignature(type: any) {
+  const signature = allSignaturesByType.get(type);
+  if (signature !== undefined) {
+    computeFullKey(signature);
+  }
+}

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

@@ -57,12 +57,14 @@ describe('ReactFreshIntegration', () => {
     }).code;
     const exportsObj = {};
     // eslint-disable-next-line no-new-func
-    new Function('React', 'exports', '__register__', '__signature__', compiled)(
-      React,
-      exportsObj,
-      __register__,
-      __signature__,
-    );
+    new Function(
+      'global',
+      'React',
+      'exports',
+      '__register__',
+      '__signature__',
+      compiled,
+    )(global, React, exportsObj, __register__, __signature__);
     return exportsObj.default;
   }
 
@@ -85,9 +87,25 @@ describe('ReactFreshIntegration', () => {
     ReactFreshRuntime.register(type, id);
   }
 
-  function __signature__(type, key, forceReset, getCustomHooks) {
-    ReactFreshRuntime.setSignature(type, key, forceReset, getCustomHooks);
-    return type;
+  function __signature__() {
+    let call = 0;
+    let savedType;
+    let hasCustomHooks;
+    return function(type, key, forceReset, getCustomHooks) {
+      switch (call++) {
+        case 0:
+          savedType = type;
+          hasCustomHooks = typeof getCustomHooks === 'function';
+          ReactFreshRuntime.setSignature(type, key, forceReset, getCustomHooks);
+          break;
+        case 1:
+          if (hasCustomHooks) {
+            ReactFreshRuntime.collectCustomHooksForSignature(savedType);
+          }
+          break;
+      }
+      return type;
+    };
   }
 
   it('reloads function declarations', () => {
@@ -833,4 +851,303 @@ describe('ReactFreshIntegration', () => {
       expect(el.textContent).toBe('G5');
     }
   });
+
+  describe('with inline requires', () => {
+    beforeEach(() => {
+      global.FakeModuleSystem = {};
+    });
+
+    afterEach(() => {
+      delete global.FakeModuleSystem;
+    });
+
+    it('remounts component if custom hook it uses changes order on first edit', () => {
+      // This test verifies that remounting works even if calls to custom Hooks
+      // were transformed with an inline requires transform, like we have on RN.
+      // Inline requires make it harder to compare previous and next signatures
+      // because useFancyState inline require always resolves to the newest version.
+      // We're not actually using inline requires in the test, but it has similar semantics.
+      if (__DEV__) {
+        render(`
+          const FakeModuleSystem = global.FakeModuleSystem;
+
+          FakeModuleSystem.useFancyState = function(initialState) {
+            return React.useState(initialState);
+          };
+
+          const App = () => {
+            const [x, setX] = FakeModuleSystem.useFancyState('X');
+            const [y, setY] = FakeModuleSystem.useFancyState('Y');
+            return <h1>A{x}{y}</h1>;
+          };
+
+          export default App;
+        `);
+        let el = container.firstChild;
+        expect(el.textContent).toBe('AXY');
+
+        patch(`
+          const FakeModuleSystem = global.FakeModuleSystem;
+
+          FakeModuleSystem.useFancyState = function(initialState) {
+            React.useEffect(() => {});
+            return React.useState(initialState);
+          };
+
+          const App = () => {
+            const [x, setX] = FakeModuleSystem.useFancyState('X');
+            const [y, setY] = FakeModuleSystem.useFancyState('Y');
+            return <h1>B{x}{y}</h1>;
+          };
+
+          export default App;
+        `);
+        // The useFancyState Hook added an effect,
+        // so we had to remount the component.
+        expect(container.firstChild).not.toBe(el);
+        el = container.firstChild;
+        expect(el.textContent).toBe('BXY');
+
+        patch(`
+          const FakeModuleSystem = global.FakeModuleSystem;
+
+          FakeModuleSystem.useFancyState = function(initialState) {
+            React.useEffect(() => {});
+            return React.useState(initialState);
+          };
+
+          const App = () => {
+            const [x, setX] = FakeModuleSystem.useFancyState('X');
+            const [y, setY] = FakeModuleSystem.useFancyState('Y');
+            return <h1>C{x}{y}</h1>;
+          };
+
+          export default App;
+        `);
+        // We didn't change anything except the header text.
+        // So we don't expect a remount.
+        expect(container.firstChild).toBe(el);
+        expect(el.textContent).toBe('CXY');
+      }
+    });
+
+    it('remounts component if custom hook it uses changes order on second edit', () => {
+      if (__DEV__) {
+        render(`
+          const FakeModuleSystem = global.FakeModuleSystem;
+
+          FakeModuleSystem.useFancyState = function(initialState) {
+            return React.useState(initialState);
+          };
+
+          const App = () => {
+            const [x, setX] = FakeModuleSystem.useFancyState('X');
+            const [y, setY] = FakeModuleSystem.useFancyState('Y');
+            return <h1>A{x}{y}</h1>;
+          };
+
+          export default App;
+        `);
+        let el = container.firstChild;
+        expect(el.textContent).toBe('AXY');
+
+        patch(`
+          const FakeModuleSystem = global.FakeModuleSystem;
+
+          FakeModuleSystem.useFancyState = function(initialState) {
+            return React.useState(initialState);
+          };
+
+          const App = () => {
+            const [x, setX] = FakeModuleSystem.useFancyState('X');
+            const [y, setY] = FakeModuleSystem.useFancyState('Y');
+            return <h1>B{x}{y}</h1>;
+          };
+
+          export default App;
+        `);
+        expect(container.firstChild).toBe(el);
+        expect(el.textContent).toBe('BXY');
+
+        patch(`
+          const FakeModuleSystem = global.FakeModuleSystem;
+
+          FakeModuleSystem.useFancyState = function(initialState) {
+            React.useEffect(() => {});
+            return React.useState(initialState);
+          };
+
+          const App = () => {
+            const [x, setX] = FakeModuleSystem.useFancyState('X');
+            const [y, setY] = FakeModuleSystem.useFancyState('Y');
+            return <h1>C{x}{y}</h1>;
+          };
+
+          export default App;
+        `);
+        // The useFancyState Hook added an effect,
+        // so we had to remount the component.
+        expect(container.firstChild).not.toBe(el);
+        el = container.firstChild;
+        expect(el.textContent).toBe('CXY');
+
+        patch(`
+          const FakeModuleSystem = global.FakeModuleSystem;
+
+          FakeModuleSystem.useFancyState = function(initialState) {
+            React.useEffect(() => {});
+            return React.useState(initialState);
+          };
+
+          const App = () => {
+            const [x, setX] = FakeModuleSystem.useFancyState('X');
+            const [y, setY] = FakeModuleSystem.useFancyState('Y');
+            return <h1>D{x}{y}</h1>;
+          };
+
+          export default App;
+        `);
+        // We didn't change anything except the header text.
+        // So we don't expect a remount.
+        expect(container.firstChild).toBe(el);
+        expect(el.textContent).toBe('DXY');
+      }
+    });
+
+    it('recovers if evaluating Hook list throws', () => {
+      if (__DEV__) {
+        render(`
+        let FakeModuleSystem = null;
+
+        global.FakeModuleSystem.useFancyState = function(initialState) {
+          return React.useState(initialState);
+        };
+
+        const App = () => {
+          FakeModuleSystem = global.FakeModuleSystem;
+          const [x, setX] = FakeModuleSystem.useFancyState('X');
+          const [y, setY] = FakeModuleSystem.useFancyState('Y');
+          return <h1>A{x}{y}</h1>;
+        };
+
+        export default App;
+      `);
+        let el = container.firstChild;
+        expect(el.textContent).toBe('AXY');
+
+        patch(`
+        let FakeModuleSystem = null;
+
+        global.FakeModuleSystem.useFancyState = function(initialState) {
+          React.useEffect(() => {});
+          return React.useState(initialState);
+        };
+
+        const App = () => {
+          FakeModuleSystem = global.FakeModuleSystem;
+          const [x, setX] = FakeModuleSystem.useFancyState('X');
+          const [y, setY] = FakeModuleSystem.useFancyState('Y');
+          return <h1>B{x}{y}</h1>;
+        };
+
+        export default App;
+      `);
+        // We couldn't evaluate the Hook signatures
+        // so we had to remount the component.
+        expect(container.firstChild).not.toBe(el);
+        el = container.firstChild;
+        expect(el.textContent).toBe('BXY');
+      }
+    });
+
+    it('remounts component if custom hook it uses changes order behind an indirection', () => {
+      if (__DEV__) {
+        render(`
+          const FakeModuleSystem = global.FakeModuleSystem;
+
+          FakeModuleSystem.useFancyState = function(initialState) {
+            return FakeModuleSystem.useIndirection(initialState);
+          };
+
+          FakeModuleSystem.useIndirection = function(initialState) {
+            return FakeModuleSystem.useOtherIndirection(initialState);
+          };
+
+          FakeModuleSystem.useOtherIndirection = function(initialState) {
+            return React.useState(initialState);
+          };
+
+          const App = () => {
+            const [x, setX] = FakeModuleSystem.useFancyState('X');
+            const [y, setY] = FakeModuleSystem.useFancyState('Y');
+            return <h1>A{x}{y}</h1>;
+          };
+
+          export default App;
+        `);
+        let el = container.firstChild;
+        expect(el.textContent).toBe('AXY');
+
+        patch(`
+          const FakeModuleSystem = global.FakeModuleSystem;
+
+          FakeModuleSystem.useFancyState = function(initialState) {
+            return FakeModuleSystem.useIndirection(initialState);
+          };
+
+          FakeModuleSystem.useIndirection = function(initialState) {
+            return FakeModuleSystem.useOtherIndirection(initialState);
+          };
+
+          FakeModuleSystem.useOtherIndirection = function(initialState) {
+            React.useEffect(() => {});
+            return React.useState(initialState);
+          };
+
+          const App = () => {
+            const [x, setX] = FakeModuleSystem.useFancyState('X');
+            const [y, setY] = FakeModuleSystem.useFancyState('Y');
+            return <h1>B{x}{y}</h1>;
+          };
+
+          export default App;
+        `);
+
+        // The useFancyState Hook added an effect,
+        // so we had to remount the component.
+        expect(container.firstChild).not.toBe(el);
+        el = container.firstChild;
+        expect(el.textContent).toBe('BXY');
+
+        patch(`
+          const FakeModuleSystem = global.FakeModuleSystem;
+
+          FakeModuleSystem.useFancyState = function(initialState) {
+            return FakeModuleSystem.useIndirection(initialState);
+          };
+
+          FakeModuleSystem.useIndirection = function(initialState) {
+            return FakeModuleSystem.useOtherIndirection(initialState);
+          };
+
+          FakeModuleSystem.useOtherIndirection = function(initialState) {
+            React.useEffect(() => {});
+            return React.useState(initialState);
+          };
+
+          const App = () => {
+            const [x, setX] = FakeModuleSystem.useFancyState('X');
+            const [y, setY] = FakeModuleSystem.useFancyState('Y');
+            return <h1>C{x}{y}</h1>;
+          };
+
+          export default App;
+        `);
+        // We didn't change anything except the header text.
+        // So we don't expect a remount.
+        expect(container.firstChild).toBe(el);
+        expect(el.textContent).toBe('CXY');
+      }
+    });
+  });
 });

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

@@ -1,14 +1,17 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
 exports[`ReactFreshBabelPlugin generates signatures for function declarations calling hooks 1`] = `
+var _s = __signature__();
 
 export default function App() {
+  _s();
+
   const [foo, setFoo] = useState(0);
   React.useEffect(() => {});
   return <h1>{foo}</h1>;
 }
 
-__signature__(App, "useState{[foo, setFoo]}\\nuseEffect{}");
+_s(App, "useState{[foo, setFoo]}\\nuseEffect{}");
 
 _c = App;
 
@@ -18,21 +21,31 @@ __register__(_c, "App");
 `;
 
 exports[`ReactFreshBabelPlugin generates signatures for function expressions calling hooks 1`] = `
+var _s = __signature__(),
+    _s2 = __signature__();
+
+export const A = _c3 = React.memo(_c2 = React.forwardRef(_c = _s((props, ref) => {
+  _s();
 
-export const A = _c3 = React.memo(_c2 = React.forwardRef(_c = __signature__((props, ref) => {
   const [foo, setFoo] = useState(0);
   React.useEffect(() => {});
   return <h1 ref={ref}>{foo}</h1>;
 }, "useState{[foo, setFoo]}\\nuseEffect{}")));
 
-export const B = _c6 = React.memo(_c5 = React.forwardRef(_c4 = __signature__(function (props, ref) {
+export const B = _c6 = React.memo(_c5 = React.forwardRef(_c4 = _s2(function (props, ref) {
+  _s2();
+
   const [foo, setFoo] = useState(0);
   React.useEffect(() => {});
   return <h1 ref={ref}>{foo}</h1>;
 }, "useState{[foo, setFoo]}\\nuseEffect{}")));
 
 function hoc() {
-  return __signature__(function Inner() {
+  var _s3 = __signature__();
+
+  return _s3(function Inner() {
+    _s3();
+
     const [foo, setFoo] = useState(0);
     React.useEffect(() => {});
     return <h1 ref={ref}>{foo}</h1>;
@@ -57,17 +70,24 @@ __register__(_c6, "B");
 `;
 
 exports[`ReactFreshBabelPlugin generates valid signature for exotic ways to call Hooks 1`] = `
+var _s2 = __signature__();
 
 import FancyHook from 'fancy';
 
 export default function App() {
+  _s2();
+
+  var _s = __signature__();
+
   function useFancyState() {
+    _s();
+
     const [foo, setFoo] = React.useState(0);
     useFancyEffect();
     return foo;
   }
 
-  __signature__(useFancyState, 'useState{[foo, setFoo]}\\nuseFancyEffect{}', true);
+  _s(useFancyState, 'useState{[foo, setFoo]}\\nuseFancyEffect{}', true);
 
   const bar = useFancyState();
   const baz = FancyHook.useThing();
@@ -76,7 +96,7 @@ export default function App() {
   return <h1>{bar}{baz}</h1>;
 }
 
-__signature__(App, 'useFancyState{bar}\\nuseThing{baz}\\nuseState{}\\nuseThePlatform{}', true, () => [FancyHook.useThing]);
+_s2(App, 'useFancyState{bar}\\nuseThing{baz}\\nuseState{}\\nuseThePlatform{}', true, () => [FancyHook.useThing]);
 
 _c = App;
 
@@ -137,27 +157,36 @@ export default function () {}
 `;
 
 exports[`ReactFreshBabelPlugin includes custom hooks into the signatures 1`] = `
+var _s = __signature__(),
+    _s2 = __signature__(),
+    _s3 = __signature__();
 
 function useFancyState() {
+  _s();
+
   const [foo, setFoo] = React.useState(0);
   useFancyEffect();
   return foo;
 }
 
-__signature__(useFancyState, "useState{[foo, setFoo]}\\nuseFancyEffect{}", false, () => [useFancyEffect]);
+_s(useFancyState, "useState{[foo, setFoo]}\\nuseFancyEffect{}", false, () => [useFancyEffect]);
 
 const useFancyEffect = () => {
+  _s2();
+
   React.useEffect(() => {});
 };
 
-__signature__(useFancyEffect, "useEffect{}");
+_s2(useFancyEffect, "useEffect{}");
 
 export default function App() {
+  _s3();
+
   const bar = useFancyState();
   return <h1>{bar}</h1>;
 }
 
-__signature__(App, "useFancyState{bar}", false, () => [useFancyState]);
+_s3(App, "useFancyState{bar}", false, () => [useFancyState]);
 
 _c = App;