From ab22f0d3003bd8cd2186dbb49b461b3a7d201da2 Mon Sep 17 00:00:00 2001
From: Pete Feltham <dev@felthy.com>
Date: Wed, 18 Jul 2018 19:00:02 +1000
Subject: [PATCH] Remove `componentDidUpdate`; remove `createState()` call

---
 src/injectSheet.js | 52 ++++++++++++++++++++--------------------------
 1 file changed, 23 insertions(+), 29 deletions(-)

diff --git a/src/injectSheet.js b/src/injectSheet.js
index f707ecc..dfe7952 100644
--- a/src/injectSheet.js
+++ b/src/injectSheet.js
@@ -1,21 +1,27 @@
 import defaultInjectSheet from 'react-jss'
 
-const managers = new WeakMap()
-
 if (process.env.NODE_ENV === 'production') {
   // eslint-disable-next-line no-console
   console.error('react-jss-hmr should never be used in production!')
 }
 
+const managers = new WeakMap()
+const deferredState = new WeakMap()
+
 export default function injectSheet(...rest) {
   const createHoc = defaultInjectSheet.apply(this, rest)
   return (InnerComponent) => {
     const Jss = createHoc(InnerComponent)
 
     class HotJss extends Jss {
+      constructor(props, context) {
+        super(props, context)
+        deferredState.set(this.manager, this.state)
+      }
+
       componentWillMount() {
-        // Note: this will never be called during hot module replacement, so we can
-        // use it as a place to store away the SheetManager
+        // This will never be called during hot module replacement, so we can use it as a place to
+        // store away the SheetManager
         super.componentWillMount()
         managers.set(Object.getPrototypeOf(this), this.manager)
       }
@@ -23,40 +29,28 @@ export default function injectSheet(...rest) {
       componentWillReceiveProps(nextProps, nextContext) {
         super.componentWillReceiveProps(nextProps, nextContext)
 
-        const prevManager = managers.get(Object.getPrototypeOf(this))
+        const key = Object.getPrototypeOf(this)
+        const prevManager = managers.get(key)
         const manager = this.manager
         if (prevManager !== manager) {
           // If the manager changed between mounting and updating, it was probably due to hot
           // module replacement.
           // In that case, we need to:
-          //   1. Call createState() again because although the constructor called it, RHL
-          //      restored the instance’s previous state so now the classes in the SheetsManager
-          //      do not match those in this.state
+          //   1. Restore the state created by the constructor, because RHL restored the instance’s
+          //      previous state so now the classes in the SheetsManager do not match those in
+          //      this.state
           //   2. Call this.manage() so that the new SheetsManager will attach the new sheets
-          //   3. Call unmanage() on the old SheetsManager to clean up the old sheets
-          const {theme} = this.state
-          const newState = this.createState({theme}, nextProps)
-          this.manage(newState)
-          this.setState(newState)
+          //   3. Call unmanage() on the previous SheetsManager to clean up the old sheets
+          const prevTheme = this.state.theme
+          const nextState = deferredState.get(manager)
+          this.manage(nextState)
+          this.setState(nextState)
           // the new sheets have been attached now, so we can detach the old ones
-          prevManager.unmanage(theme)
-        }
-      }
-
-      componentDidUpdate(prevProps, prevState) {
-        const key = Object.getPrototypeOf(this)
-        const prevManager = managers.get(key)
-        const manager = this.manager
-
-        const isHmrUpdate = prevManager && prevManager !== manager
-        if (isHmrUpdate) {
+          prevManager.unmanage(prevTheme)
+          // now start tracking the new manager, ready for next time
           managers.set(key, manager)
         }
-        if (!isHmrUpdate || prevState.dynamicSheet) {
-          // don't call when we don't have a previous dynamic sheet, otherwise it will throw if we
-          // added a new dynamic sheet via hmr
-          super.componentDidUpdate(prevProps, prevState)
-        }
+        deferredState.delete(manager)
       }
     }