Better history. Works even when non-current screens change

1 files changed, 78 insertions, 34 deletions

diff --git a/src/Rahm/Desktop/History.hs b/src/Rahm/Desktop/History.hs
index e39171f..ffcb10e 100644
--- a/src/Rahm/Desktop/History.hs
+++ b/src/Rahm/Desktop/History.hs
@@ -18,14 +18,14 @@ import Data.Default (Default (..))
 import Data.Foldable (find, toList)
 import Data.Map (Map)
 import qualified Data.Map as Map
-import Data.Maybe (fromMaybe, isJust)
+import Data.Maybe (catMaybes, fromMaybe, isJust)
 import Data.Sequence (Seq (..))
 import qualified Data.Sequence as Seq (length, (!?))
-import Rahm.Desktop.Common (Location (Location), focusLocation, getCurrentScreen, getCurrentWorkspace)
+import Rahm.Desktop.Common (Location (Location), focusLocation, getCurrentScreen, getCurrentWorkspace, locationWorkspace)
 import Rahm.Desktop.Hooks.WindowChange
 import Rahm.Desktop.Logger
 import Text.Printf (printf)
-import XMonad (ExtensionClass (extensionType, initialValue), ScreenId, StateExtension (..), X)
+import XMonad (ExtensionClass (extensionType, initialValue), ScreenId, StateExtension (..), Window, X)
 import XMonad.StackSet
 import qualified XMonad.Util.ExtensibleState as XS
   ( get,
@@ -58,7 +58,7 @@ zipperDbgPrint _ = "<empty>"
 pushZipper :: a -> BoundedSeqZipper a -> BoundedSeqZipper a
 pushZipper e (BoundedSeqZipper maxSize _ (tail :|> _))
   | maxSize <= Seq.length tail =
-      BoundedSeqZipper maxSize mempty (e :<| tail)
+    BoundedSeqZipper maxSize mempty (e :<| tail)
 pushZipper e (BoundedSeqZipper maxSize _ tail) =
   BoundedSeqZipper maxSize mempty (e :<| tail)
 
@@ -127,43 +127,87 @@ dbgLogHistory = do
   forM_ (Map.toList byScreen) $ \(screenId, hist) ->
     logs Trace "%s    ->    %s\n" (show screenId) (zipperDbgPrint hist)
 
+data ScreenDiff = ScreenDiff
+  { scrId :: ScreenId,
+    oldLocation :: Location,
+    newLocation :: Location
+  }
+
 historyHook :: WindowStack -> WindowStack -> X ()
 -- History hook where the 'from' location workspace does not match the 'to'
 -- location workspace.
 historyHook lastWindowSet currentWindowSet = do
-  let (sc1, ws1, win1) = getWindowsetData lastWindowSet
-      (sc2, ws2, win2) = getWindowsetData currentWindowSet
-      l1 = Location ws1 win1
-
-  case () of
-    -- We moved to a previously invisible workspace
-    () | not (ws2 `visibleIn` lastWindowSet) -> do
-      logs Trace "Jumped to hidden workspace"
-      XS.modify $ \(History byScreen) ->
-        History
-          ( Map.alter
-              (Just . pushZipper l1 . fromMaybe emptyZipper)
-              sc2
-              byScreen
-          )
-
-    -- We moved to a workspace that was on a different monitor, but was still
-    -- visible. In this case, we'll swap the history for the current screen with
-    -- the screen that the workspace was previously on. This will keep
-    -- per-screen history somewhat persistent
-    ()
-      | ws1 /= ws2 && sc1 == sc2,
-        (Just oldScreen) <- screenOf ws2 lastWindowSet -> do
-          logs Trace "Just Swapping Screens"
-          XS.modify $ \(History byScreen) ->
-            History (mapSwap oldScreen sc2 byScreen)
-
-    -- This is typically the case when changing focus to a different monitor,
-    -- but did not actually swap anything.
-    _ -> return ()
+  (History hist) <- XS.get
+  forM_ (getScreenDiffs lastWindowSet currentWindowSet) $
+    -- Read as "the screen <sid> went from <oloc> to <nloc>"
+    \(ScreenDiff sid oloc nloc) ->
+      let (ows, nws) = (locationWorkspace oloc, locationWorkspace nloc)
+
+       -- The goal here is to preserve history in as intuitive a way as possible
+       -- When the stackset changes, for each screen that changed in the last
+       -- windowchange, one of 2 situations are possibel:
+       --
+       --   1. The workspace on the screen was swapped with an already visible
+       --      screen
+       --
+       --   2. The workspace on the screen was swapped with a hidden workspace.
+       --
+       --   In the case of 1, we want to treat it as if the screen was
+       --   "reseated" to a different monitor, preserving the history for that
+       --   screen on its new screen.
+       --
+       --   In case of 2, we want to add the old workspace to the history of the
+       --   screen that changed.
+       in case () of
+            () | nws `visibleIn` lastWindowSet,
+                 (Just oscr) <- screenOf nws lastWindowSet ->
+              -- The last workspace was on a different screen. Swap the current
+              -- screen's history with the history from the last screen the
+              -- workspace was on.
+              XS.modify $ \(History byScreen) ->
+                History
+                  ( Map.alter
+                      (const $ Map.lookup oscr hist)
+                      sid
+                      byScreen
+                  )
+            -- The new workspace was not originally visible, add to history
+            () | not (nws `visibleIn` lastWindowSet) ->
+              XS.modify $ \(History byScreen) ->
+                History
+                  ( Map.alter
+                      (Just . pushZipper oloc . fromMaybe emptyZipper)
+                      sid
+                      byScreen
+                  )
+
+            -- This is typically not a possible case. It's only possible when a
+            -- screen is unplugged. If that's the case, do nothing.
+            _ -> return ()
 
   dbgLogHistory
   where
+    -- Returns a list of "screen diffs", which are a record of which screens
+    -- changed and how they changed.
+    getScreenDiffs os ns =
+      catMaybes $
+        Map.elems $
+          Map.intersectionWithKey
+            ( \screenId
+               (Screen ow@(Workspace ot _ _) _ _)
+               (Screen nw@(Workspace nt _ _) _ _) ->
+                  case () of
+                    () | ot == nt -> Nothing
+                    _ -> Just (ScreenDiff screenId (wsToLoc ow) (wsToLoc nw))
+            )
+            (screenMap os)
+            (screenMap ns)
+
+    wsToLoc (Workspace t _ (fmap focus -> win)) = Location t win
+
+    screenMap (StackSet ocur ovis _ _) =
+      Map.fromList $ map (\s -> (screen s, s)) (ocur : ovis)
+
     mapSwap k1 k2 map =
       Map.alter (const $ Map.lookup k2 map) k1 $
         Map.alter (const $ Map.lookup k1 map) k2 map