From b454f1e0be143353b76b0c69fc3736b9cdaa3f0f Mon Sep 17 00:00:00 2001
From: Tomas Janousek <tomi@nomi.cz>
Date: Wed, 16 Oct 2024 13:59:12 +0100
Subject: [PATCH] X.H.Rescreen: Move error handling to rescreenHook

This handles errors in hooks set using `rescreenHook` as well, not just
those set using the individual adders/setters.

Fixes: 2e3254a9080c ("X.H.Rescreen: Catch exceptions in user-provided hooks in add*Hook")
---
 XMonad/Hooks/Rescreen.hs | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/XMonad/Hooks/Rescreen.hs b/XMonad/Hooks/Rescreen.hs
index 4a33657b..594e101b 100644
--- a/XMonad/Hooks/Rescreen.hs
+++ b/XMonad/Hooks/Rescreen.hs
@@ -101,21 +101,28 @@ instance Monoid RescreenConfig where
 -- 'rescreenWorkspacesHook', which has a replace rather than sequence
 -- semantics), also just once.
 rescreenHook :: RescreenConfig -> XConfig l -> XConfig l
-rescreenHook = XC.once $ \c -> c
-    { startupHook = startupHook c <> rescreenStartupHook
-    , handleEventHook = handleEventHook c <> rescreenEventHook }
+rescreenHook = XC.once hook . catchUserCode
+  where
+    hook c = c
+        { startupHook = startupHook c <> rescreenStartupHook
+        , handleEventHook = handleEventHook c <> rescreenEventHook }
+    catchUserCode rc@RescreenConfig{..} = rc
+        { afterRescreenHook = userCodeDef () afterRescreenHook
+        , randrChangeHook = userCodeDef () randrChangeHook
+        , rescreenWorkspacesHook = flip catchX rescreen <$> rescreenWorkspacesHook
+        }
 
 -- | Shortcut for 'rescreenHook'.
 addAfterRescreenHook :: X () -> XConfig l -> XConfig l
-addAfterRescreenHook h = rescreenHook def{ afterRescreenHook = userCodeDef () h }
+addAfterRescreenHook h = rescreenHook def{ afterRescreenHook = h }
 
 -- | Shortcut for 'rescreenHook'.
 addRandrChangeHook :: X () -> XConfig l -> XConfig l
-addRandrChangeHook h = rescreenHook def{ randrChangeHook = userCodeDef () h }
+addRandrChangeHook h = rescreenHook def{ randrChangeHook = h }
 
 -- | Shortcut for 'rescreenHook'.
 setRescreenWorkspacesHook :: X () -> XConfig l -> XConfig l
-setRescreenWorkspacesHook h = rescreenHook def{ rescreenWorkspacesHook = pure (catchX h rescreen) }
+setRescreenWorkspacesHook h = rescreenHook def{ rescreenWorkspacesHook = pure h }
 
 -- | Startup hook to listen for @RRScreenChangeNotify@ events.
 rescreenStartupHook :: X ()