2 files changed, 23 insertions, 6 deletions

diff --git a/src/samples/thmviewer/display.cpp b/src/samples/thmviewer/display.cpp
index 52fa3cf8..c0e6c7e1 100644
--- a/src/samples/thmviewer/display.cpp
+++ b/src/samples/thmviewer/display.cpp
@@ -329,7 +329,7 @@ static BOOL DisplayOnCreate(
 {
     HRESULT hr = S_OK;
 
-    hr = ThemeLoadControls(pTheme, NULL, 0);
+    hr = ThemeLoadControls(pTheme);
     ExitOnFailure(hr, "Failed to load theme controls");
 
     // Pre-populate some control types with data.
diff --git a/src/samples/thmviewer/thmviewer.cpp b/src/samples/thmviewer/thmviewer.cpp
index f83182d3..cffa3851 100644
--- a/src/samples/thmviewer/thmviewer.cpp
+++ b/src/samples/thmviewer/thmviewer.cpp
@@ -14,10 +14,6 @@ enum THMVWR_CONTROL
     THMVWR_CONTROL_TREE = THEME_FIRST_ASSIGN_CONTROL_ID,
 };
 
-static THEME_ASSIGN_CONTROL_ID vrgInitControls[] = {
-    { THMVWR_CONTROL_TREE, L"Tree" },
-};
-
 // Internal functions
 
 static HRESULT ProcessCommandLine(
@@ -52,6 +48,10 @@ static void OnNewTheme(
     __in HWND hWnd,
     __in HANDLE_THEME* pHandle
     );
+static BOOL OnThemeLoadingControl(
+    __in const THEME_LOADINGCONTROL_ARGS* pArgs,
+    __in THEME_LOADINGCONTROL_RESULTS* pResults
+    );
 static void CALLBACK ThmviewerTraceError(
     __in_z LPCSTR szFile,
     __in int iLine,
@@ -353,7 +353,7 @@ static LRESULT CALLBACK MainWndProc(
 
     case WM_CREATE:
         {
-            HRESULT hr = ThemeLoadControls(vpTheme, vrgInitControls, countof(vrgInitControls));
+            HRESULT hr = ThemeLoadControls(vpTheme);
             if (FAILED(hr))
             {
                 return -1;
@@ -400,6 +400,9 @@ static LRESULT CALLBACK MainWndProc(
         }
         }
         break;
+
+    case WM_THMUTIL_LOADING_CONTROL:
+        return OnThemeLoadingControl(reinterpret_cast<THEME_LOADINGCONTROL_ARGS*>(wParam), reinterpret_cast<THEME_LOADINGCONTROL_RESULTS*>(lParam));
     }
 
     return ThemeDefWindowProc(vpTheme, hWnd, uMsg, wParam, lParam);
@@ -541,3 +544,17 @@ static void OnNewTheme(
         ThemeSendControlMessage(pTheme, THMVWR_CONTROL_TREE, TVM_SELECTITEM, TVGN_CARET, reinterpret_cast<LPARAM>(htiSelected));
     }
 }
+
+static BOOL OnThemeLoadingControl(
+    __in const THEME_LOADINGCONTROL_ARGS* pArgs,
+    __in THEME_LOADINGCONTROL_RESULTS* pResults
+    )
+{
+    if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, pArgs->pThemeControl->sczName, -1, L"Tree", -1))
+    {
+        pResults->wId = THMVWR_CONTROL_TREE;
+    }
+
+    pResults->hr = S_OK;
+    return TRUE;
+}