2 files changed, 84 insertions, 6 deletions

diff --git a/src/Rahm/Desktop/Keys.hs b/src/Rahm/Desktop/Keys.hs
index b57d310..fb49394 100644
--- a/src/Rahm/Desktop/Keys.hs
+++ b/src/Rahm/Desktop/Keys.hs
@@ -64,6 +64,8 @@ import Rahm.Desktop.Submap
 import Rahm.Desktop.Swallow
 import Rahm.Desktop.SwapMaster (swapMaster)
 import Rahm.Desktop.Workspaces
+import Rahm.Desktop.Theater
+
 import qualified Rahm.Desktop.StackSet as W
 import Rahm.Desktop.History
 
@@ -281,10 +283,18 @@ keymap = runKeys $ do
       runMaybeT_ $ (lift . gotoWorkspaceFn) =<< readNextWorkspace
 
     shiftMod $
-      doc "Swap a workspace with another workspace." $
-      runMaybeT_ $
-        lift . windows . uncurry W.swapWorkspaces =<<
-          (,) <$> readNextWorkspaceName <*> readNextWorkspaceName
+      doc "Restore the theater marked with the next typed character." $
+      runMaybeT_ $ do
+        mapNextString $ \_ str -> lift $
+          case str of
+            [ch] | isAlpha ch -> restoreTheater [ch]
+            _ -> return ()
+
+    -- shiftMod $
+    --   doc "Swap a workspace with another workspace." $
+    --   runMaybeT_ $
+    --     lift . windows . uncurry W.swapWorkspaces =<<
+    --       (,) <$> readNextWorkspaceName <*> readNextWorkspaceName
 
     controlMod $
       doc "Move the current focused window to another workspace and view that workspace" $
@@ -400,12 +410,12 @@ keymap = runKeys $ do
       doc "For mosaic layout, shrink the size-share of the current window" $
       sendMessage =<< shrinkPositionAlt
 
-  bind xK_m $
+  bind xK_m $ do
     justMod $
       doc "Mark the current window with the next typed character." $ do
       locs <- fromMaybe [] <$> runMaybeT readNextLocationSet
       let wins = mapMaybe locationWindow locs
-      withBorderWidth 4 wins $
+      withBorderWidth 2 wins $
         withBorderColor "#00ffff" wins $ do
           runMaybeT_ $ do
            mapNextString $ \_ str -> lift $
@@ -413,6 +423,14 @@ keymap = runKeys $ do
               [ch] | isAlpha ch -> markAllLocations str locs
               _ -> return ()
 
+    shiftMod $
+      doc "Mark the current theater with the next typed character." $
+      runMaybeT_ $ do
+        mapNextString $ \_ str -> lift $
+          case str of
+            [ch] | isAlpha ch -> saveCurrentTheater str
+            _ -> return ()
+
   bind xK_plus $ do
     justMod $
       doc "Increase the number of windows in the master region." $
diff --git a/src/Rahm/Desktop/Theater.hs b/src/Rahm/Desktop/Theater.hs
new file mode 100644
index 0000000..b0404b7
--- /dev/null
+++ b/src/Rahm/Desktop/Theater.hs
@@ -0,0 +1,60 @@
+module Rahm.Desktop.Theater where
+
+-- A "Theater" is a mapping from screen -> workspace. This is used to save the
+-- state of the current screen -> workspace and thus restore it.
+
+-- import XMonad.Operations
+import Data.Maybe (fromMaybe)
+import Control.Monad (forM_)
+import XMonad (X(..))
+import qualified XMonad.StackSet as W
+import qualified XMonad as X
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Default
+import qualified XMonad.Util.ExtensibleState as XS
+
+newtype Theater si wi = Theater (Map si wi)
+  deriving (Read, Show)
+
+newtype Theaters = Theaters {
+  theaters :: Map String (Theater X.ScreenId X.WorkspaceId)
+} deriving (Read, Show)
+
+instance Default Theaters where
+  def = Theaters mempty
+
+instance X.ExtensionClass Theaters where
+  initialValue = def
+  extensionType = X.PersistentExtension
+
+saveCurrentTheater :: String -> X ()
+saveCurrentTheater name =
+  X.withWindowSet $ \windowSet ->
+    XS.modify $ \(Theaters m) ->
+      Theaters $ flip (Map.insert name) m $
+        Theater $ Map.fromList $
+          map (\(W.Screen ws sid _) -> (sid, W.tag ws)) $ W.screens windowSet  
+
+restoreTheater :: String -> X ()
+restoreTheater name = do
+  (Theaters theaters) <- XS.get
+  forM_ (Map.lookup name theaters) $ \(Theater screenToWorkspace) ->
+    X.windows $ \ws@(W.StackSet cur vis hidden float) ->
+      let workspacesById = Map.fromList $ map (\ws -> (W.tag ws, ws)) (W.workspaces ws)
+
+          newScreenWorkspace scr =
+            fromMaybe scr $ do
+              wid <- Map.lookup (W.screen scr) screenToWorkspace
+              workspace <- Map.lookup wid workspacesById
+              return $ scr { W.workspace = workspace }
+
+          newScreens = map newScreenWorkspace (cur : vis)
+          newVisibleWorkspaces = map (W.tag . W.workspace) newScreens
+          newHiddenWorkspaces =
+            filter (\ws -> not (W.tag ws `elem` newVisibleWorkspaces)) $
+              W.workspaces ws
+
+          (newCur:newVisible) = newScreens
+        in
+          W.StackSet newCur newVisible newHiddenWorkspaces float