[FancyZones] Update primary desktop data on virtual desktop switch (#7994)

* Update primary desktop data after virtual desktop switch

* Don't remove zeroed-GUID inside RemoveDeletedDesktops method

* Minor refactoring in VirtualDesktopUtils

* Use std::vector::empty check instead of comparing size with 0

src/modules/fancyzones/lib/FancyZones.cpp

@@ -1276,8 +1276,9 @@ void FancyZones::RegisterVirtualDesktopUpdates(std::vector<GUID>& ids) noexcept
 
     m_workAreaHandler.RegisterUpdates(ids);
     std::vector<std::wstring> active{};
-    if (VirtualDesktopUtils::GetVirtualDesktopIds(active))
+    if (VirtualDesktopUtils::GetVirtualDesktopIds(active) && !active.empty())
     {
+        FancyZonesDataInstance().UpdatePrimaryDesktopData(active[0]);
         FancyZonesDataInstance().RemoveDeletedDesktops(active);
     }
 }

src/modules/fancyzones/lib/FancyZonesData.cpp

@@ -262,17 +262,18 @@ void FancyZonesData::RemoveDeletedDesktops(const std::vector<std::wstring>& acti
     for (auto it = std::begin(deviceInfoMap); it != std::end(deviceInfoMap);)
     {
         std::wstring desktopId = ExtractVirtualDesktopId(it->first);
+        if (desktopId != NonLocalizable::DefaultGuid)
+        {
             auto foundId = active.find(desktopId);
             if (foundId == std::end(active))
             {
                 RemoveDesktopAppZoneHistory(desktopId);
                 it = deviceInfoMap.erase(it);
+                continue;
+            }
         }
-        else
-        {
         ++it;
     }
-    }
     SaveFancyZonesData();
 }
 

src/modules/fancyzones/lib/VirtualDesktopUtils.cpp

@@ -8,6 +8,7 @@ namespace NonLocalizable
     const wchar_t RegCurrentVirtualDesktop[] = L"CurrentVirtualDesktop";
     const wchar_t RegVirtualDesktopIds[] = L"VirtualDesktopIDs";
     const wchar_t RegKeyVirtualDesktops[] = L"Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\VirtualDesktops";
+    const wchar_t RegKeyVirtualDesktopsFromSession[] = L"Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\SessionInfo\\%d\\VirtualDesktops";
 }
 
 namespace VirtualDesktopUtils
@@ -53,20 +54,21 @@ namespace VirtualDesktopUtils
     bool GetDesktopIdFromCurrentSession(GUID* desktopId)
     {
         DWORD sessionId;
-        ProcessIdToSessionId(GetCurrentProcessId(), &sessionId);
+        if (!ProcessIdToSessionId(GetCurrentProcessId(), &sessionId))
+        {
+            return false;
+        }
 
         wchar_t sessionKeyPath[256]{};
-        RETURN_IF_FAILED(
+        if (FAILED(StringCchPrintfW(sessionKeyPath, ARRAYSIZE(sessionKeyPath), NonLocalizable::RegKeyVirtualDesktopsFromSession, sessionId)))
-            StringCchPrintfW(
+        {
-                sessionKeyPath,
+            return false;
-                ARRAYSIZE(sessionKeyPath),
+        }
-                L"Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\SessionInfo\\%d\\VirtualDesktops",
-                sessionId));
 
         wil::unique_hkey key{};
-        GUID value{};
         if (RegOpenKeyExW(HKEY_CURRENT_USER, sessionKeyPath, 0, KEY_ALL_ACCESS, &key) == ERROR_SUCCESS)
         {
+            GUID value{};
             DWORD size = sizeof(GUID);
             if (RegQueryValueExW(key.get(), NonLocalizable::RegCurrentVirtualDesktop, 0, nullptr, reinterpret_cast<BYTE*>(&value), &size) == ERROR_SUCCESS)
             {
@@ -86,17 +88,10 @@ namespace VirtualDesktopUtils
         {
             return true;
         }
-        // First fallback scenario is to try obtaining virtual desktop id through IVirtualDesktopManager
+        // Fallback scenario is to get array of virtual desktops stored in registry, but not kept per session.
-        // interface. Use foreground window (the window with which the user is currently working) to determine
+        // Note that we are taking first element from virtual desktop array, which is primary desktop.
-        // current virtual desktop.
+        // If user has more than one virtual desktop, previous function should return correct value, as desktop
-        else if (GetWindowDesktopId(GetForegroundWindow(), desktopId))
+        // switch occurred in current session.
-        {
-            return true;
-        }
-        // Second fallback scenario is to get array of virtual desktops stored in registry, but not kept per
-        // session. Note that we are taking first element from virtual desktop array, which is primary desktop.
-        // If user has more than one virtual desktop, one of previous functions should return correct value,
-        // as desktop switch occurred in current session.
         else
         {
             std::vector<GUID> ids{};