From 9ccd97798f122eea09ab2bd93bac35cd5a6d9106 Mon Sep 17 00:00:00 2001
From: vldmr11080 <57061786+vldmr11080@users.noreply.github.com>
Date: Tue, 17 Nov 2020 12:39:31 +0100
Subject: [PATCH] [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     |  3 +-
 src/modules/fancyzones/lib/FancyZonesData.cpp | 17 +++++-----
 .../fancyzones/lib/VirtualDesktopUtils.cpp    | 33 ++++++++-----------
 3 files changed, 25 insertions(+), 28 deletions(-)

diff --git a/src/modules/fancyzones/lib/FancyZones.cpp b/src/modules/fancyzones/lib/FancyZones.cpp
index 239227d582..f86302af60 100644
--- a/src/modules/fancyzones/lib/FancyZones.cpp
+++ b/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);
     }
 }
diff --git a/src/modules/fancyzones/lib/FancyZonesData.cpp b/src/modules/fancyzones/lib/FancyZonesData.cpp
index 435b44d345..281cb2fa91 100644
--- a/src/modules/fancyzones/lib/FancyZonesData.cpp
+++ b/src/modules/fancyzones/lib/FancyZonesData.cpp
@@ -262,16 +262,17 @@ 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);
-        auto foundId = active.find(desktopId);
-        if (foundId == std::end(active))
+        if (desktopId != NonLocalizable::DefaultGuid)
         {
-            RemoveDesktopAppZoneHistory(desktopId);
-            it = deviceInfoMap.erase(it);
-        }
-        else
-        {
-            ++it;
+            auto foundId = active.find(desktopId);
+            if (foundId == std::end(active))
+            {
+                RemoveDesktopAppZoneHistory(desktopId);
+                it = deviceInfoMap.erase(it);
+                continue;
+            }
         }
+        ++it;
     }
     SaveFancyZonesData();
 }
diff --git a/src/modules/fancyzones/lib/VirtualDesktopUtils.cpp b/src/modules/fancyzones/lib/VirtualDesktopUtils.cpp
index 5e7915d2d9..50491529b6 100644
--- a/src/modules/fancyzones/lib/VirtualDesktopUtils.cpp
+++ b/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(
-            StringCchPrintfW(
-                sessionKeyPath,
-                ARRAYSIZE(sessionKeyPath),
-                L"Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\SessionInfo\\%d\\VirtualDesktops",
-                sessionId));
+        if (FAILED(StringCchPrintfW(sessionKeyPath, ARRAYSIZE(sessionKeyPath), NonLocalizable::RegKeyVirtualDesktopsFromSession, sessionId)))
+        {
+            return false;
+        }
 
         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
-        // interface. Use foreground window (the window with which the user is currently working) to determine
-        // current virtual desktop.
-        else if (GetWindowDesktopId(GetForegroundWindow(), desktopId))
-        {
-            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.
+        // 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, previous function should return correct value, as desktop
+        // switch occurred in current session.
         else
         {
             std::vector<GUID> ids{};