1 files changed, 27 insertions, 53 deletions

diff --git a/src/Internal/Marking.hs b/src/Internal/Marking.hs
index f9083d2..229ea02 100644
--- a/src/Internal/Marking.hs
+++ b/src/Internal/Marking.hs
@@ -11,6 +11,8 @@ import System.IO
 import Control.Exception
 import System.Environment
 
+import qualified XMonad.Util.ExtensibleState as XS
+
 import qualified Data.Map as Map
 
 {- Submodule that handles marking windows so they can be jumped back to. -}
@@ -23,65 +25,38 @@ data MarkState =
   , markLast :: Maybe Window
   } deriving (Read, Show)
 
-data MarkContext = MarkContext (IORef MarkState)
-
-readMarkState :: IO MarkState
-readMarkState = do
-  -- TODO(rahm) The mark state should use the extensible state constructs in the
-  -- XState rather than relying on IO.
-  dir <- getEnv "HOME"
-  let markstate = dir </> ".xmonad" </> "markstate"
-  catch
-    (read <$> (hGetContents =<< openFile markstate ReadMode))
-    (\(e :: IOError) -> return (MarkState mempty Nothing))
-
-saveMarkState :: MarkState -> X ()
-saveMarkState ms = do
-  dir <- io $ getEnv "HOME"
-  let markstate = dir </> ".xmonad" </> "markstate"
-  liftIO $ writeFile markstate (show ms)
-
 
-withNewMarkContext :: (MarkContext -> IO a) -> IO a
-withNewMarkContext fn = do
-  ioref <- newIORef =<< readMarkState
-  fn (MarkContext ioref)
+instance ExtensionClass MarkState where
+  initialValue = MarkState Map.empty Nothing
 
-markCurrentWindow :: MarkContext -> Mark -> X ()
-markCurrentWindow (MarkContext ioref) mark = do
+markCurrentWindow :: Mark -> X ()
+markCurrentWindow mark = do
   withFocused $ \win ->
-    liftIO $
-      modifyIORef ioref $ \state@(MarkState {markStateMap = ms}) ->
+    XS.modify $ \state@(MarkState {markStateMap = ms}) ->
         state {
           markStateMap = Map.insert mark win ms
         }
 
-  saveMarkState =<< liftIO (readIORef ioref)
-
-saveLastMark :: MarkContext -> X ()
-saveLastMark (MarkContext ioref) =
-  withFocused $ \win -> do
-      liftIO $ modifyIORef ioref (\state -> state { markLast = Just win })
+saveLastMark :: X ()
+saveLastMark =
+  withFocused $ \win ->
+      XS.modify $ \state -> state { markLast = Just win }
 
-jumpToLast :: MarkContext -> X ()
-jumpToLast ctx@(MarkContext ioref) = do
-  m <- markLast <$> (liftIO $ readIORef ioref)
-  saveLastMark ctx
+jumpToLast :: X ()
+jumpToLast = do
+  m <- markLast <$> XS.get
+  saveLastMark
   mapM_ focus m
 
-  saveMarkState =<< liftIO (readIORef ioref)
-
-jumpToMark :: MarkContext -> Mark -> X ()
-jumpToMark ctx@(MarkContext ioref) mark = do
-  MarkState {markStateMap = m} <- liftIO $ readIORef ioref
+jumpToMark :: Mark -> X ()
+jumpToMark mark = do
+  MarkState {markStateMap = m} <- XS.get
   case Map.lookup mark m of
     Nothing -> return ()
     Just w -> do
-      saveLastMark ctx
+      saveLastMark
       focus w
 
-  saveMarkState =<< liftIO (readIORef ioref)
-
 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
@@ -117,20 +92,19 @@ swapWithFocused winToSwap stackSet =
           mapWindows (
             \w -> if w == winToSwap then focused else w) stackSet
 
-swapWithLastMark :: MarkContext -> X ()
-swapWithLastMark ctx@(MarkContext ioref) = do
-  MarkState {markStateMap = m} <- liftIO $ readIORef ioref
-  m <- markLast <$> (liftIO $ readIORef ioref)
-  saveLastMark ctx
+swapWithLastMark :: X ()
+swapWithLastMark = do
+  m <- markLast <$> XS.get
+  saveLastMark
 
   case m of
     Nothing -> return ()
     Just win -> windows $ swapWithFocused win
 
-swapWithMark :: MarkContext -> Mark -> X ()
-swapWithMark ctx@(MarkContext ioref) mark = do
-  MarkState {markStateMap = m} <- liftIO $ readIORef ioref
-  saveLastMark ctx
+swapWithMark :: Mark -> X ()
+swapWithMark mark = do
+  MarkState {markStateMap = m} <- XS.get
+  saveLastMark
 
   case Map.lookup mark m of
     Nothing -> return ()