From 3e38810a495b6ba0e10f686f7691fd6dc94be3be Mon Sep 17 00:00:00 2001
From: eight <eight04@gmail.com>
Date: Thu, 11 Oct 2018 18:13:24 +0800
Subject: [PATCH] Fix: make sure icons are refreshed at startup

---
 background/background.js | 42 +++++++++++++++++++++++++++++++++-------
 content/apply.js         | 11 ++++++++---
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/background/background.js b/background/background.js
index bdb2232f..57e3e281 100644
--- a/background/background.js
+++ b/background/background.js
@@ -142,11 +142,27 @@ prefs.subscribe([
 ], () => debounce(refreshIconBadgeColor));
 
 prefs.subscribe([
-  'show-badge',
+  'show-badge'
+], () => debounce(refreshIconBadgeText));
+
+prefs.subscribe([
+  'disableAll',
   'iconset',
 ], () => debounce(refreshAllIcons));
 
-prefs.initializing.then(refreshIconBadgeColor);
+prefs.initializing.then(() => {
+  refreshIconBadgeColor();
+  refreshAllIconsBadgeText();
+  refreshAllIcons();
+});
+
+navigatorUtil.onUrlChange(({tabId, frameId}, type) => {
+  if (type === 'committed' && !frameId) {
+    // it seems that the tab icon would be reset when pressing F5. We
+    // invalidate the cache here so it would be refreshed.
+    tabIcons.delete(tabId);
+  }
+});
 
 // *************************************************************************
 chrome.runtime.onInstalled.addListener(({reason}) => {
@@ -348,16 +364,21 @@ function updateIconBadge(tabId, count) {
   if (tabIcon.count === count) {
     return;
   }
+  const oldCount = tabIcon.count;
   tabIcon.count = count;
-  iconUtil.setBadgeText({
-    text: prefs.get('show-badge') && count ? String(count) : '',
-    tabId
-  });
-  if (!count) {
+  refreshIconBadgeText(tabId, tabIcon);
+  if (Boolean(oldCount) !== Boolean(count)) {
     refreshIcon(tabId, tabIcon);
   }
 }
 
+function refreshIconBadgeText(tabId, icon) {
+  iconUtil.setBadgeText({
+    text: prefs.get('show-badge') && icon.count ? String(icon.count) : '',
+    tabId
+  });
+}
+
 function refreshIcon(tabId, icon) {
   const disableAll = prefs.get('disableAll');
   const iconset = prefs.get('iconset') === 1 ? 'light/' : '';
@@ -392,6 +413,13 @@ function refreshAllIcons() {
   for (const [tabId, icon] of tabIcons) {
     refreshIcon(tabId, icon);
   }
+  refreshIcon(null, {}); // default icon
+}
+
+function refreshAllIconsBadgeText() {
+  for (const [tabId, icon] of tabIcons) {
+    refreshIconBadgeText(tabId, icon);
+  }
 }
 
 function onRuntimeMessage(msg, sender) {
diff --git a/content/apply.js b/content/apply.js
index 16a43728..977026a8 100644
--- a/content/apply.js
+++ b/content/apply.js
@@ -74,6 +74,7 @@ const APPLY = (() => {
     };
 
     function injectPageScript() {
+      // FIXME: does it work with XML?
       const script = document.createElement('script');
       const scriptContent = EVENT_NAME => {
         document.currentScript.remove();
@@ -252,7 +253,7 @@ const APPLY = (() => {
       method: 'invokeAPI',
       name: 'updateIconBadge',
       args: [count]
-    });
+    }).catch(() => {});
   }
 
   function applyStyleState({id, enabled}) {
@@ -291,7 +292,9 @@ const APPLY = (() => {
 
   function applyStyles(styles, done) {
     if (!styles.length) {
-      done();
+      if (done) {
+        done();
+      }
       return;
     }
 
@@ -326,7 +329,9 @@ const APPLY = (() => {
 
     updateExposeIframes();
     updateCount();
-    done();
+    if (done) {
+      done();
+    }
   }
 
   function applySections(id, code) {