From e2c6c45a2636b24b7efd1dacd2227689a0fe08c6 Mon Sep 17 00:00:00 2001
From: Anton Khirnov <anton@khirnov.net>
Date: Tue, 16 Nov 2021 16:09:18 +0100
Subject: rc/workspace: handle screen removal cleanly

---
 rc.lua         |  9 ++++++++-
 urgent_wgt.lua |  3 ++-
 utils.lua      |  7 +++++++
 workspace.lua  | 30 ++++++++++++++++++++++++++++++
 4 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/rc.lua b/rc.lua
index 9ff61ce..c72747a 100644
--- a/rc.lua
+++ b/rc.lua
@@ -79,7 +79,7 @@ mypromptbox = {}
 local wsp = workspace.Workspace:new(layouts)
 
 awful.screen.connect_for_each_screen(function(s)
-    print('connected screen')
+    print('Setting up screen:')
     for k, v in pairs(s.outputs) do
         print(k)
     end
@@ -126,6 +126,13 @@ awful.screen.connect_for_each_screen(function(s)
     mywibox[s] = awful.wibar({ position = "right", screen = s, width = panel_width })
     mywibox[s]:set_widget(layout)
 
+    -- clean up on screen removal
+    s:connect_signal("removed",
+        function (s)
+            mywibox[s]     = nil
+            mypromptbox[s] = nil
+        end)
+
     -- show desktop <n> on <n>th screen
     gears.timer.delayed_call(function(s) wsp:view(s, s.index, 1) end, s)
 end)
diff --git a/urgent_wgt.lua b/urgent_wgt.lua
index 1b2f60c..a50bd47 100644
--- a/urgent_wgt.lua
+++ b/urgent_wgt.lua
@@ -2,12 +2,13 @@ local awful     = require("awful")
 local beautiful = require("beautiful")
 local shape     = require("gears.shape")
 local wibox     = require("wibox")
+local utils     = require("utils")
 
 local M = {}
 
 local function client_get_desk_idx(c)
     if c.first_tag then
-        local desk_idx = tonumber('0x' .. c.first_tag.name:sub(1, 2))
+        local desk_idx = utils.tag_desk_idx(c.first_tag)
         if desk_idx ~= nil and desk_idx >= 0 then
             return desk_idx
         else
diff --git a/utils.lua b/utils.lua
index 4e3bb90..eb698dc 100644
--- a/utils.lua
+++ b/utils.lua
@@ -3,6 +3,13 @@ local M = {}
 local awful = require("awful")
 local naughty = require("naughty")
 
+function M.tag_desk_idx(t)
+    return tonumber('0x' .. t.name:sub(1, 2))
+end
+function M.tag_page_idx(t)
+    return tonumber('0x' .. t.name:sub(3, 4))
+end
+
 -- mapping from logical screen indices to indices corresponding to their
 -- physical layout
 function M.screen_physical(n)
diff --git a/workspace.lua b/workspace.lua
index 51eff90..82e5511 100644
--- a/workspace.lua
+++ b/workspace.lua
@@ -74,6 +74,26 @@ function Workspace:_apply_state()
     timer.delayed_call(awful.screen.focus, orig_focus)
 end
 
+-- save pages when a screen is removed
+function Workspace:_tag_request_screen(page)
+    local desk_idx = utils.tag_desk_idx(page)
+    local page_idx = utils.tag_page_idx(page)
+
+    if page.selected then
+        self.desktops[desk_idx].indices_free:push(page_idx)
+        page.selected = false
+    end
+
+    self.desktops[desk_idx].screen_map[page_idx] = nil
+
+    for s_other in screen do
+        if s_other ~= page.screen then
+            page.screen = s_other
+            break
+        end
+    end
+end
+
 function Workspace:view(screen, desktop_idx, page_idx)
     local desktop = self.desktops[desktop_idx]
 
@@ -203,6 +223,16 @@ function Workspace:new(layouts)
         o:_add_desktop(i, "Desktop " .. i, 10)
     end
 
+    -- handle screen removal
+    tag.connect_signal("request::screen",
+        function(t)
+            o:_tag_request_screen(t)
+        end)
+    screen.connect_signal("removed",
+        function (s)
+            o.screen_state[s] = nil
+        end)
+
     return o
 end
 
-- 
cgit v1.2.3