Introduce handling of new/old work areas

src/modules/fancyzones/lib/FancyZones.cpp

@@ -13,6 +13,7 @@
 #include <functional>
 #include <common/common.h>
 #include <lib\util.h>
+#include <unordered_set>
 
 enum class DisplayChangeType
 {
@@ -136,7 +137,12 @@ struct FancyZones : public winrt::implements<FancyZones, IFancyZones, IFancyZone
     void MoveSizeStartInternal(HWND window, HMONITOR monitor, POINT const& ptScreen, require_write_lock) noexcept;
     void MoveSizeEndInternal(HWND window, POINT const& ptScreen, require_write_lock) noexcept;
     void MoveSizeUpdateInternal(HMONITOR monitor, POINT const& ptScreen, require_write_lock) noexcept;
+
     void HandleVirtualDesktopUpdates(HANDLE fancyZonesDestroyedEvent) noexcept;
+    void RegisterVirtualDesktopUpdates(std::unordered_set<GUID>& currentVirtualDesktopIds) noexcept;
+    void RegisterNewWorkArea(GUID virtualDesktopId, HMONITOR monitor) noexcept;
+    bool IsNewWorkArea(GUID virtualDesktopId, HMONITOR monitor) noexcept;
+
     void OnEditorExitEvent() noexcept;
 
     const HINSTANCE m_hinstance{};
@@ -152,7 +158,7 @@ struct FancyZones : public winrt::implements<FancyZones, IFancyZones, IFancyZone
     winrt::com_ptr<IZoneWindow> m_zoneWindowMoveSize; // "Active" ZoneWindow, where the move/size is happening. Will update as drag moves between monitors.
     winrt::com_ptr<IFancyZonesSettings> m_settings{};
     GUID m_currentVirtualDesktopId{}; // UUID of the current virtual desktop. Is GUID_NULL until first VD switch per session.
-    std::unordered_map<GUID, bool> m_virtualDesktopIds;
+    std::unordered_map<GUID, std::vector<HMONITOR>> m_processedWorkAreas; // Work area is defined by monitor and virtual desktop id.
     wil::unique_handle m_terminateEditorEvent; // Handle of FancyZonesEditor.exe we launch and wait on
     wil::unique_handle m_terminateVirtualDesktopTrackerEvent;
 
@@ -607,23 +613,22 @@ void FancyZones::AddZoneWindow(HMONITOR monitor, PCWSTR deviceId) noexcept
     wil::unique_cotaskmem_string virtualDesktopId;
     if (SUCCEEDED_LOG(StringFromCLSID(m_currentVirtualDesktopId, &virtualDesktopId)))
     {
-        std::wstring uniqueId = ZoneWindowUtils::GenerateUniqueId(monitor, deviceId, virtualDesktopId.get());
-        bool newVirtualDesktop = true;
+        std::wstring uniqueId = ZoneWindowUtils::GenerateUniqueId(deviceId, virtualDesktopId.get());
+        JSONHelpers::FancyZonesDataInstance().SetActiveDeviceId(uniqueId);
 
-        auto it = m_virtualDesktopIds.find(m_currentVirtualDesktopId);
-        if (it != end(m_virtualDesktopIds))
-        {
-            newVirtualDesktop = it->second;
-            JSONHelpers::FancyZonesDataInstance().SetActiveDeviceId(uniqueId);
-        }
+        const bool newWorkArea = IsNewWorkArea(m_currentVirtualDesktopId, monitor);
+        const bool flash = m_settings->GetSettings().zoneSetChange_flashZones && newWorkArea;
 
-        const bool flash = m_settings->GetSettings().zoneSetChange_flashZones && newVirtualDesktop;
         auto zoneWindow = MakeZoneWindow(this, m_hinstance, monitor, uniqueId, flash);
         if (zoneWindow)
         {
             m_zoneWindowMap[monitor] = std::move(zoneWindow);
         }
-        m_virtualDesktopIds[m_currentVirtualDesktopId] = false;
+        if (newWorkArea)
+        {
+            RegisterNewWorkArea(m_currentVirtualDesktopId, monitor);
+            JSONHelpers::FancyZonesDataInstance().SaveFancyZonesData();
+        }
     }
 }
 
@@ -961,31 +966,73 @@ void FancyZones::HandleVirtualDesktopUpdates(HANDLE fancyZonesDestroyedEvent) no
         {
             return;
         }
-        const int guidSize = sizeof(GUID);
-        std::unordered_map<GUID, bool> temp;
+        const size_t guidSize = sizeof(GUID);
+        std::unordered_set<GUID> temp;
         temp.reserve(bufferCapacity / guidSize);
         for (size_t i = 0; i < bufferCapacity; i += guidSize)
         {
             GUID* guid = reinterpret_cast<GUID*>(buffer.get() + i);
-            temp[*guid] = true;
+            temp.insert(*guid);
         }
-        std::unique_lock writeLock(m_lock);
-        for (auto it = begin(m_virtualDesktopIds); it != end(m_virtualDesktopIds);)
+        RegisterVirtualDesktopUpdates(temp);
+    }
+}
+
+void FancyZones::RegisterVirtualDesktopUpdates(std::unordered_set<GUID>& currentVirtualDesktopIds) noexcept
+{
+    std::unique_lock writeLock(m_lock);
+    bool modified{ false };
+    for (auto it = begin(m_processedWorkAreas); it != end(m_processedWorkAreas);)
+    {
+        auto iter = currentVirtualDesktopIds.find(it->first);
+        if (iter == currentVirtualDesktopIds.end())
         {
-            auto iter = temp.find(it->first);
-            if (iter == temp.end())
+            // if we couldn't find the GUID in currentVirtualDesktopIds, we must remove it from both m_processedWorkAreas and deviceInfoMap
+            wil::unique_cotaskmem_string virtualDesktopId;
+            if (SUCCEEDED_LOG(StringFromCLSID(it->first, &virtualDesktopId)))
             {
-                it = m_virtualDesktopIds.erase(it); // virtual desktop closed, remove it from map
-            }
-            else
-            {
-                temp.erase(it->first); // virtual desktop already in map, skip it
-                ++it;
+                modified |= JSONHelpers::FancyZonesDataInstance().RemoveDevicesByVirtualDesktopId(virtualDesktopId.get());
             }
+            it = m_processedWorkAreas.erase(it);
+        }
+        else
+        {
+            currentVirtualDesktopIds.erase(it->first); // virtual desktop already in map, skip it
+            ++it;
         }
-        // register new virtual desktops, if any
-        m_virtualDesktopIds.insert(begin(temp), end(temp));
     }
+    if (modified)
+    {
+        JSONHelpers::FancyZonesDataInstance().SaveFancyZonesData();
+    }
+    // register new virtual desktops, if any
+    for (const auto& id : currentVirtualDesktopIds)
+    {
+        m_processedWorkAreas[id] = std::vector<HMONITOR>();
+    }
+}
+
+void FancyZones::RegisterNewWorkArea(GUID virtualDesktopId, HMONITOR monitor) noexcept
+{
+    if (!m_processedWorkAreas.contains(virtualDesktopId))
+    {
+        m_processedWorkAreas[virtualDesktopId] = { monitor };
+    }
+    else
+    {
+        m_processedWorkAreas[virtualDesktopId].push_back(monitor);
+    }
+}
+
+bool FancyZones::IsNewWorkArea(GUID virtualDesktopId, HMONITOR monitor) noexcept
+{
+    auto it = m_processedWorkAreas.find(virtualDesktopId);
+    if (it != m_processedWorkAreas.end())
+    {
+        // virtual desktop exists, check if it's processed on given monitor
+        return std::find(it->second.begin(), it->second.end(), monitor) == it->second.end();
+    }
+    return true;
 }
 
 void FancyZones::OnEditorExitEvent() noexcept

src/modules/fancyzones/lib/JsonHelpers.cpp

@@ -21,6 +21,13 @@ namespace
     constexpr int c_blankCustomModelId = 0xFFFA;
 
     const wchar_t* FANCY_ZONES_DATA_FILE = L"zones-settings.json";
+    const wchar_t* DEFAULT_GUID = L"{00000000-0000-0000-0000-000000000000}";
+
+    std::wstring ExtractVirtualDesktopId(const std::wstring& deviceId)
+    {
+        // Format: <device-id>_<virtual-desktop-id>
+        return deviceId.substr(deviceId.rfind('_') + 1);
+    }
 }
 
 namespace JSONHelpers
@@ -160,12 +167,34 @@ namespace JSONHelpers
         if (!deviceInfoMap.contains(deviceId))
         {
             // Creates default entry in map when ZoneWindow is created
-            deviceInfoMap[deviceId] = DeviceInfoData{ ZoneSetData{ L"null", ZoneSetLayoutType::Blank }  };
+            deviceInfoMap[deviceId] = DeviceInfoData{ ZoneSetData{ L"null", ZoneSetLayoutType::Blank } };
 
             MigrateDeviceInfoFromRegistry(deviceId);
         }
     }
 
+    bool FancyZonesData::RemoveDevicesByVirtualDesktopId(const std::wstring& virtualDesktopId)
+    {
+        if (virtualDesktopId == DEFAULT_GUID)
+        {
+            return false;
+        }
+        bool modified{ false };
+        for (auto it = deviceInfoMap.begin(); it != deviceInfoMap.end();)
+        {
+            if (ExtractVirtualDesktopId(it->first) == virtualDesktopId)
+            {
+                it = deviceInfoMap.erase(it);
+                modified = true;
+            }
+            else
+            {
+                ++it;
+            }
+        }
+        return modified;
+    }
+
     void FancyZonesData::CloneDeviceInfo(const std::wstring& source, const std::wstring& destination)
     {
         // Clone information from source device if destination device is uninitialized (Blank).
@@ -386,7 +415,8 @@ namespace JSONHelpers
 
         for (const auto& [deviceID, deviceData] : deviceInfoMap)
         {
-            if (deviceData.activeZoneSet.type != ZoneSetLayoutType::Blank) {
+            if (deviceData.activeZoneSet.type != ZoneSetLayoutType::Blank)
+            {
                 DeviceInfosJSON.Append(DeviceInfoJSON::DeviceInfoJSON::ToJson(DeviceInfoJSON{ deviceID, deviceData }));
             }
         }

src/modules/fancyzones/lib/JsonHelpers.h

@@ -196,6 +196,7 @@ namespace JSONHelpers
         }
 
         void AddDevice(const std::wstring& deviceId);
+        bool RemoveDevicesByVirtualDesktopId(const std::wstring& virtualDesktopId);
         void CloneDeviceInfo(const std::wstring& source, const std::wstring& destination);
 
         int GetAppLastZoneIndex(HWND window, const std::wstring_view& deviceId, const std::wstring_view& zoneSetId) const;

src/modules/fancyzones/lib/ZoneWindow.cpp

@@ -63,19 +63,16 @@ namespace ZoneWindowUtils
         return customZoneSetsTmpFileName;
     }
 
-    std::wstring GenerateUniqueId(HMONITOR monitor, PCWSTR deviceId, PCWSTR virtualDesktopId)
+    std::wstring GenerateUniqueId(PCWSTR deviceId, PCWSTR virtualDesktopId)
     {
-        wchar_t uniqueId[256]{}; // Parsed deviceId + resolution + virtualDesktopId
+        // Format: <device-id>_<virtual-desktop-id>
+        wchar_t uniqueId[256]{};
 
-        MONITORINFOEXW mi;
-        mi.cbSize = sizeof(mi);
-        if (virtualDesktopId && GetMonitorInfo(monitor, &mi))
+        if (virtualDesktopId)
         {
             wchar_t parsedId[256]{};
             ParseDeviceId(deviceId, parsedId, 256);
-
-            Rect const monitorRect(mi.rcMonitor);
-            StringCchPrintf(uniqueId, ARRAYSIZE(uniqueId), L"%s_%d_%d_%s", parsedId, monitorRect.width(), monitorRect.height(), virtualDesktopId);
+            StringCchPrintf(uniqueId, ARRAYSIZE(uniqueId), L"%s_%s", parsedId, virtualDesktopId);
         }
         return std::wstring{ uniqueId };
     }

src/modules/fancyzones/lib/ZoneWindow.h

@@ -7,7 +7,7 @@ namespace ZoneWindowUtils
     const std::wstring& GetActiveZoneSetTmpPath();
     const std::wstring& GetAppliedZoneSetTmpPath();
     const std::wstring& GetCustomZoneSetsTmpPath();
-    std::wstring GenerateUniqueId(HMONITOR monitor, PCWSTR deviceId, PCWSTR virtualDesktopId);
+    std::wstring GenerateUniqueId(PCWSTR deviceId, PCWSTR virtualDesktopId);
 }
 
 interface __declspec(uuid("{7F017528-8110-4FB3-BE41-F472969C2560}")) IZoneWindow : public IUnknown

src/modules/fancyzones/tests/UnitTests/ZoneWindow.Spec.cpp

@@ -173,11 +173,11 @@ namespace FancyZonesUnitTests
         TEST_METHOD(CreateZoneWindowNoDeviceId)
         {
             // Generate unique id without device id
-            std::wstring uniqueId = ZoneWindowUtils::GenerateUniqueId(m_monitor, nullptr, m_virtualDesktopId.c_str());
+            std::wstring uniqueId = ZoneWindowUtils::GenerateUniqueId(nullptr, m_virtualDesktopId.c_str());
             m_zoneWindow = MakeZoneWindow(m_hostPtr, m_hInst, m_monitor, uniqueId, false);
 
             const std::wstring expectedWorkArea = std::to_wstring(m_monitorInfo.rcMonitor.right) + L"_" + std::to_wstring(m_monitorInfo.rcMonitor.bottom);
-            const std::wstring expectedUniqueId = L"FallbackDevice_" + std::to_wstring(m_monitorInfo.rcMonitor.right) + L"_" + std::to_wstring(m_monitorInfo.rcMonitor.bottom) + L"_" + m_virtualDesktopId;
+            const std::wstring expectedUniqueId = L"FallbackDevice_" + m_virtualDesktopId;
 
             Assert::IsNotNull(m_zoneWindow.get());
             Assert::IsFalse(m_zoneWindow->IsDragEnabled());
@@ -189,10 +189,9 @@ namespace FancyZonesUnitTests
         TEST_METHOD(CreateZoneWindowNoDesktopId)
         {
             // Generate unique id without virtual desktop id
-            std::wstring uniqueId = ZoneWindowUtils::GenerateUniqueId(m_monitor, m_deviceId.c_str(), nullptr);
+            std::wstring uniqueId = ZoneWindowUtils::GenerateUniqueId(m_deviceId.c_str(), nullptr);
             m_zoneWindow = MakeZoneWindow(m_hostPtr, m_hInst, m_monitor, uniqueId, false);
 
-            const std::wstring expectedWorkArea = std::to_wstring(m_monitorInfo.rcMonitor.right) + L"_" + std::to_wstring(m_monitorInfo.rcMonitor.bottom);
             Assert::IsNotNull(m_zoneWindow.get());
             Assert::IsFalse(m_zoneWindow->IsDragEnabled());
             Assert::IsTrue(m_zoneWindow->UniqueId().empty());