From 10ed61dd9e82ad00bfa31872ffc140dfeebe9d42 Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@thesnail.org>
Date: Thu, 28 Dec 2017 11:33:41 +1100
Subject: [PATCH] Use store revisions to ensure that stories re-render on HMR.

For #2587
---
 app/react/src/client/preview/render.js     | 27 +++++++++++++++-------
 lib/core/src/client/preview/client_api.js  |  1 +
 lib/core/src/client/preview/story_store.js | 14 ++++++++++-
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/app/react/src/client/preview/render.js b/app/react/src/client/preview/render.js
index 4795dffc5ede..b53dbd4c68cb 100644
--- a/app/react/src/client/preview/render.js
+++ b/app/react/src/client/preview/render.js
@@ -14,6 +14,7 @@ const logger = console;
 let rootEl = null;
 let previousKind = '';
 let previousStory = '';
+let previousRevision = -1;
 
 if (isBrowser) {
   rootEl = document.getElementById('root');
@@ -46,6 +47,7 @@ export function renderMain(data, storyStore) {
   const noPreview = <NoPreview />;
   const { selectedKind, selectedStory } = data;
 
+  const revision = storyStore.getRevision();
   const story = storyStore.getStory(selectedKind, selectedStory);
   if (!story) {
     ReactDOM.render(noPreview, rootEl);
@@ -56,16 +58,25 @@ export function renderMain(data, storyStore) {
   // renderMain() gets executed after each action. Actions will cause the whole
   // story to re-render without this check.
   //    https://github.com/storybooks/react-storybook/issues/116
-  if (selectedKind !== previousKind || previousStory !== selectedStory) {
-    // We need to unmount the existing set of components in the DOM node.
-    // Otherwise, React may not recrease instances for every story run.
-    // This could leads to issues like below:
-    //    https://github.com/storybooks/react-storybook/issues/81
-    previousKind = selectedKind;
-    previousStory = selectedStory;
-    ReactDOM.unmountComponentAtNode(rootEl);
+  // However, we do want the story to re-render if the store itself has changed
+  // (which happens at the moment when HMR occurs)
+  if (
+    revision === previousRevision &&
+    selectedKind === previousKind &&
+    previousStory === selectedStory
+  ) {
+    return null;
   }
 
+  // We need to unmount the existing set of components in the DOM node.
+  // Otherwise, React may not recrease instances for every story run.
+  // This could leads to issues like below:
+  //    https://github.com/storybooks/react-storybook/issues/81
+  previousRevision = revision;
+  previousKind = selectedKind;
+  previousStory = selectedStory;
+  ReactDOM.unmountComponentAtNode(rootEl);
+
   const context = {
     kind: selectedKind,
     story: selectedStory,
diff --git a/lib/core/src/client/preview/client_api.js b/lib/core/src/client/preview/client_api.js
index 73e074b345ab..f40c68e56562 100644
--- a/lib/core/src/client/preview/client_api.js
+++ b/lib/core/src/client/preview/client_api.js
@@ -47,6 +47,7 @@ export default class ClientApi {
     if (m && m.hot) {
       m.hot.dispose(() => {
         this._storyStore.removeStoryKind(kind);
+        this._storyStore.incrementRevision();
       });
     }
 
diff --git a/lib/core/src/client/preview/story_store.js b/lib/core/src/client/preview/story_store.js
index 0f7fcff98705..33424250ccfe 100644
--- a/lib/core/src/client/preview/story_store.js
+++ b/lib/core/src/client/preview/story_store.js
@@ -12,6 +12,15 @@ export default class StoryStore extends EventEmitter {
   constructor() {
     super();
     this._data = {};
+    this._revision = 0;
+  }
+
+  getRevision() {
+    return this._revision;
+  }
+
+  incrementRevision() {
+    this._revision += 1;
   }
 
   addStory(kind, name, fn, fileName) {
@@ -88,7 +97,10 @@ export default class StoryStore extends EventEmitter {
   }
 
   dumpStoryBook() {
-    const data = this.getStoryKinds().map(kind => ({ kind, stories: this.getStories(kind) }));
+    const data = this.getStoryKinds().map(kind => ({
+      kind,
+      stories: this.getStories(kind),
+    }));
 
     return data;
   }