[Fresh] Fix edge case with early function call (#17824)

Author: gaearon

packages/react-refresh/src/ReactFreshRuntime.js

@@ -601,9 +601,14 @@ export function _getMountedRootCount() {
 //   'useState{[foo, setFoo]}(0)',
 //   () => [useCustomHook], /* Lazy to avoid triggering inline requires */
 // );
+type SignatureStatus = 'needsSignature' | 'needsCustomHooks' | 'resolved';
 export function createSignatureFunctionForTransform() {
   if (__DEV__) {
-    let call = 0;
+    // We'll fill in the signature in two steps.
+    // First, we'll know the signature itself. This happens outside the component.
+    // Then, we'll know the references to custom Hooks. This happens inside the component.
+    // After that, the returned function will be a fast path no-op.
+    let status: SignatureStatus = 'needsSignature';
     let savedType;
     let hasCustomHooks;
     return function<T>(
@@ -612,16 +617,25 @@ export function createSignatureFunctionForTransform() {
       forceReset?: boolean,
       getCustomHooks?: () => Array<Function>,
     ): T {
-      switch (call++) {
-        case 0:
-          savedType = type;
-          hasCustomHooks = typeof getCustomHooks === 'function';
-          setSignature(type, key, forceReset, getCustomHooks);
+      switch (status) {
+        case 'needsSignature':
+          if (type !== undefined) {
+            // If we received an argument, this is the initial registration call.
+            savedType = type;
+            hasCustomHooks = typeof getCustomHooks === 'function';
+            setSignature(type, key, forceReset, getCustomHooks);
+            // The next call we expect is from inside a function, to fill in the custom Hooks.
+            status = 'needsCustomHooks';
+          }
           break;
-        case 1:
+        case 'needsCustomHooks':
           if (hasCustomHooks) {
             collectCustomHooksForSignature(savedType);
           }
+          status = 'resolved';
+          break;
+        case 'resolved':
+          // Do nothing. Fast path for all future renders.
           break;
       }
       return type;

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

@@ -605,6 +605,86 @@ describe('ReactFreshIntegration', () => {
       }
     });
 
+    it('does not get confused when component is called early', () => {
+      if (__DEV__) {
+        render(`
+          // This isn't really a valid pattern but it's close enough
+          // to simulate what happens when you call ReactDOM.render
+          // in the same file. We want to ensure this doesn't confuse
+          // the runtime.
+          App();
+
+          function App() {
+            const [x, setX] = useFancyState('X');
+            const [y, setY] = useFancyState('Y');
+            return <h1>A{x}{y}</h1>;
+          };
+
+          function useFancyState(initialState) {
+            // No real Hook calls to avoid triggering invalid call invariant.
+            // We only want to verify that we can still call this function early.
+            return initialState;
+          }
+
+          export default App;
+        `);
+        let el = container.firstChild;
+        expect(el.textContent).toBe('AXY');
+
+        patch(`
+          // This isn't really a valid pattern but it's close enough
+          // to simulate what happens when you call ReactDOM.render
+          // in the same file. We want to ensure this doesn't confuse
+          // the runtime.
+          App();
+
+          function App() {
+            const [x, setX] = useFancyState('X');
+            const [y, setY] = useFancyState('Y');
+            return <h1>B{x}{y}</h1>;
+          };
+
+          function useFancyState(initialState) {
+            // No real Hook calls to avoid triggering invalid call invariant.
+            // We only want to verify that we can still call this function early.
+            return initialState;
+          }
+
+          export default App;
+        `);
+        // Same state variables, so no remount.
+        expect(container.firstChild).toBe(el);
+        expect(el.textContent).toBe('BXY');
+
+        patch(`
+          // This isn't really a valid pattern but it's close enough
+          // to simulate what happens when you call ReactDOM.render
+          // in the same file. We want to ensure this doesn't confuse
+          // the runtime.
+          App();
+
+          function App() {
+            const [y, setY] = useFancyState('Y');
+            const [x, setX] = useFancyState('X');
+            return <h1>B{x}{y}</h1>;
+          };
+
+          function useFancyState(initialState) {
+            // No real Hook calls to avoid triggering invalid call invariant.
+            // We only want to verify that we can still call this function early.
+            return initialState;
+          }
+
+          export default App;
+        `);
+        // Hooks were re-ordered. This causes a remount.
+        // Therefore, Hook calls don't accidentally share state.
+        expect(container.firstChild).not.toBe(el);
+        el = container.firstChild;
+        expect(el.textContent).toBe('BXY');
+      }
+    });
+
     it('does not get confused by Hooks defined inline', () => {
       // This is not a recommended pattern but at least it shouldn't break.
       if (__DEV__) {