[Fresh] Implement missing features

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,13 +20,15 @@ let act;
 describe('ReactFresh', () => {
   let container;
   let lastRoot;
+  let findHostNodesForHotUpdate;
   let scheduleHotUpdate;
 
   beforeEach(() => {
     global.__REACT_DEVTOOLS_GLOBAL_HOOK__ = {
       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]);
+    });
+  }
 });