1 files changed, 24 insertions, 15 deletions
diff --git a/src/Rahm/Desktop/Hooks/WindowChange.hs b/src/Rahm/Desktop/Hooks/WindowChange.hs
index a7d403e..4cb2371 100644
--- a/src/Rahm/Desktop/Hooks/WindowChange.hs
+++ b/src/Rahm/Desktop/Hooks/WindowChange.hs
@@ -1,28 +1,41 @@
 module Rahm.Desktop.Hooks.WindowChange where
 
 import Control.Monad (when)
+import Control.Monad.State (gets)
 import Data.Default (Default (..))
 import Rahm.Desktop.Common
   ( Location (Location),
+    getCurrentScreen,
     getCurrentWorkspace,
   )
 import qualified Rahm.Desktop.StackSet as W (peek)
 import XMonad
   ( ExtensionClass (..),
+    ScreenDetail,
+    ScreenId,
     StateExtension (PersistentExtension),
+    Window,
+    WorkspaceId,
     X,
     XConfig (logHook),
+    windowset,
     withWindowSet,
   )
+import XMonad.StackSet
 import qualified XMonad.Util.ExtensibleState as XS (get, put)
 
-newtype LastLocation = LastLocation (Maybe Location)
+type WindowStack = StackSet WorkspaceId () Window ScreenId ScreenDetail
+
+-- Type of hook. Takes the last WindowStack and the new WindowStack
+type StackChangeHook = WindowStack -> WindowStack -> X ()
+
+newtype LastState = LastState (Maybe WindowStack)
   deriving (Read, Show)
 
-instance Default LastLocation where
-  def = LastLocation Nothing
+instance Default LastState where
+  def = LastState def
 
-instance ExtensionClass LastLocation where
+instance ExtensionClass LastState where
   initialValue = def
   extensionType = PersistentExtension
 
@@ -32,20 +45,16 @@ instance ExtensionClass LastLocation where
 -- the new window.
 --
 -- If the first window is Nothing, this is the first time XMonad started.
-withLocationChangeHook :: (Maybe Location -> Location -> X ()) -> XConfig l -> XConfig l
-withLocationChangeHook fn config =
+withStackChangeHook :: StackChangeHook -> XConfig l -> XConfig l
+withStackChangeHook fn config =
   config
     { logHook = do
         logHook config
 
-        currentLocation <-
-          Location <$> getCurrentWorkspace <*> withWindowSet (return . W.peek)
-
-        LastLocation last <- XS.get
-
-        when (last /= Just currentLocation) $
-          fn last currentLocation
+        current <- gets (mapLayout (const ()) . windowset)
+        LastState last <- XS.get
+        XS.put (LastState $ Just current)
 
-        XS.put $ LastLocation $ Just currentLocation
-        return ()
+        when (Just current /= last) $
+          mapM_ (`fn` current) last
     }