feat: support hooks update on HMR, fixes #1256

Author: theKashey

README.md

@@ -53,6 +53,22 @@ export default hot(App);
 
 4.  If you need hooks support, use React-🔥-Dom
 
+### Hook support
+
+To enable hot hook update you have to enable it (not setting by default for now):
+
+```js
+import { setConfig } from 'react-hot-loader';
+
+setConfig({
+  hotHooks: true,
+});
+```
+
+With this option set **all** `useEffects`, `useCallbacks` and `useMemo` would be updated on Hot Module Replacement.
+
+Please try it so we can enable it by default. (or not)
+
 ## React-🔥-Dom
 
 React-🔥-Dom ([hot-loader/react-dom](https://github.com/hot-loader/react-dom)) replaces the "react-dom" package of the same version, but with additional patches to support hot reloading.

src/configuration.js

@@ -11,6 +11,9 @@ const configuration = {
   // Allows SFC to be used, enables "intermediate" components used by Relay, should be disabled for Preact
   allowSFC: true,
 
+  // Allow hot reload of effect hooks
+  hotHooks: false,
+
   // Disable "hot-replacement-render"
   disableHotRenderer: false,
 

src/global/generation.js

@@ -1,13 +1,20 @@
 import { forEachKnownClass } from '../proxy/createClassProxy';
 
+// this counter tracks `register` invocations.
+// works good, but code splitting is breaking it
 let generation = 1;
+
+// these counters are aimed to mitigate the "first render"
 let hotComparisonCounter = 0;
 let hotComparisonRuns = 0;
 const nullFunction = () => ({});
+
+// these callbacks would be called on component update
 let onHotComparisonOpen = nullFunction;
 let onHotComparisonElement = nullFunction;
 let onHotComparisonClose = nullFunction;
 
+// inversion of control
 export const setComparisonHooks = (open, element, close) => {
   onHotComparisonOpen = open;
   onHotComparisonElement = element;
@@ -43,12 +50,19 @@ export const configureGeneration = (counter, runs) => {
   hotComparisonRuns = runs;
 };
 
+// TODO: shall it be called from incrementHotGeneration?
 export const enterHotUpdate = () => {
   Promise.resolve(incrementHot()).then(() => setTimeout(decrementHot, 0));
 };
 
+// TODO: deprecate?
 export const increment = () => {
   enterHotUpdate();
   return generation++;
 };
 export const get = () => generation;
+
+// These counters tracks HMR generations, and probably should be used instead of the old one
+let hotReplacementGeneration = 0;
+export const incrementHotGeneration = () => hotReplacementGeneration++;
+export const getHotGeneration = () => hotReplacementGeneration;

src/reactHotLoader.js

@@ -7,7 +7,7 @@ import {
   isForwardType,
   isContextType,
 } from './internal/reactUtils';
-import { increment as incrementGeneration } from './global/generation';
+import { increment as incrementGeneration, getHotGeneration } from './global/generation';
 import {
   updateProxyById,
   resetProxies,
@@ -26,6 +26,13 @@ import { hotComponentCompare } from './reconciler/componentComparator';
 
 const forceSimpleSFC = { proxy: { pureSFC: true } };
 
+const hookWrapper = hook => (cb, deps) => {
+  if (configuration.hotHooks) {
+    return hook(cb, deps ? [...deps, getHotGeneration()] : deps);
+  }
+  return hook(cb, deps);
+};
+
 const reactHotLoader = {
   IS_REACT_MERGE_ENABLED: false,
   register(type, uniqueLocalName, fileName, options = {}) {
@@ -75,6 +82,7 @@ const reactHotLoader = {
       incrementGeneration();
     }
     if (isForwardType({ type })) {
+      reactHotLoader.register(type.render, `${uniqueLocalName}:render`, fileName, forceSimpleSFC);
       updateFunctionProxyById(id, type, updateForward);
       incrementGeneration();
     }
@@ -158,6 +166,13 @@ const reactHotLoader = {
       React.Children.only.isPatchedByReactHotLoader = true;
     }
 
+    if (React.useEffect && !React.useState.isPatchedByReactHotLoader) {
+      React.useEffect = hookWrapper(React.useEffect);
+      React.useLayoutEffect = hookWrapper(React.useLayoutEffect);
+      React.useCallback = hookWrapper(React.useCallback);
+      React.useMemo = hookWrapper(React.useMemo);
+    }
+
     // reactHotLoader.reset()
   },
 };

src/reconciler/proxies.js

@@ -2,6 +2,7 @@ import createProxy, { PROXY_KEY } from '../proxy';
 import { resetClassProxies } from '../proxy/createClassProxy';
 import { isCompositeComponent, isReactClass } from '../internal/reactUtils';
 import configuration from '../configuration';
+import { incrementHotGeneration } from '../global/generation';
 
 const merge = require('lodash/merge');
 
@@ -58,6 +59,8 @@ export const updateProxyById = (id, type, options = {}) => {
     );
   } else {
     proxiesByID[id].update(type);
+    // proxy could be registered again only in case of HMR
+    incrementHotGeneration();
   }
   return proxiesByID[id];
 };