From 1b9f2da08b1436b11d23009486b6a513bdfacf75 Mon Sep 17 00:00:00 2001
From: Anton Korzunov <thekashey@gmail.com>
Date: Thu, 4 Jul 2019 19:59:49 +1000
Subject: [PATCH] fix: make type comparison stronger

---
 src/internal/reactUtils.js            | 30 +++++++++++++++
 src/reconciler/componentComparator.js | 55 ++++++++++++++++-----------
 2 files changed, 63 insertions(+), 22 deletions(-)

diff --git a/src/internal/reactUtils.js b/src/internal/reactUtils.js
index fed2e844b..a6fa68618 100644
--- a/src/internal/reactUtils.js
+++ b/src/internal/reactUtils.js
@@ -66,4 +66,34 @@ export const isForwardType = ({ type }) =>
   type && typeof type === 'object' && '$$typeof' in type && type.$$typeof === ForwardType && ForwardType;
 export const isContextType = type => isContextConsumer(type) || isContextProvider(type);
 
+export const getElementType = type => {
+  const element = { type };
+
+  if (isContextConsumer(element)) {
+    return 'Consumer';
+  }
+  if (isContextProvider(element)) {
+    return 'Provider';
+  }
+  if (isLazyType(element)) {
+    return 'Lazy';
+  }
+  if (isMemoType(element)) {
+    return 'Memo';
+  }
+  if (isForwardType(element)) {
+    return 'Forward';
+  }
+
+  if (isReactClass(type)) {
+    return 'Class';
+  }
+
+  if (typeof element === 'function') {
+    return 'FC';
+  }
+
+  return 'unknown';
+};
+
 export const getContextProvider = type => type && type._context;
diff --git a/src/reconciler/componentComparator.js b/src/reconciler/componentComparator.js
index 04601b886..a48ed74e1 100644
--- a/src/reconciler/componentComparator.js
+++ b/src/reconciler/componentComparator.js
@@ -1,6 +1,7 @@
 import { getIdByType, getProxyByType, getSignature, isColdType, updateProxyById } from './proxies';
 import { hotComparisonOpen } from '../global/generation';
 import {
+  getElementType,
   isContextType,
   isForwardType,
   isLazyType,
@@ -20,31 +21,36 @@ const getInnerComponentType = component => {
 };
 
 function haveEqualSignatures(prevType, nextType) {
-  const prevSignature = getSignature(prevType);
-  const nextSignature = getSignature(nextType);
+  try {
+    const prevSignature = getSignature(prevType);
+    const nextSignature = getSignature(nextType);
 
-  if (prevSignature === undefined && nextSignature === undefined) {
-    return true;
-  }
-  if (prevSignature === undefined || nextSignature === undefined) {
-    return false;
-  }
-  if (prevSignature.key !== nextSignature.key) {
-    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.
-  const prevCustomHooks = prevSignature.getCustomHooks();
-  const nextCustomHooks = nextSignature.getCustomHooks();
-  if (prevCustomHooks.length !== nextCustomHooks.length) {
-    return false;
-  }
+    if (prevSignature === undefined && nextSignature === undefined) {
+      return true;
+    }
+    if (prevSignature === undefined || nextSignature === undefined) {
+      return false;
+    }
+    if (prevSignature.key !== nextSignature.key) {
+      return false;
+    }
 
-  for (let i = 0; i < nextCustomHooks.length; i++) {
-    if (!haveEqualSignatures(prevCustomHooks[i], nextCustomHooks[i])) {
+    // 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.
+    const prevCustomHooks = prevSignature.getCustomHooks();
+    const nextCustomHooks = nextSignature.getCustomHooks();
+    if (prevCustomHooks.length !== nextCustomHooks.length) {
       return false;
     }
+
+    for (let i = 0; i < nextCustomHooks.length; i++) {
+      if (!haveEqualSignatures(prevCustomHooks[i], nextCustomHooks[i])) {
+        return false;
+      }
+    }
+  } catch (e) {
+    logger.error('React-Hot-Loader: error occurred while comparing hook signature', e);
+    return false;
   }
 
   return true;
@@ -99,7 +105,12 @@ const areDeepSwappable = (oldType, newType) => {
 const compareComponents = (oldType, newType, setNewType, baseType) => {
   let defaultResult = oldType === newType;
 
-  if ((oldType && !newType) || (!oldType && newType) || typeof oldType !== typeof newType) {
+  if (
+    (oldType && !newType) ||
+    (!oldType && newType) ||
+    typeof oldType !== typeof newType ||
+    getElementType(oldType) !== getElementType(newType)
+  ) {
     return defaultResult;
   }