1 files changed, 69 insertions, 1 deletions

diff --git a/src/Rahm/Desktop/StackSet.hs b/src/Rahm/Desktop/StackSet.hs
index 8db16c1..652dafe 100644
--- a/src/Rahm/Desktop/StackSet.hs
+++ b/src/Rahm/Desktop/StackSet.hs
@@ -7,6 +7,12 @@ module Rahm.Desktop.StackSet (
   shiftWin,
   screenRotateBackward,
   screenRotateForward,
+  mapWindows,
+  swapWindows,
+  getLocationWorkspace,
+  WindowLocation(..),
+  windowMemberOfWorkspace,
+  findWindow,
   module W) where
 
 import Prelude hiding (head)
@@ -15,7 +21,37 @@ import Data.List (find)
 import XMonad.StackSet as W hiding (greedyView, shiftWin)
 import qualified XMonad.StackSet
 import Data.Default
-import Data.Maybe (fromMaybe)
+import Data.Maybe (fromMaybe, catMaybes, listToMaybe)
+import qualified Data.Map as Map
+
+data WindowLocation i l a s sd =
+  OnScreen (Screen i l a s sd) |
+  OnHiddenWorkspace (Workspace i l a) |
+  Floating
+
+getLocationWorkspace :: WindowLocation i l a s sd -> Maybe (Workspace i l a)
+getLocationWorkspace (OnScreen (Screen w _ _)) = Just w
+getLocationWorkspace (OnHiddenWorkspace w) = Just w
+getLocationWorkspace _ = Nothing
+
+mapWindows :: (Ord a, Ord b) => (a -> b) -> StackSet i l a s sd -> StackSet i l b s sd
+mapWindows fn (StackSet cur vis hid float) =
+  StackSet
+    (mapWindowsScreen cur)
+    (map mapWindowsScreen vis)
+    (map mapWindowsWorkspace hid)
+    (Map.mapKeys fn float)
+  where
+    mapWindowsScreen (Screen work a b) = Screen (mapWindowsWorkspace work) a b
+    mapWindowsWorkspace (Workspace t l stack) =
+      Workspace t l (fmap (fmap fn) stack)
+
+swapWindows :: (Ord a) => a -> a -> StackSet i l a s d -> StackSet i l a s d
+swapWindows wa wb = mapWindows $ \w ->
+  case w of
+    _ | w == wa -> wb
+    _ | w == wb -> wa
+    _ -> w
 
 masterWindow :: StackSet i l a s sd -> Maybe a
 masterWindow = head . integrate' . stack . workspace . current
@@ -67,3 +103,35 @@ screenRotateForward (W.StackSet current visible others floating) = do
       in W.StackSet current' visible' others floating
 
   where rcycle l = last l : l
+
+{- Finds a Window and returns the screen its on and the workspace its on.
+ - Returns nothing if the window doesn't exist.
+ -
+ - If the window is not a screen Just (Nothing, workspace) is returned.
+ - If the window is a floating window Just (Nothing, Nothing) is returned. -}
+findWindow ::
+    (Eq a) => StackSet i l a s sd -> a -> Maybe (WindowLocation i l a s sd)
+findWindow (StackSet cur vis hid float) win =
+    listToMaybe . catMaybes $
+      map findWindowScreen (cur : vis) ++
+        map findWindowWorkspace hid ++
+          [findWindowFloat]
+
+  where
+    findWindowScreen s@(Screen ws _ _) =
+      if windowMemberOfWorkspace ws win
+        then Just (OnScreen s)
+        else Nothing
+
+    findWindowWorkspace w =
+      if windowMemberOfWorkspace w win
+        then Just (OnHiddenWorkspace w)
+        else Nothing
+
+    findWindowFloat =
+      if win `elem` Map.keys float
+        then Just Floating
+        else Nothing
+
+windowMemberOfWorkspace :: (Eq a) => Workspace i l a -> a -> Bool
+windowMemberOfWorkspace (Workspace _ _ s) w = w `elem` integrate' s