1 files changed, 106 insertions, 5 deletions

diff --git a/src/Rahm/Desktop/Keys/Wml.hs b/src/Rahm/Desktop/Keys/Wml.hs
index 6c46361..0d0691f 100644
--- a/src/Rahm/Desktop/Keys/Wml.hs
+++ b/src/Rahm/Desktop/Keys/Wml.hs
@@ -17,9 +17,15 @@ module Rahm.Desktop.Keys.Wml
     readWindowsetMacro,
     readNextWorkspace,
     readNextLocationSet,
+    readNextLocationSet',
     moveLocationToWorkspace,
     moveWindowToWorkspaceFn,
+    getAndResetWindowSelection,
     gotoWorkspaceFn,
+    toggleWindowInSelection,
+    addWindowToSelection,
+    clearWindowSelection,
+    removeWindowFromSelection,
     readMacroString,
     justWorkspace,
     justWorkspaceWithPreferredWindow,
@@ -37,10 +43,11 @@ module Rahm.Desktop.Keys.Wml
     locationSetForKeys,
     readNextWorkspaceName,
     workspaceName,
+    wmlLogHook,
   )
 where
 
-import Control.Monad (forM_, join, when)
+import Control.Monad (forM_, join, void, when)
 import Control.Monad.Trans.Class (MonadTrans (lift))
 import Control.Monad.Trans.Maybe (MaybeT (..), mapMaybeT)
 import Control.Monad.Trans.State as S
@@ -50,17 +57,19 @@ import Control.Monad.Trans.State as S
     put,
   )
 import Data.Char (isAlpha, isAlphaNum, isDigit, ord)
-import Data.List (intercalate, sortOn)
+import Data.List (find, intercalate, sortOn)
 import Data.List.Safe (head, last)
 import Data.Map (Map)
 import qualified Data.Map as Map
-  ( empty,
+  ( delete,
+    elems,
+    empty,
     insert,
     keys,
     lookup,
     member,
   )
-import Data.Maybe (catMaybes, fromMaybe)
+import Data.Maybe (catMaybes, fromJust, fromMaybe)
 import Data.Ord (Down (..))
 import Data.Typeable (cast)
 import Rahm.Desktop.Common
@@ -70,6 +79,7 @@ import Rahm.Desktop.Common
     getCurrentWorkspace,
     gotoWorkspace,
     moveLocationToWorkspace,
+    setBorderColor,
     windowsInWorkspace,
   )
 import Rahm.Desktop.History
@@ -122,6 +132,7 @@ import XMonad
     WorkspaceId,
     X,
     asks,
+    clearArea,
     directories,
     io,
     killWindow,
@@ -172,6 +183,10 @@ saveMacros = do
   macros <- getMacros
   io $ writeFile dataFile $ show macros
 
+selColor = "#b8b880"
+
+selFocusColor = "#ffff00"
+
 insertWorkspaceMacroString :: (KeyMask, KeySym) -> KeyString -> X ()
 insertWorkspaceMacroString k ks = do
   macros <- getMacros
@@ -195,6 +210,68 @@ insertWindowSetMacroString k ks = do
 instance ExtensionClass MaybeMacros where
   initialValue = NoMacros
 
+newtype WindowSelect = WindowSelect (Map Window (X ()))
+
+instance ExtensionClass WindowSelect where
+  initialValue = WindowSelect mempty
+
+toggleWindowInSelection :: Window -> X ()
+toggleWindowInSelection win = do
+  (WindowSelect sel) <- XS.get
+  case Map.lookup win sel of
+    Nothing -> do
+      foc <- withWindowSet (return . W.peek)
+
+      cleanup <-
+        setBorderColor
+          ( if Just win == foc
+              then selFocusColor
+              else selColor
+          )
+          [win]
+
+      XS.put $ WindowSelect $ Map.insert win cleanup sel
+    (Just cleanup) -> do
+      XS.put $ WindowSelect $ Map.delete win sel
+      cleanup
+
+addWindowToSelection :: Window -> X ()
+addWindowToSelection win = do
+  (WindowSelect sel) <- XS.get
+  case Map.lookup win sel of
+    Nothing -> do
+      foc <- withWindowSet (return . W.peek)
+
+      cleanup <-
+        setBorderColor
+          ( if Just win == foc
+              then selFocusColor
+              else selColor
+          )
+          [win]
+
+      XS.put $ WindowSelect $ Map.insert win cleanup sel
+    _ -> return ()
+
+removeWindowFromSelection :: Window -> X ()
+removeWindowFromSelection win = do
+  (WindowSelect sel) <- XS.get
+  case Map.lookup win sel of
+    (Just cleanup) -> do
+      XS.put $ WindowSelect $ Map.delete win sel
+      cleanup
+    _ -> return ()
+
+clearWindowSelection :: X ()
+clearWindowSelection = void getAndResetWindowSelection
+
+getAndResetWindowSelection :: X [Window]
+getAndResetWindowSelection = do
+  (WindowSelect mp) <- XS.get
+  sequence_ (Map.elems mp)
+  XS.put (initialValue :: WindowSelect)
+  return (Map.keys mp)
+
 data Workspace = forall a.
   (Typeable a) =>
   Workspace
@@ -522,7 +599,19 @@ readNextWorkspace =
     mt = lift . fromX
 
 readNextLocationSet :: (KeyFeeder m) => MaybeT m [Location]
-readNextLocationSet =
+readNextLocationSet = do
+  (WindowSelect mp) <- MaybeT (Just <$> fromX XS.get)
+  case Map.keys mp of
+    [] -> readNextLocationSet'
+    wins -> do
+      lift $ fromX $ addStringToPendingBuffer "<sel> "
+      fromMaybeTX $
+        mapM windowLocation
+          =<< MaybeT (Just <$> fromX getAndResetWindowSelection)
+
+-- Like readNextLocationSet, but ignores the window selection.
+readNextLocationSet' :: (KeyFeeder m) => MaybeT m [Location]
+readNextLocationSet' =
   readNextKey $ \mask sym str -> do
     macros <- (lift . fromX) $ windowsetMacros <$> getMacros
 
@@ -543,6 +632,10 @@ readNextLocationSet =
       --     (: []) <$> MaybeT (fromX $ pastHistory (ord ch - 0x30))
       -- The current window.
       (_, _, ".") -> (: []) <$> mt getCurrentLocation
+      -- The selected windows in the selection set.
+      (_, _, "#") ->
+        fromMaybeTX $
+          mapM windowLocation =<< MaybeT (Just <$> fromX getAndResetWindowSelection)
       -- The window on the far-left of the screens.
       (_, _, "^") -> (: []) <$> fromMaybeTX farLeftWindow
       -- The windows on the far-right of the screens.
@@ -618,3 +711,11 @@ readNextLocationSet =
   where
     mt :: (KeyFeeder m) => X a -> MaybeT m a
     mt = lift . fromX
+
+wmlLogHook :: X ()
+wmlLogHook = do
+  -- Reset the border colors for the selected window.
+  (WindowSelect (Map.keys -> sel)) <- XS.get
+  foc <- (withWindowSet (return . fromMaybe (0 :: Window) . W.peek) :: X Window)
+  void $ setBorderColor selColor (filter (/= foc) sel)
+  mapM_ (setBorderColor selFocusColor . (: [])) (find (== foc) sel)