[FancyZones] Fix window cycling on multiple monitors (#25145)

src/modules/fancyzones/FancyZonesLib/FancyZones.cpp

@@ -134,7 +134,7 @@ public:
 
     LRESULT WndProc(HWND, UINT, WPARAM, LPARAM) noexcept;
     void OnDisplayChange(DisplayChangeType changeType) noexcept;
-    void AddWorkArea(HMONITOR monitor, const FancyZonesDataTypes::WorkAreaId& id, bool updateWindowsPositions) noexcept;
+    bool AddWorkArea(HMONITOR monitor, const FancyZonesDataTypes::WorkAreaId& id) noexcept;
 
 protected:
     static LRESULT CALLBACK s_WndProc(HWND, UINT, WPARAM, LPARAM) noexcept;
@@ -722,7 +722,7 @@ void FancyZones::OnDisplayChange(DisplayChangeType changeType) noexcept
     UpdateWorkAreas(updateWindowsPositions);
 }
 
-void FancyZones::AddWorkArea(HMONITOR monitor, const FancyZonesDataTypes::WorkAreaId& id, bool updateWindowsPositions) noexcept
+bool FancyZones::AddWorkArea(HMONITOR monitor, const FancyZonesDataTypes::WorkAreaId& id) noexcept
 {
     wil::unique_cotaskmem_string virtualDesktopIdStr;
     if (!SUCCEEDED(StringFromCLSID(VirtualDesktop::instance().GetCurrentVirtualDesktopId(), &virtualDesktopIdStr)))
@@ -741,16 +741,14 @@ void FancyZones::AddWorkArea(HMONITOR monitor, const FancyZonesDataTypes::WorkAr
     }
         
     auto workArea = WorkArea::Create(m_hinstance, id, m_workAreaHandler.GetParent(monitor), rect);
-    if (workArea)
+    if (!workArea)
     {
-        workArea->InitSnappedWindows();
+        Logger::error(L"Failed to create work area {}", id.toString());
-        if (updateWindowsPositions)
+        return false;
-        {
-            workArea->UpdateWindowPositions();
-        }
-
-        m_workAreaHandler.AddWorkArea(monitor, std::move(workArea));
     }
+    
+    m_workAreaHandler.AddWorkArea(monitor, std::move(workArea));
+    return true;
 }
 
 LRESULT CALLBACK FancyZones::s_WndProc(HWND window, UINT message, WPARAM wparam, LPARAM lparam) noexcept
@@ -780,7 +778,7 @@ void FancyZones::UpdateWorkAreas(bool updateWindowPositions) noexcept
         workAreaId.virtualDesktopId = VirtualDesktop::instance().GetCurrentVirtualDesktopId();
         workAreaId.monitorId = { .deviceId = { .id = ZonedWindowProperties::MultiMonitorName, .instanceId = ZonedWindowProperties::MultiMonitorInstance } };
 
-        AddWorkArea(nullptr, workAreaId, updateWindowPositions);
+        AddWorkArea(nullptr, workAreaId);
     }
     else
     {
@@ -791,7 +789,76 @@ void FancyZones::UpdateWorkAreas(bool updateWindowPositions) noexcept
             workAreaId.virtualDesktopId = VirtualDesktop::instance().GetCurrentVirtualDesktopId();
             workAreaId.monitorId = monitor;
 
-            AddWorkArea(monitor.monitor, workAreaId, updateWindowPositions);
+            AddWorkArea(monitor.monitor, workAreaId);
+        }
+    }
+
+    // init previously snapped windows
+    std::unordered_map<HWND, ZoneIndexSet> windowsToSnap{};
+    for (const auto& window : VirtualDesktop::instance().GetWindowsFromCurrentDesktop())
+    {
+        auto indexes = FancyZonesWindowProperties::RetrieveZoneIndexProperty(window);
+        if (indexes.size() == 0)
+        {
+            continue;
+        }
+
+        windowsToSnap.insert({ window, indexes });
+    }
+
+    if (FancyZonesSettings::settings().spanZonesAcrossMonitors) // one work area across monitors
+    {
+        const auto workArea = m_workAreaHandler.GetWorkArea(nullptr);
+        if (workArea)
+        {
+            for (const auto& [window, zones] : windowsToSnap)
+            {
+                workArea->SnapWindow(window, zones, false);
+            }
+        }
+    }
+    else
+    {
+        // first, snap windows to the monitor where they're placed
+        for (auto iter = windowsToSnap.begin(); iter != windowsToSnap.end();)
+        {
+            const auto window = iter->first;
+            const auto zones = iter->second;
+            const auto monitor = MonitorFromWindow(window, MONITOR_DEFAULTTONULL);
+            const auto workAreaForMonitor = m_workAreaHandler.GetWorkArea(monitor);
+            if (workAreaForMonitor && workAreaForMonitor->GetWindowZoneIndexes(window) == zones)
+            {
+                workAreaForMonitor->SnapWindow(window, zones, false);
+                iter = windowsToSnap.erase(iter);
+            }
+            else
+            {
+                ++iter;
+            }
+        }
+
+        // snap rest of the windows to other work areas (in case they were moved after the monitor unplug)
+        for (const auto& [window, zones] : windowsToSnap)
+        {
+            for (const auto& [_, workArea] : m_workAreaHandler.GetAllWorkAreas())
+            {
+                const auto savedIndexes = workArea->GetWindowZoneIndexes(window);
+                if (savedIndexes == zones)
+                {
+                    workArea->SnapWindow(window, zones, false);
+                }
+            }
+        }
+    }
+
+    if (updateWindowPositions)
+    {
+        for (const auto& [_, workArea] : m_workAreaHandler.GetAllWorkAreas())
+        {
+            if (workArea)
+            {
+                workArea->UpdateWindowPositions();
+            }
         }
     }
 }

src/modules/fancyzones/FancyZonesLib/WorkArea.cpp

@@ -537,43 +537,6 @@ void WorkArea::InitLayout(const FancyZonesDataTypes::WorkAreaId& parentUniqueId)
     CalculateZoneSet();
 }
 
-void WorkArea::InitSnappedWindows()
-{
-    Logger::info(L"Init work area windows: {}", m_uniqueId.toString());
-
-    for (const auto& window : VirtualDesktop::instance().GetWindowsFromCurrentDesktop())
-    {
-        auto indexes = FancyZonesWindowProperties::RetrieveZoneIndexProperty(window);
-        if (indexes.size() == 0)
-        {
-            continue;
-        }
-
-        if (!m_uniqueId.monitorId.monitor) // one work area across monitors
-        {
-            MoveWindowIntoZoneByIndexSet(window, indexes, false);
-        }
-        else
-        {
-            const auto monitor = MonitorFromWindow(window, MONITOR_DEFAULTTONULL);
-            if (monitor && m_uniqueId.monitorId.monitor == monitor)
-            {
-                // prioritize snapping on the current monitor if the window was snapped to several work areas
-                MoveWindowIntoZoneByIndexSet(window, indexes, false);
-            }
-            else
-            {
-                // if the window is not snapped on the current monitor, then check the others
-                auto savedIndexes = GetWindowZoneIndexes(window);
-                if (savedIndexes == indexes)
-                {
-                    MoveWindowIntoZoneByIndexSet(window, indexes, false);
-                }
-            }
-        }
-    }
-}
-
 void WorkArea::CalculateZoneSet()
 {
     const auto appliedLayout = AppliedLayouts::instance().GetDeviceLayout(m_uniqueId);

src/modules/fancyzones/FancyZonesLib/WorkArea.h

@@ -51,7 +51,6 @@ public:
     bool MoveWindowIntoZoneByDirectionAndPosition(HWND window, DWORD vkCode, bool cycle);
     bool ExtendWindowByDirectionAndPosition(HWND window, DWORD vkCode);
 
-    void InitSnappedWindows();
     void SnapWindow(HWND window, const ZoneIndexSet& zones, bool extend = false);
     void UnsnapWindow(HWND window);