Fix memory leak / assertion error when switching to schemas with custom icons

WeaselUI/WeaselPanel.cpp

@@ -25,18 +25,34 @@
 
 #pragma comment(lib, "Shcore.lib")
 
-template <class t0, class t1, class t2>
-inline void LoadIconNecessary(t0& a, t1& b, t2& c, int d) {
-  if (a == b)
+namespace {
+
+void LoadCustomIconIfNecessary(std::wstring& loaded_path,
+                               const std::wstring& desired_path,
+                               CIcon& icon) {
+  if (loaded_path == desired_path)
+    return;
+  loaded_path = desired_path;
+  if (!icon.IsNull())
+    icon.DestroyIcon();
+  if (desired_path.empty())
     return;
-  a = b;
-  if (b.empty())
-    c.LoadIconW(d, STATUS_ICON_SIZE, STATUS_ICON_SIZE, LR_DEFAULTCOLOR);
-  else
-    c = (HICON)LoadImage(NULL, b.c_str(), IMAGE_ICON, STATUS_ICON_SIZE,
-                         STATUS_ICON_SIZE, LR_LOADFROMFILE);
+  HICON hIcon = static_cast<HICON>(
+      ::LoadImageW(NULL, desired_path.c_str(), IMAGE_ICON,  //
+                   STATUS_ICON_SIZE, STATUS_ICON_SIZE, LR_LOADFROMFILE));
+  if (hIcon)
+    icon.Attach(hIcon);
 }
 
+HICON GetCustomIconOrBuiltIn(const std::wstring& custom_icon_path,
+                             const CIcon& custom_icon,
+                             const CIcon& built_in_icon) {
+  return (!custom_icon_path.empty() && !custom_icon.IsNull()) ? custom_icon
+                                                              : built_in_icon;
+}
+
+}  // namespace
+
 static inline void ReconfigRoundInfo(IsToRoundStruct& rd,
                                      const int& i,
                                      const int& m_candidateCount) {
@@ -1081,28 +1097,37 @@ void WeaselPanel::DoPaint(CDCHandle dc) {
     // status icon (I guess Metro IME stole my idea :)
     if (m_layout->ShouldDisplayStatusIcon()) {
       // decide if custom schema zhung icon to show
-      LoadIconNecessary(m_current_zhung_icon, m_style.current_zhung_icon,
-                        m_iconEnabled, IDI_ZH);
-      LoadIconNecessary(m_current_ascii_icon, m_style.current_ascii_icon,
-                        m_iconAlpha, IDI_EN);
-      LoadIconNecessary(m_current_half_icon, m_style.current_half_icon,
-                        m_iconHalf, IDI_HALF_SHAPE);
-      LoadIconNecessary(m_current_full_icon, m_style.current_full_icon,
-                        m_iconFull, IDI_FULL_SHAPE);
+      LoadCustomIconIfNecessary(m_current_zhung_icon,
+                                m_style.current_zhung_icon,
+                                m_customIconEnabled);
+      LoadCustomIconIfNecessary(m_current_ascii_icon,
+                                m_style.current_ascii_icon, m_customIconAlpha);
+      LoadCustomIconIfNecessary(m_current_half_icon, m_style.current_half_icon,
+                                m_customIconHalf);
+      LoadCustomIconIfNecessary(m_current_full_icon, m_style.current_full_icon,
+                                m_customIconFull);
       CRect iconRect(m_layout->GetStatusIconRect());
       if (m_istorepos && !m_ctx.aux.str.empty())
         iconRect.OffsetRect(0, m_offsety_aux);
       else if (m_istorepos && !m_layout->IsInlinePreedit() &&
                !m_ctx.preedit.str.empty())
         iconRect.OffsetRect(0, m_offsety_preedit);
 
-      CIcon& icon(
-          m_status.disabled ? m_iconDisabled
-          : m_status.ascii_mode
-              ? m_iconAlpha
-              : (m_status.type == SCHEMA
-                     ? m_iconEnabled
-                     : (m_status.full_shape ? m_iconFull : m_iconHalf)));
+      HICON icon = m_status.disabled ? m_iconDisabled.m_hIcon
+                   : m_status.ascii_mode
+                       ? GetCustomIconOrBuiltIn(m_current_ascii_icon,
+                                                m_customIconAlpha, m_iconAlpha)
+                       : (m_status.type == SCHEMA
+                              ? GetCustomIconOrBuiltIn(m_current_zhung_icon,
+                                                       m_customIconEnabled,
+                                                       m_iconEnabled)
+                              : (m_status.full_shape
+                                     ? GetCustomIconOrBuiltIn(
+                                           m_current_full_icon,
+                                           m_customIconFull, m_iconFull)
+                                     : GetCustomIconOrBuiltIn(
+                                           m_current_half_icon,
+                                           m_customIconHalf, m_iconHalf)));
       memDC.DrawIconEx(iconRect.left, iconRect.top, icon, 0, 0);
       drawn = true;
     }

WeaselUI/WeaselPanel.h

@@ -124,6 +124,10 @@ class WeaselPanel
   CIcon m_iconAlpha;
   CIcon m_iconFull;
   CIcon m_iconHalf;
+  CIcon m_customIconEnabled;
+  CIcon m_customIconAlpha;
+  CIcon m_customIconFull;
+  CIcon m_customIconHalf;
   std::wstring m_current_zhung_icon;
   std::wstring m_current_ascii_icon;
   std::wstring m_current_half_icon;