[Fresh] Implement missing features (#15860)

* Fix existing test

This test included a change in variable name. It wasn't needed, and distracts from the actual thing being tested (the annotation).

* Reset state on edits to initial state argument

* Add a way to check whether there are hot updates

prepareUpdate() now returns null if there are none.

* Add a way to query host nodes for families

packages/react-fresh/src/ReactFreshBabelPlugin.js

@@ -175,6 +175,17 @@ export default function(babel) {
       // TODO: if there is no LHS, consider some other heuristic.
       key = hookCallPath.parentPath.get('id').getSource();
     }
+
+    // Some built-in Hooks reset on edits to arguments.
+    const args = hookCallPath.get('arguments');
+    if (hookName === 'useState' && args.length > 0) {
+      // useState second argument is initial state.
+      key += '(' + args[0].getSource() + ')';
+    } else if (hookName === 'useReducer' && args.length > 1) {
+      // useReducer second argument is initial state.
+      key += '(' + args[1].getSource() + ')';
+    }
+
     hookCallsForFn.push({
       name: hookName,
       callee: hookCallPath.node.callee,

packages/react-fresh/src/ReactFreshRuntime.js

@@ -115,7 +115,11 @@ function resolveFamily(type) {
   return familiesByType.get(type);
 }
 
-export function prepareUpdate(): HotUpdate {
+export function prepareUpdate(): HotUpdate | null {
+  if (pendingUpdates.length === 0) {
+    return null;
+  }
+
   const staleFamilies = new Set();
   const updatedFamilies = new Set();
 
@@ -206,3 +210,7 @@ export function collectCustomHooksForSignature(type: any) {
     computeFullKey(signature);
   }
 }
+
+export function getFamilyByID(id: string): Family | void {
+  return allFamiliesByID.get(id);
+}

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

@@ -20,6 +20,7 @@ let act;
 describe('ReactFresh', () => {
   let container;
   let lastRoot;
+  let findHostNodesForHotUpdate;
   let scheduleHotUpdate;
 
   beforeEach(() => {
@@ -27,6 +28,7 @@ describe('ReactFresh', () => {
       supportsFiber: true,
       inject: injected => {
         scheduleHotUpdate = injected.scheduleHotUpdate;
+        findHostNodesForHotUpdate = injected.findHostNodesForHotUpdate;
       },
       onCommitFiberRoot: (id, root) => {
         lastRoot = root;
@@ -2934,4 +2936,112 @@ describe('ReactFresh', () => {
       expect(finalEl.textContent).toBe('1');
     }
   });
+
+  it('can find host nodes for a family', () => {
+    if (__DEV__) {
+      render(() => {
+        function Child({children}) {
+          return <div className="Child">{children}</div>;
+        }
+        __register__(Child, 'Child');
+
+        function Parent({children}) {
+          return (
+            <div className="Parent">
+              <div>
+                <Child />
+              </div>
+              <div>
+                <Child />
+              </div>
+            </div>
+          );
+        }
+        __register__(Parent, 'Parent');
+
+        function App() {
+          return (
+            <div className="App">
+              <Parent />
+              <Cls>
+                <Parent />
+              </Cls>
+              <Indirection>
+                <Empty />
+              </Indirection>
+            </div>
+          );
+        }
+        __register__(App, 'App');
+
+        class Cls extends React.Component {
+          render() {
+            return this.props.children;
+          }
+        }
+
+        function Indirection({children}) {
+          return children;
+        }
+
+        function Empty() {
+          return null;
+        }
+        __register__(Empty, 'Empty');
+
+        function Frag() {
+          return (
+            <React.Fragment>
+              <div className="Frag">
+                <div />
+              </div>
+              <div className="Frag">
+                <div />
+              </div>
+            </React.Fragment>
+          );
+        }
+        __register__(Frag, 'Frag');
+
+        return App;
+      });
+
+      const parentFamily = ReactFreshRuntime.getFamilyByID('Parent');
+      const childFamily = ReactFreshRuntime.getFamilyByID('Child');
+      const emptyFamily = ReactFreshRuntime.getFamilyByID('Empty');
+
+      testFindNodesForFamilies(
+        [parentFamily],
+        container.querySelectorAll('.Parent'),
+      );
+
+      testFindNodesForFamilies(
+        [childFamily],
+        container.querySelectorAll('.Child'),
+      );
+
+      // When searching for both Parent and Child,
+      // we'll stop visual highlighting at the Parent.
+      testFindNodesForFamilies(
+        [parentFamily, childFamily],
+        container.querySelectorAll('.Parent'),
+      );
+
+      // When we can't find host nodes, use the closest parent.
+      testFindNodesForFamilies(
+        [emptyFamily],
+        container.querySelectorAll('.App'),
+      );
+    }
+  });
+
+  function testFindNodesForFamilies(families, expectedNodes) {
+    const foundNodes = Array.from(
+      findHostNodesForHotUpdate(lastRoot, families),
+    );
+    expect(foundNodes.length).toEqual(expectedNodes.length);
+    foundNodes.forEach((node, i) => {
+      expect(node).toBe(expectedNodes[i]);
+    });
+  }
 });

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

@@ -73,6 +73,8 @@ describe('ReactFreshIntegration', () => {
     act(() => {
       ReactDOM.render(<Component />, container);
     });
+    // Module initialization shouldn't be counted as a hot update.
+    expect(ReactFreshRuntime.prepareUpdate()).toBe(null);
   }
 
   function patch(source) {
@@ -275,9 +277,10 @@ describe('ReactFreshIntegration', () => {
     if (__DEV__) {
       render(`
         const {useState} = React;
+        const S = 1;
 
         export default function App() {
-          const [foo, setFoo] = useState(1);
+          const [foo, setFoo] = useState(S);
           return <h1>A{foo}</h1>;
         }
       `);
@@ -286,9 +289,10 @@ describe('ReactFreshIntegration', () => {
 
       patch(`
         const {useState} = React;
+        const S = 2;
 
         export default function App() {
-          const [foo, setFoo] = useState('ignored');
+          const [foo, setFoo] = useState(S);
           return <h1>B{foo}</h1>;
         }
       `);
@@ -298,16 +302,17 @@ describe('ReactFreshIntegration', () => {
 
       patch(`
         const {useState} = React;
+        const S = 3;
 
         export default function App() {
-          const [bar, setBar] = useState(2);
+          const [bar, setBar] = useState(S);
           return <h1>C{bar}</h1>;
         }
       `);
       // Different state variable name, so state is reset.
       expect(container.firstChild).not.toBe(el);
       const newEl = container.firstChild;
-      expect(newEl.textContent).toBe('C2');
+      expect(newEl.textContent).toBe('C3');
     }
   });
 
@@ -315,10 +320,11 @@ describe('ReactFreshIntegration', () => {
     if (__DEV__) {
       render(`
         const {useState} = React;
+        const S = 1;
 
         function hoc(Wrapped) {
           return function Generated() {
-            const [foo, setFoo] = useState(1);
+            const [foo, setFoo] = useState(S);
             return <Wrapped value={foo} />;
           };
         }
@@ -332,10 +338,11 @@ describe('ReactFreshIntegration', () => {
 
       patch(`
         const {useState} = React;
+        const S = 2;
 
         function hoc(Wrapped) {
           return function Generated() {
-            const [foo, setFoo] = useState('ignored');
+            const [foo, setFoo] = useState(S);
             return <Wrapped value={foo} />;
           };
         }
@@ -350,10 +357,11 @@ describe('ReactFreshIntegration', () => {
 
       patch(`
         const {useState} = React;
+        const S = 3;
 
         function hoc(Wrapped) {
           return function Generated() {
-            const [bar, setBar] = useState(2);
+            const [bar, setBar] = useState(S);
             return <Wrapped value={bar} />;
           };
         }
@@ -365,7 +373,7 @@ describe('ReactFreshIntegration', () => {
       // Different state variable name, so state is reset.
       expect(container.firstChild).not.toBe(el);
       const newEl = container.firstChild;
-      expect(newEl.textContent).toBe('C2');
+      expect(newEl.textContent).toBe('C3');
     }
   });
 
@@ -373,10 +381,11 @@ describe('ReactFreshIntegration', () => {
     if (__DEV__) {
       render(`
         const {useState} = React;
+        const S = 1;
 
         function hoc(Wrapped) {
           return function Generated() {
-            const [foo, setFoo] = useState(1);
+            const [foo, setFoo] = useState(S);
             return <Wrapped value={foo} />;
           };
         }
@@ -392,10 +401,11 @@ describe('ReactFreshIntegration', () => {
 
       patch(`
         const {useState} = React;
+        const S = 2;
 
         function hoc(Wrapped) {
           return function Generated() {
-            const [foo, setFoo] = useState('ignored');
+            const [foo, setFoo] = useState(S);
             return <Wrapped value={foo} />;
           };
         }
@@ -412,10 +422,11 @@ describe('ReactFreshIntegration', () => {
 
       patch(`
         const {useState} = React;
+        const S = 3;
 
         function hoc(Wrapped) {
           return function Generated() {
-            const [bar, setBar] = useState(2);
+            const [bar, setBar] = useState(S);
             return <Wrapped value={bar} />;
           };
         }
@@ -429,7 +440,7 @@ describe('ReactFreshIntegration', () => {
       // Different state variable name, so state is reset.
       expect(container.firstChild).not.toBe(el);
       const newEl = container.firstChild;
-      expect(newEl.textContent).toBe('C2');
+      expect(newEl.textContent).toBe('C3');
     }
   });
 
@@ -754,6 +765,114 @@ describe('ReactFreshIntegration', () => {
   });
 
   it('resets state on every edit with @hot reset annotation', () => {
+    if (__DEV__) {
+      render(`
+        const {useState} = React;
+        const S = 1;
+
+        export default function App() {
+          const [foo, setFoo] = useState(S);
+          return <h1>A{foo}</h1>;
+        }
+      `);
+      let el = container.firstChild;
+      expect(el.textContent).toBe('A1');
+
+      patch(`
+        const {useState} = React;
+        const S = 2;
+
+        export default function App() {
+          const [foo, setFoo] = useState(S);
+          return <h1>B{foo}</h1>;
+        }
+      `);
+      // Same state variable name, so state is preserved.
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('B1');
+
+      patch(`
+        const {useState} = React;
+        const S = 3;
+
+        /* @hot reset */
+
+        export default function App() {
+          const [foo, setFoo] = useState(S);
+          return <h1>C{foo}</h1>;
+        }
+      `);
+      // Found remount annotation, so state is reset.
+      expect(container.firstChild).not.toBe(el);
+      el = container.firstChild;
+      expect(el.textContent).toBe('C3');
+
+      patch(`
+        const {useState} = React;
+        const S = 4;
+
+        export default function App() {
+
+          // @hot reset
+
+          const [foo, setFoo] = useState(S);
+          return <h1>D{foo}</h1>;
+        }
+      `);
+      // Found remount annotation, so state is reset.
+      expect(container.firstChild).not.toBe(el);
+      el = container.firstChild;
+      expect(el.textContent).toBe('D4');
+
+      patch(`
+        const {useState} = React;
+        const S = 5;
+
+        export default function App() {
+          const [foo, setFoo] = useState(S);
+          return <h1>E{foo}</h1>;
+        }
+      `);
+      // There is no remount annotation anymore,
+      // so preserve the previous state.
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('E4');
+
+      patch(`
+        const {useState} = React;
+        const S = 6;
+
+        export default function App() {
+          const [foo, setFoo] = useState(S);
+          return <h1>F{foo}</h1>;
+        }
+      `);
+      // Continue editing.
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('F4');
+
+      patch(`
+        const {useState} = React;
+        const S = 7;
+
+        export default function App() {
+
+          /* @hot reset */
+
+          const [foo, setFoo] = useState(S);
+          return <h1>G{foo}</h1>;
+        }
+      `);
+      // Force remount one last time.
+      expect(container.firstChild).not.toBe(el);
+      el = container.firstChild;
+      expect(el.textContent).toBe('G7');
+    }
+  });
+
+  // This is best effort for simple cases.
+  // We won't attempt to resolve identifiers.
+  it('resets state when useState initial state is edited', () => {
     if (__DEV__) {
       render(`
         const {useState} = React;
@@ -770,85 +889,68 @@ describe('ReactFreshIntegration', () => {
         const {useState} = React;
 
         export default function App() {
-          const [foo, setFoo] = useState('ignored');
+          const [foo, setFoo] = useState(1);
           return <h1>B{foo}</h1>;
         }
       `);
-      // Same state variable name, so state is preserved.
+      // Same initial state, so it's preserved.
       expect(container.firstChild).toBe(el);
       expect(el.textContent).toBe('B1');
 
       patch(`
         const {useState} = React;
 
-        /* @hot reset */
-
         export default function App() {
-          const [bar, setBar] = useState(2);
-          return <h1>C{bar}</h1>;
+          const [foo, setFoo] = useState(2);
+          return <h1>C{foo}</h1>;
         }
       `);
-      // Found remount annotation, so state is reset.
+      // Different initial state, so state is reset.
       expect(container.firstChild).not.toBe(el);
       el = container.firstChild;
       expect(el.textContent).toBe('C2');
+    }
+  });
 
-      patch(`
-        const {useState} = React;
+  // This is best effort for simple cases.
+  // We won't attempt to resolve identifiers.
+  it('resets state when useReducer initial state is edited', () => {
+    if (__DEV__) {
+      render(`
+        const {useReducer} = React;
 
         export default function App() {
-
-          // @hot reset
-
-          const [bar, setBar] = useState(3);
-          return <h1>D{bar}</h1>;
+          const [foo, setFoo] = useReducer(x => x, 1);
+          return <h1>A{foo}</h1>;
         }
       `);
-      // Found remount annotation, so state is reset.
+      let el = container.firstChild;
+      expect(el.textContent).toBe('A1');
+
+      patch(`
+        const {useReducer} = React;
+
+        export default function App() {
+          const [foo, setFoo] = useReducer(x => x, 1);
+          return <h1>B{foo}</h1>;
+        }
+      `);
+      // Same initial state, so it's preserved.
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('B1');
+
+      patch(`
+        const {useReducer} = React;
+
+        export default function App() {
+          const [foo, setFoo] = useReducer(x => x, 2);
+          return <h1>C{foo}</h1>;
+        }
+      `);
+      // Different initial state, so state is reset.
       expect(container.firstChild).not.toBe(el);
       el = container.firstChild;
-      expect(el.textContent).toBe('D3');
-
-      patch(`
-        const {useState} = React;
-
-        export default function App() {
-          const [bar, setBar] = useState(4);
-          return <h1>E{bar}</h1>;
-        }
-      `);
-      // There is no remount annotation anymore,
-      // so preserve the previous state.
-      expect(container.firstChild).toBe(el);
-      expect(el.textContent).toBe('E3');
-
-      patch(`
-        const {useState} = React;
-
-        export default function App() {
-          const [bar, setBar] = useState(4);
-          return <h1>F{bar}</h1>;
-        }
-      `);
-      // Continue editing.
-      expect(container.firstChild).toBe(el);
-      expect(el.textContent).toBe('F3');
-
-      patch(`
-        const {useState} = React;
-
-        export default function App() {
-
-          /* @hot reset */
-
-          const [bar, setBar] = useState(5);
-          return <h1>G{bar}</h1>;
-        }
-      `);
-      // Force remount one last time.
-      expect(container.firstChild).not.toBe(el);
-      el = container.firstChild;
-      expect(el.textContent).toBe('G5');
+      expect(el.textContent).toBe('C2');
     }
   });
 

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

@@ -11,7 +11,7 @@ export default function App() {
   return <h1>{foo}</h1>;
 }
 
-_s(App, "useState{[foo, setFoo]}\\nuseEffect{}");
+_s(App, "useState{[foo, setFoo](0)}\\nuseEffect{}");
 
 _c = App;
 
@@ -30,7 +30,7 @@ export const A = _c3 = React.memo(_c2 = React.forwardRef(_c = _s((props, ref) =>
   const [foo, setFoo] = useState(0);
   React.useEffect(() => {});
   return <h1 ref={ref}>{foo}</h1>;
-}, "useState{[foo, setFoo]}\\nuseEffect{}")));
+}, "useState{[foo, setFoo](0)}\\nuseEffect{}")));
 
 export const B = _c6 = React.memo(_c5 = React.forwardRef(_c4 = _s2(function (props, ref) {
   _s2();
@@ -38,7 +38,7 @@ export const B = _c6 = React.memo(_c5 = React.forwardRef(_c4 = _s2(function (pro
   const [foo, setFoo] = useState(0);
   React.useEffect(() => {});
   return <h1 ref={ref}>{foo}</h1>;
-}, "useState{[foo, setFoo]}\\nuseEffect{}")));
+}, "useState{[foo, setFoo](0)}\\nuseEffect{}")));
 
 function hoc() {
   var _s3 = __signature__();
@@ -49,7 +49,7 @@ function hoc() {
     const [foo, setFoo] = useState(0);
     React.useEffect(() => {});
     return <h1 ref={ref}>{foo}</h1>;
-  }, "useState{[foo, setFoo]}\\nuseEffect{}");
+  }, "useState{[foo, setFoo](0)}\\nuseEffect{}");
 }
 
 export let C = hoc();
@@ -87,7 +87,7 @@ export default function App() {
     return foo;
   }
 
-  _s(useFancyState, 'useState{[foo, setFoo]}\\nuseFancyEffect{}', true);
+  _s(useFancyState, 'useState{[foo, setFoo](0)}\\nuseFancyEffect{}', true);
 
   const bar = useFancyState();
   const baz = FancyHook.useThing();
@@ -169,7 +169,7 @@ function useFancyState() {
   return foo;
 }
 
-_s(useFancyState, "useState{[foo, setFoo]}\\nuseFancyEffect{}", false, () => [useFancyEffect]);
+_s(useFancyState, "useState{[foo, setFoo](0)}\\nuseFancyEffect{}", false, () => [useFancyEffect]);
 
 const useFancyEffect = () => {
   _s2();

packages/react-reconciler/src/ReactFiberHotReloading.js

@@ -10,6 +10,7 @@
 import type {ReactElement} from 'shared/ReactElementType';
 import type {Fiber} from './ReactFiber';
 import type {FiberRoot} from './ReactFiberRoot';
+import type {Instance} from './ReactFiberHostConfig';
 
 import {
   flushSync,
@@ -21,6 +22,9 @@ import {
   ClassComponent,
   FunctionComponent,
   ForwardRef,
+  HostComponent,
+  HostPortal,
+  HostRoot,
   MemoComponent,
   SimpleMemoComponent,
 } from 'shared/ReactWorkTags';
@@ -287,3 +291,134 @@ function scheduleFibersWithFamiliesRecursively(
     }
   }
 }
+
+export function findHostNodesForHotUpdate(
+  root: FiberRoot,
+  families: Array<Family>,
+): Set<Instance> {
+  if (__DEV__) {
+    const hostNodes = new Set();
+    const types = new Set(families.map(family => family.current));
+    findHostNodesForMatchingFibersRecursively(root.current, types, hostNodes);
+    return hostNodes;
+  } else {
+    throw new Error(
+      'Did not expect findHostNodesForHotUpdate to be called in production.',
+    );
+  }
+}
+
+function findHostNodesForMatchingFibersRecursively(
+  fiber: Fiber,
+  types: Set<any>,
+  hostNodes: Set<Instance>,
+) {
+  if (__DEV__) {
+    const {child, sibling, tag, type} = fiber;
+
+    let candidateType = null;
+    switch (tag) {
+      case FunctionComponent:
+      case SimpleMemoComponent:
+      case ClassComponent:
+        candidateType = type;
+        break;
+      case ForwardRef:
+        candidateType = type.render;
+        break;
+      default:
+        break;
+    }
+
+    let didMatch = false;
+    if (candidateType !== null) {
+      if (types.has(candidateType)) {
+        didMatch = true;
+      }
+    }
+
+    if (didMatch) {
+      // We have a match. This only drills down to the closest host components.
+      // There's no need to search deeper because for the purpose of giving
+      // visual feedback, "flashing" outermost parent rectangles is sufficient.
+      findHostNodesForFiberShallowly(fiber, hostNodes);
+    } else {
+      // If there's no match, maybe there will be one further down in the child tree.
+      if (child !== null) {
+        findHostNodesForMatchingFibersRecursively(child, types, hostNodes);
+      }
+    }
+
+    if (sibling !== null) {
+      findHostNodesForMatchingFibersRecursively(sibling, types, hostNodes);
+    }
+  }
+}
+
+function findHostNodesForFiberShallowly(
+  fiber: Fiber,
+  hostNodes: Set<Instance>,
+): void {
+  if (__DEV__) {
+    const foundHostNodes = findChildHostNodesForFiberShallowly(
+      fiber,
+      hostNodes,
+    );
+    if (foundHostNodes) {
+      return;
+    }
+    // If we didn't find any host children, fallback to closest host parent.
+    let node = fiber;
+    while (true) {
+      switch (node.tag) {
+        case HostComponent:
+          hostNodes.add(node.stateNode);
+          return;
+        case HostPortal:
+          hostNodes.add(node.stateNode.containerInfo);
+          return;
+        case HostRoot:
+          hostNodes.add(node.stateNode.containerInfo);
+          return;
+      }
+      if (node.return === null) {
+        throw new Error('Expected to reach root first.');
+      }
+      node = node.return;
+    }
+  }
+}
+
+function findChildHostNodesForFiberShallowly(
+  fiber: Fiber,
+  hostNodes: Set<Instance>,
+): boolean {
+  if (__DEV__) {
+    let node: Fiber = fiber;
+    let foundHostNodes = false;
+    while (true) {
+      if (node.tag === HostComponent) {
+        // We got a match.
+        foundHostNodes = true;
+        hostNodes.add(node.stateNode);
+        // There may still be more, so keep searching.
+      } else if (node.child !== null) {
+        node.child.return = node;
+        node = node.child;
+        continue;
+      }
+      if (node === fiber) {
+        return foundHostNodes;
+      }
+      while (node.sibling === null) {
+        if (node.return === null || node.return === fiber) {
+          return foundHostNodes;
+        }
+        node = node.return;
+      }
+      node.sibling.return = node.return;
+      node = node.sibling;
+    }
+  }
+  return false;
+}

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -70,7 +70,10 @@ import {StrictMode} from './ReactTypeOfMode';
 import {Sync} from './ReactFiberExpirationTime';
 import {revertPassiveEffectsChange} from 'shared/ReactFeatureFlags';
 import {requestCurrentSuspenseConfig} from './ReactFiberSuspenseConfig';
-import {scheduleHotUpdate} from './ReactFiberHotReloading';
+import {
+  scheduleHotUpdate,
+  findHostNodesForHotUpdate,
+} from './ReactFiberHotReloading';
 
 type OpaqueRoot = FiberRoot;
 
@@ -472,6 +475,7 @@ export function injectIntoDevTools(devToolsConfig: DevToolsConfig): boolean {
 
   return injectInternals({
     ...devToolsConfig,
+    findHostNodesForHotUpdate: __DEV__ ? findHostNodesForHotUpdate : null,
     scheduleHotUpdate: __DEV__ ? scheduleHotUpdate : null,
     overrideHookState,
     overrideProps,