[Fresh] Support classes by force-remounting them on edit (#15801)

* Remount classes during hot reload

* Fix a crash when Hook isn't in scope inside the signature

* Minor tweaks

* Support a comment annotation to force state reset

* Refactoring: pass a function instead of WeakMap

This hides the implementation a little bit and reduces how much React knows about the underlying mechanism.

* Refactor: use forceReset to remount unknown Hooks

We already have the logic to reset a component, so let's just reuse it instead of that special case.

packages/react-fresh/src/ReactFreshBabelPlugin.js

@@ -223,11 +223,65 @@ export default function(babel) {
     };
   }
 
-  function createArgumentsForSignature(node, signature) {
+  let hasForceResetCommentByFile = new WeakMap();
+
+  // We let user do /* @hot reset */ to reset state in the whole file.
+  function hasForceResetComment(path) {
+    const file = path.hub.file;
+    let hasForceReset = hasForceResetCommentByFile.get(file);
+    if (hasForceReset !== undefined) {
+      return hasForceReset;
+    }
+
+    hasForceReset = false;
+    const comments = file.ast.comments;
+    for (let i = 0; i < comments.length; i++) {
+      const cmt = comments[i];
+      if (cmt.value.indexOf('@hot reset') !== -1) {
+        hasForceReset = true;
+        break;
+      }
+    }
+
+    hasForceResetCommentByFile.set(file, hasForceReset);
+    return hasForceReset;
+  }
+
+  function createArgumentsForSignature(node, signature, scope) {
     const {key, customHooks} = signature;
+
+    let forceReset = hasForceResetComment(scope.path);
+    let customHooksInScope = [];
+    customHooks.forEach(callee => {
+      // Check if a correponding binding exists where we emit the signature.
+      let bindingName;
+      switch (callee.type) {
+        case 'MemberExpression':
+          if (callee.object.type === 'Identifier') {
+            bindingName = callee.object.name;
+          }
+          break;
+        case 'Identifier':
+          bindingName = callee.name;
+          break;
+      }
+      if (scope.hasBinding(bindingName)) {
+        customHooksInScope.push(callee);
+      } else {
+        // We don't have anything to put in the array because Hook is out of scope.
+        // Since it could potentially have been edited, remount the component.
+        forceReset = true;
+      }
+    });
+
     const args = [node, t.stringLiteral(key)];
-    if (customHooks.length > 0) {
-      args.push(t.arrowFunctionExpression([], t.arrayExpression(customHooks)));
+    if (forceReset || customHooksInScope.length > 0) {
+      args.push(t.booleanLiteral(forceReset));
+    }
+    if (customHooksInScope.length > 0) {
+      args.push(
+        t.arrowFunctionExpression([], t.arrayExpression(customHooksInScope)),
+      );
     }
     return args;
   }
@@ -376,7 +430,11 @@ export default function(babel) {
             t.expressionStatement(
               t.callExpression(
                 t.identifier('__signature__'),
-                createArgumentsForSignature(id, signature),
+                createArgumentsForSignature(
+                  id,
+                  signature,
+                  insertAfterPath.scope,
+                ),
               ),
             ),
           );
@@ -418,7 +476,11 @@ export default function(babel) {
               t.expressionStatement(
                 t.callExpression(
                   t.identifier('__signature__'),
-                  createArgumentsForSignature(path.parent.id, signature),
+                  createArgumentsForSignature(
+                    path.parent.id,
+                    signature,
+                    insertAfterPath.scope,
+                  ),
                 ),
               ),
             );
@@ -428,7 +490,7 @@ export default function(babel) {
             path.replaceWith(
               t.callExpression(
                 t.identifier('__signature__'),
-                createArgumentsForSignature(node, signature),
+                createArgumentsForSignature(node, signature, path.scope),
               ),
             );
             // Result: let Foo = hoc(__signature(() => {}, ...))

packages/react-fresh/src/ReactFreshRuntime.js

@@ -16,6 +16,7 @@ import {REACT_MEMO_TYPE, REACT_FORWARD_REF_TYPE} from 'shared/ReactSymbols';
 
 type Signature = {|
   key: string,
+  forceReset: boolean,
   getCustomHooks: () => Array<Function>,
 |};
 
@@ -45,6 +46,9 @@ function haveEqualSignatures(prevType, nextType) {
   if (prevSignature.key !== nextSignature.key) {
     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.
@@ -63,6 +67,24 @@ function haveEqualSignatures(prevType, nextType) {
   return true;
 }
 
+function isReactClass(type) {
+  return type.prototype && type.prototype.isReactComponent;
+}
+
+function canPreserveStateBetween(prevType, nextType) {
+  if (isReactClass(prevType) || isReactClass(nextType)) {
+    return false;
+  }
+  if (haveEqualSignatures(prevType, nextType)) {
+    return true;
+  }
+  return false;
+}
+
+function resolveFamily(type) {
+  return familiesByType.get(type);
+}
+
 export function prepareUpdate(): HotUpdate {
   const staleFamilies = new Set();
   const updatedFamilies = new Set();
@@ -78,15 +100,15 @@ export function prepareUpdate(): HotUpdate {
     family.current = nextType;
 
     // Determine whether this should be a re-render or a re-mount.
-    if (haveEqualSignatures(prevType, nextType)) {
+    if (canPreserveStateBetween(prevType, nextType)) {
       updatedFamilies.add(family);
     } else {
       staleFamilies.add(family);
     }
   });
 
   return {
-    familiesByType,
+    resolveFamily,
     updatedFamilies,
     staleFamilies,
   };
@@ -135,10 +157,12 @@ export function register(type: any, id: string): void {
 export function setSignature(
   type: any,
   key: string,
+  forceReset?: boolean = false,
   getCustomHooks?: () => Array<Function>,
 ): void {
   allSignaturesByType.set(type, {
     key,
+    forceReset,
     getCustomHooks: getCustomHooks || (() => []),
   });
 }