1 files changed, 80 insertions, 33 deletions

diff --git a/src/Rahm/Desktop/Layout/PinWindow.hs b/src/Rahm/Desktop/Layout/PinWindow.hs
index 959cc33..fe98afd 100644
--- a/src/Rahm/Desktop/Layout/PinWindow.hs
+++ b/src/Rahm/Desktop/Layout/PinWindow.hs
@@ -3,6 +3,8 @@
 module Rahm.Desktop.Layout.PinWindow where
 
 import Control.Arrow (Arrow (second))
+import Control.Exception (throw)
+import Control.Monad (unless, when)
 import Control.Monad.Trans (lift)
 import Control.Monad.Trans.Maybe (MaybeT (..))
 import Data.Foldable (find)
@@ -20,7 +22,7 @@ import XMonad
 import qualified XMonad.StackSet as W (filter)
 import qualified XMonad.Util.ExtensibleState as XS
 
-newtype PinWindowState = PinWindowState (Map ScreenId (Set Window))
+newtype PinWindowState = PinWindowState (Map ScreenId [(Window, W.RationalRect)])
   deriving (Show, Read)
 
 instance Default PinWindowState where
@@ -35,59 +37,85 @@ newtype PinWindowLayout (l :: * -> *) (a :: *) = PinWindowLayout (l a)
 
 instance (LayoutClass l a, sid ~ ScreenId, a ~ Window) => LayoutClass (PinWindowLayout l) a where
   runLayout (W.Workspace t (PinWindowLayout l) stack) rect = do
+    -- Clean up window id's thare should not be pinned anymore because the
+    -- windows were unmanaged.
+    (PinWindowState pinned') <- XS.get
+    ws <- gets windowset
+    XS.put $ PinWindowState (cleanupPinned pinned' ws)
     (PinWindowState pinned) <- XS.get
-    let allPinned = mconcat $ Map.elems pinned
 
-    (W.StackSet cur vis _ floating) <- gets windowset
-
-    (stack', winAndRect) <- case find ((== t) . (W.tag . W.workspace)) (cur : vis) of
-      Just (W.Screen ws sid _) -> do
-        let winsToMove = fromMaybe mempty (Map.lookup sid pinned)
-
-        modifyWindowSet $
-          appEndo $
-            mconcat $ map (Endo . W.shiftWinNoFocus (W.tag ws)) (Set.toList winsToMove)
-
-        updatedWorkspace <- withWindowSet $ return . W.findWorkspace t
-
-        return
-          ( maybe stack W.stack updatedWorkspace,
-            map (second (scaleRationalRect rect)) $
-              mapMaybe (\w -> (w,) <$> Map.lookup w floating) (Set.toList winsToMove)
-          )
-      Nothing -> return (stack, [])
+    let myScreen = find ((== t) . W.tag . W.workspace) (W.screens ws)
+        pinnedRects =
+          maybe
+            []
+            (map $ second (scaleRationalRect rect))
+            (((`Map.lookup` pinned) . W.screen) =<< myScreen)
+        allPinned = Set.fromList $ map fst $ concat $ Map.elems pinned
+        pinnedOnMyScreen = map fst pinnedRects
+
+        windowsToMove =
+          filter (\win -> W.findTag win ws /= Just t) pinnedOnMyScreen
+
+    modifyWindowSet $
+      appEndo . mconcat $
+        -- Move the windows that are supposed to now be on this workspace here
+        map (Endo . W.shiftWinNoFocus t) windowsToMove
+          ++
+          -- once again, sink the windows that are supposed to be here. Make
+          -- sure they don't float.
+          map (Endo . W.sink) (Set.toList allPinned)
 
     (rects, maybeNewLayout) <-
-      runLayout (W.Workspace t l (W.filter (not . (`Map.member` floating)) =<< stack')) rect
+      runLayout
+        (W.Workspace t l (W.filter (not . (`Set.member` allPinned)) =<< stack))
+        rect
+
+    return (pinnedRects ++ rects, PinWindowLayout <$> maybeNewLayout)
+    where
+      cleanupPinned mp ss =
+        let aw = Set.fromList (W.allWindows ss)
+         in Map.map (filter ((`Set.member` aw) . fst)) mp
 
-    return
-      ( winAndRect ++ rects,
-        PinWindowLayout <$> maybeNewLayout
-      )
+  handleMessage (PinWindowLayout l) a = do
+    maybeNewLayout <- handleMessage l a
+    return (PinWindowLayout <$> maybeNewLayout)
 
 pinWindow :: Window -> X ()
 pinWindow win = runMaybeT_ $ do
   lift $ logs Debug "Pinning window %d" win
 
-  ws@(W.StackSet cur vis _ _) <- gets windowset
+  ws@(W.StackSet cur vis _ flt) <- gets windowset
   t <- hoist (W.findTag win ws)
   scr <- hoist $ find ((== t) . (W.tag . W.workspace)) (cur : vis)
+  rect <- hoist $ Map.lookup win flt
 
-  lift $
+  lift $ do
     XS.modify $ \(PinWindowState mp) ->
       PinWindowState $
         Map.alter
-          (Just . maybe (Set.singleton win) (Set.insert win))
+          (Just . maybe [(win, rect)] ((win, rect) :))
           (W.screen scr)
           mp
+
+    -- Don't float the window anymore.
+    modifyWindowSet $ W.sink win
   where
     hoist = MaybeT . return
 
 unpinWindow :: Window -> X ()
-unpinWindow win = do
-  XS.modify $ \(PinWindowState mp) ->
-    PinWindowState $
-      Map.map (Set.delete win) mp
+unpinWindow win = runMaybeT_ $ do
+  (PinWindowState mp) <- lift XS.get
+  (win, rect) <- hoist $ find ((== win) . fst) (concat $ Map.elems mp)
+
+  lift $ do
+    XS.put $
+      PinWindowState $
+        Map.map (filter ((/= win) . fst)) mp
+
+    -- refloat the window.
+    modifyWindowSet $ W.float win rect
+  where
+    hoist = MaybeT . return
 
 toggleWindowPin :: Window -> X ()
 toggleWindowPin win = do
@@ -99,4 +127,23 @@ toggleWindowPin win = do
 isWindowPinned :: Window -> X Bool
 isWindowPinned win = do
   (PinWindowState mp) <- XS.get
-  return $ any (Set.member win) (Map.elems mp)
+  return $ any (any $ (== win) . fst) (Map.elems mp)
+
+pinnedWindows :: X [Window]
+pinnedWindows = do
+  (PinWindowState s) <- XS.get
+  return $ map fst $ concat $ Map.elems s
+
+pinnable :: l a -> PinWindowLayout l a
+pinnable = PinWindowLayout
+
+-- Unpins the window, executes the action, then repins the window. Useful for
+-- window shifts and whatnot.
+withWindowsUnpinned :: [Window] -> X () -> X ()
+withWindowsUnpinned wins fn =
+  ( do
+      mapM_ unpinWindow wins
+      fn
+      mapM_ pinWindow wins
+  )
+    `catchX` mapM_ pinWindow wins