Introduce handling of new/old work areas (#1296)

* Introduce handling of new/old work areas (defined by virtual desktop id and monitor)

* Save data to JSON storage only once after we register all deleted virtual desktops

* Update saving mechanism

* Remove resolution from unique identifier

* Update comments related to new unique identifier

* Don't handle NULL GUID as it is default one

* Revert "Remove resolution from unique identifier"

This reverts commit 5a36651217.

* Revert "Update comments related to new unique identifier"

This reverts commit 8e89802549.

* removed MigrateDeviceInfoFromRegistry

Co-authored-by: Seraphima Zykova <zykovas91@gmail.com>

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
 {
@@ -152,7 +153,12 @@ private:
     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{};
@@ -168,7 +174,7 @@ private:
     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;
 
@@ -638,22 +644,21 @@ void FancyZones::AddZoneWindow(HMONITOR monitor, PCWSTR deviceId) noexcept
     if (SUCCEEDED_LOG(StringFromCLSID(m_currentVirtualDesktopId, &virtualDesktopId)))
     {
         std::wstring uniqueId = ZoneWindowUtils::GenerateUniqueId(monitor, deviceId, virtualDesktopId.get());
-        bool newVirtualDesktop = true;
+        JSONHelpers::FancyZonesDataInstance().SetActiveDeviceId(uniqueId);
 
-        auto it = m_virtualDesktopIds.find(m_currentVirtualDesktopId);
+        const bool newWorkArea = IsNewWorkArea(m_currentVirtualDesktopId, monitor);
-        if (it != end(m_virtualDesktopIds))
+        const bool flash = m_settings->GetSettings().zoneSetChange_flashZones && newWorkArea;
-        {
-            newVirtualDesktop = it->second;
-            JSONHelpers::FancyZonesDataInstance().SetActiveDeviceId(uniqueId);
-        }
 
-        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();
+        }
     }
 }
 
@@ -991,33 +996,75 @@ void FancyZones::HandleVirtualDesktopUpdates(HANDLE fancyZonesDestroyedEvent) no
         {
             return;
         }
-        const int guidSize = sizeof(GUID);
+        const size_t guidSize = sizeof(GUID);
-        std::unordered_map<GUID, bool> temp;
+        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);
+        RegisterVirtualDesktopUpdates(temp);
-        for (auto it = begin(m_virtualDesktopIds); it != end(m_virtualDesktopIds);)
-        {
-            auto iter = temp.find(it->first);
-            if (iter == temp.end())
-            {
-                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;
-            }
-        }
-        // register new virtual desktops, if any
-        m_virtualDesktopIds.insert(begin(temp), end(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())
+        {
+            // 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)))
+            {
+                modified |= JSONHelpers::FancyZonesDataInstance().RemoveDevicesByVirtualDesktopId(virtualDesktopId.get());
+            }
+            it = m_processedWorkAreas.erase(it);
+        }
+        else
+        {
+            currentVirtualDesktopIds.erase(it->first); // virtual desktop already in map, skip it
+            ++it;
+        }
+    }
+    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
 {
     // Colect information about changes in zone layout after editor exited.

src/modules/fancyzones/lib/JsonHelpers.cpp

@@ -22,6 +22,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>_<resolution>_<virtual-desktop-id>
+        return deviceId.substr(deviceId.rfind('_') + 1);
+    }
 }
 
 namespace JSONHelpers
@@ -174,11 +181,31 @@ namespace JSONHelpers
         {
             // Creates default entry in map when ZoneWindow is created
             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)
     {
         std::scoped_lock lock{ dataLock };
@@ -558,32 +585,6 @@ namespace JSONHelpers
         }
     }
 
-    void FancyZonesData::MigrateDeviceInfoFromRegistry(const std::wstring& deviceId)
-    {
-        std::scoped_lock lock{ dataLock };
-        wchar_t key[256];
-        StringCchPrintf(key, ARRAYSIZE(key), L"%s\\%s", RegistryHelpers::REG_SETTINGS, deviceId.c_str());
-
-        wchar_t activeZoneSetId[256];
-        activeZoneSetId[0] = '\0';
-        DWORD bufferSize = sizeof(activeZoneSetId);
-        DWORD showSpacing = 1;
-        DWORD spacing = 16;
-        DWORD zoneCount = 3;
-        DWORD size = sizeof(DWORD);
-
-        SHRegGetUSValueW(key, L"ActiveZoneSetId", nullptr, &activeZoneSetId, &bufferSize, FALSE, nullptr, 0);
-        SHRegGetUSValueW(key, L"ShowSpacing", nullptr, &showSpacing, &size, FALSE, nullptr, 0);
-        SHRegGetUSValueW(key, L"Spacing", nullptr, &spacing, &size, FALSE, nullptr, 0);
-        SHRegGetUSValueW(key, L"ZoneCount", nullptr, &zoneCount, &size, FALSE, nullptr, 0);
-
-        if (appliedZoneSetsMap.contains(std::wstring{ activeZoneSetId }))
-        {
-            deviceInfoMap[deviceId] = DeviceInfoData{ appliedZoneSetsMap.at(std::wstring{ activeZoneSetId }), static_cast<bool>(showSpacing), static_cast<int>(spacing), static_cast<int>(zoneCount) };
-            SaveFancyZonesData();
-        }
-    }
-
     void FancyZonesData::MigrateCustomZoneSetsFromRegistry()
     {
         std::scoped_lock lock{ dataLock };

src/modules/fancyzones/lib/JsonHelpers.h

@@ -222,6 +222,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;
@@ -247,8 +248,6 @@ namespace JSONHelpers
         void LoadFancyZonesData();
         void SaveFancyZonesData() const;
 
-        void MigrateDeviceInfoFromRegistry(const std::wstring& deviceId);
-
     private:
         void TmpMigrateAppliedZoneSetsFromRegistry();
         void MigrateCustomZoneSetsFromRegistry();