2 files changed, 30 insertions, 17 deletions
diff --git a/src/Rahm/Desktop/Keys.hs b/src/Rahm/Desktop/Keys.hs
index f4ccb72..d153a18 100644
--- a/src/Rahm/Desktop/Keys.hs
+++ b/src/Rahm/Desktop/Keys.hs
@@ -296,10 +296,9 @@ keymap = runKeys $ do
             withBorderColorM "#00ffff" l1 $ do
               lift $ addStringToPendingBuffer " "
               l2 <- mapMaybe (\(Location _ w) -> w) <$> readNextLocationSet
-              let rewriteMap = Map.fromList $ zip l1 l2 ++ zip l2 l1
               lift $ do
                 setAlternateWindows l1
-                windows $ mapWindows (\w -> fromMaybe w (Map.lookup w rewriteMap))
+                windows $ W.swapWindows $ zip l1 l2 ++ zip l2 l1
 
   bind xK_BackSpace $ do
     -- The only raw keybinding. Meant to get a terminal to unbrick XMonad if
diff --git a/src/Rahm/Desktop/SwapMaster.hs b/src/Rahm/Desktop/SwapMaster.hs
index c7e02bc..68072a9 100644
--- a/src/Rahm/Desktop/SwapMaster.hs
+++ b/src/Rahm/Desktop/SwapMaster.hs
@@ -5,11 +5,16 @@ import Control.Monad (void)
 import Control.Monad.State (gets)
 import Control.Monad.Trans (lift)
 import Control.Monad.Trans.Maybe (MaybeT (..))
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Rahm.Desktop.Common (runMaybeT_)
 import qualified Rahm.Desktop.StackSet as W
-  ( focusMaster,
+  ( currentTag,
+    focusMaster,
     masterWindow,
     peek,
     swapWindows,
+    tag,
   )
 import XMonad
   ( ExtensionClass (..),
@@ -18,33 +23,42 @@ import XMonad
     windows,
     windowset,
   )
-import qualified XMonad.Util.ExtensibleState as XS (get, put)
+import qualified XMonad.Util.ExtensibleState as XS (get, modify, put)
 
 newtype LastWindow = LastWindow
-  { lastWindow :: Maybe Window
+  { lastWindow :: Map String Window
   }
   deriving (Show, Read)
 
 instance ExtensionClass LastWindow where
-  initialValue = LastWindow Nothing
+  initialValue = LastWindow mempty
 
 hoist :: (Monad m) => Maybe a -> MaybeT m a
 hoist = MaybeT . return
 
 swapMaster :: X ()
-swapMaster = void $
-  runMaybeT $ do
+swapMaster =
+  runMaybeT_ $ do
     ss <- gets windowset
+    let ct = W.currentTag ss
 
-    focused <- hoist $ W.peek ss
-    master <- hoist $ W.masterWindow ss
-
-    if focused == master
-      then do
-        lw <- MaybeT $ lastWindow <$> XS.get
-        lift $ windows (W.swapWindows [(focused, lw)])
-      else lift $ windows (W.swapWindows [(focused, master)])
+    (focused, master) <-
+      hoist $ do
+        a <- W.peek ss
+        b <- W.masterWindow ss
+        return (a, b)
 
     lift $ do
-      XS.put (LastWindow $ Just master)
+      st <- lastWindow <$> XS.get
+      windows . W.swapWindows $
+        case focused == master of
+          True
+            | (Just lw) <- Map.lookup ct st ->
+              [(focused, lw)]
+          False -> [(master, focused)]
+          _ -> []
+
+      XS.modify $
+        \(LastWindow m) ->
+          LastWindow $ Map.insert (W.currentTag ss) master m
       windows W.focusMaster