Add extensible state and macros. Fix some bugs with keys.

4 files changed, 252 insertions, 46 deletions

diff --git a/src/Config.hs b/src/Config.hs
index 04d13c6..e71f48a 100644
--- a/src/Config.hs
+++ b/src/Config.hs
@@ -1,20 +1,14 @@
 module Config (config) where
 
+import Text.Printf
 import Control.Monad.IO.Class
 import Control.Monad.Loops
 import Wetterhorn.Core.KeyEvent qualified as KeyEvent
 import Wetterhorn.Core.Keys
 import Wetterhorn.Core.W
+import Wetterhorn.Keys.Macros
 import Wetterhorn.Layout.Full
 
-alsoLog :: KeyContinuation -> W ()
-alsoLog kh =
-  putKeyHandler
-    ( \ke -> do
-        liftIO $ putStrLn $ (: []) $ KeyEvent.codepoint ke
-        kh ke
-    )
-
 config :: Config WindowLayout
 config =
   defaultConfig
@@ -25,15 +19,21 @@ config =
           },
       layout = WindowLayout Full,
       resetHook = do
-        useKeysWithContinuation alsoLog $ do
+        useKeysWithContinuation recordMacroContinuation $ do
           kp <- nextKeyPress
 
+          bind kp (Mod1 .+ 'q') macroKeyBind
+
+          bind kp (weak $ Mod1 .+ '@') replayMacroKeybind
+
           bind kp (Mod1 .+ 'r') (shellExec "wofi --show run")
 
           bind kp (Shift .+ Mod1 .+ 'R') requestHotReload
 
           bind kp (Mod1 .+ 't') (shellExec "alacritty")
 
+          bind kp (Mod1 .+ 'n') (return () :: W ())
+
           bind kp (weak $ Mod1 .+ '∫') (shellExec "gxmessage hi")
 
           bind kp (Mod1 .+ 'p') $ do
diff --git a/src/Wetterhorn/Core/Keys.hs b/src/Wetterhorn/Core/Keys.hs
index b979048..f7fb497 100644
--- a/src/Wetterhorn/Core/Keys.hs
+++ b/src/Wetterhorn/Core/Keys.hs
@@ -1,17 +1,14 @@
 module Wetterhorn.Core.Keys where
 
-import Control.Monad (forever, join, void, when)
+import Control.Monad (forever, void, when)
 import Control.Monad.Cont.Class
-import Control.Monad.Fix (fix)
 import Control.Monad.IO.Class
-import Control.Monad.Reader (ReaderT (runReaderT))
-import Control.Monad.Reader.Class
 import Control.Monad.State (MonadState (get, put), MonadTrans (lift), StateT, evalStateT, gets)
 import Control.Monad.Trans.Cont
 import Data.Bits
 import Data.Word
 import Wetterhorn.Core.KeyEvent
-import Wetterhorn.Core.KeyEvent qualified as KeyEvent
+import qualified Wetterhorn.Core.KeyEvent as KeyEvent
 import Wetterhorn.Core.W
 import Wetterhorn.Foreign.WlRoots (wlrSeatKeyboardNotifyKey, wlrSeatSetKeyboard)
 
@@ -73,29 +70,14 @@ instance Wlike KeysM where
 
 type KeyContinuation = KeyEvent -> W ()
 
--- Return type in the keysM monad.
--- data KeysMR a = NextKey (KeysM a) | Lift a | Continue
-
--- keysWithHandler :: (KeyContinuation -> W ()) -> KeysM a -> KeyContinuation
--- keysWithHandler nextAction keysM = fix $ \top ke -> keys' top keysM ke
---   where
---     keys' top (KeysM fn) ke = do
---       e <- fn top ke
---       case e of
---         NextKey keysM' -> nextAction (keys' top keysM')
---         Lift _ -> return ()
---         _ -> nextAction top
-
--- keys :: KeysM a -> KeyEvent -> W ()
--- keys = keysWithHandler putKeyHandler
-
 useKeysWithContinuation :: (KeyContinuation -> W ()) -> KeysM () -> W ()
-useKeysWithContinuation continuation km@(KeysM c) =
-  evalStateT (evalContT (forever c)) (KeysState km continuation)
+useKeysWithContinuation continuation (forever -> km@(KeysM c)) =
+  evalStateT (evalContT c) (KeysState km continuation)
 
 useKeys :: KeysM () -> W ()
 useKeys = useKeysWithContinuation putKeyHandler
 
+-- | Returns the next key event.
 nextKeyEvent :: KeysM KeyEvent
 nextKeyEvent = do
   st <- KeysM $ lift get
@@ -106,8 +88,20 @@ nextKeyEvent = do
             handleContinuation st (\kp -> evalStateT (keyHandler kp) st)
       )
 
+-- | Discards the rest of the continuation and stars again from the top. Useful
+-- for keybinds where once the key is handled, just continue to the top.
 continue :: KeysM ()
-continue = join (KeysM (gets keysTop))
+continue = do
+  st <- KeysM $ lift get
+  let (KeysM topCont) = keysTop st
+
+  -- This shift discards the rest of the computation and instead returns to the
+  -- top of the handler.
+  KeysM $ shiftT (\_ -> lift (evalContT topCont))
+
+-- | Returns the "top" continuation.
+getTop :: KeysM (KeysM ())
+getTop = KeysM (gets keysTop)
 
 putKeyHandler :: KeyContinuation -> W ()
 putKeyHandler handler = do
diff --git a/src/Wetterhorn/Core/W.hs b/src/Wetterhorn/Core/W.hs
index dfa0753..79da19b 100644
--- a/src/Wetterhorn/Core/W.hs
+++ b/src/Wetterhorn/Core/W.hs
@@ -3,23 +3,30 @@
 module Wetterhorn.Core.W where
 
 import Control.Arrow (Arrow (first))
+import Control.Monad ((<=<))
 import Control.Monad.RWS (MonadIO (liftIO), MonadReader (..), MonadState, modify)
 import Control.Monad.Reader (ReaderT (runReaderT))
-import Control.Monad.State (StateT (runStateT))
+import Control.Monad.State (StateT (runStateT), gets, modify')
 import Control.Monad.Trans.Maybe
 import Data.Data (Typeable, cast)
+import Data.Default.Class (Default, def)
 import Data.Kind (Constraint, Type)
+import Data.Map (Map)
+import qualified Data.Map as M
+import Data.Maybe (fromMaybe, mapMaybe)
+import Data.Proxy
 import Data.Set (Set)
-import Data.Set qualified as Set
+import qualified Data.Set as Set
 import Foreign (Ptr, StablePtr, intPtrToPtr, ptrToIntPtr)
 import Text.Read
+import Type.Reflection (someTypeRep)
 import Wetterhorn.Core.KeyEvent
 import Wetterhorn.Core.SurfaceEvent
 import Wetterhorn.Foreign.ForeignInterface (ForeignInterface)
-import Wetterhorn.Foreign.ForeignInterface qualified as ForeignInterface
+import qualified Wetterhorn.Foreign.ForeignInterface as ForeignInterface
 import Wetterhorn.Foreign.WlRoots (Surface, WlrSeat)
 import Wetterhorn.StackSet hiding (layout)
-import Wetterhorn.StackSet qualified as StackSet
+import qualified Wetterhorn.StackSet as StackSet
 
 data RationalRect = RationalRect Rational Rational Rational Rational
 
@@ -97,7 +104,7 @@ handleWindowMessage m (WindowLayout l) = WindowLayout <$> handleMessage m l
 readWindowLayout :: WindowLayout -> String -> WindowLayout
 readWindowLayout (WindowLayout l) s
   | (Just x) <- readLayout s =
-      WindowLayout (asTypeOf x l)
+    WindowLayout (asTypeOf x l)
 readWindowLayout l _ = l
 
 -- | Serializes a window layout to a string.
@@ -151,6 +158,49 @@ data Config l = Config
     resetHook :: W ()
   }
 
+-- | Typeclass defining the set of types which can be used as state extensions
+-- to the W monad. These state extensions may be persistent or not.
+--
+-- There are default implementations for all methods if the type implements
+-- Read, Show and Default,
+class (Typeable a) => ExtensionClass a where
+  -- | The initial value used for the first time an extension is 'gotten' or
+  -- demarshalling fails.
+  initialValue :: a
+
+  -- | Transforms a type into a string. If the type cannot be marshalled, this
+  -- function should return Nothing.
+  --
+  -- If a type cannot be marshalled, it cannot persist across hot reloads.
+  marshalExtension :: a -> Maybe String
+
+  -- | Reads an extension from a string. If this type is not marshallable or
+  -- reading fails, this function should return Nothing.
+  demarshalExtension :: String -> Maybe a
+
+  -- | If the type implements Default, use the default implementation.
+  default initialValue :: (Default a) => a
+  initialValue = def
+
+  -- | If the type implements Show, use show for the marshalling.
+  default marshalExtension :: (Show a) => a -> Maybe String
+  marshalExtension = Just . show
+
+  -- | If the type implements Read, use read for the demarshalling.
+  default demarshalExtension :: (Read a) => String -> Maybe a
+  demarshalExtension = readMaybe
+
+data StateExtension where
+  StateExtension :: (ExtensionClass a) => a -> StateExtension
+
+-- State as it is marshalled.
+data MarshalledState
+  = MarshalledState
+      (StackSet ScreenId ScreenDetail Tag String Window)
+      (Set Window)
+      [(String, String)]
+  deriving (Show, Read)
+
 data State = State
   { -- The datastructure containing the state of the windows.
     mapped :: StackSet ScreenId ScreenDetail Tag WindowLayout Window,
@@ -159,7 +209,9 @@ data State = State
     -- | Current set of hooks. The initial hooks are provided by the
     -- configuration, but the hooks can change during operation. This is how key
     -- sequences can be mapped.
-    currentHooks :: Hooks
+    currentHooks :: Hooks,
+    -- | Map from the typerep string to the state extension.
+    extensibleState :: Map String (Either String StateExtension)
   }
 
 -- | Initializes a "cold" state from a configuration. A cold state is the
@@ -173,6 +225,7 @@ initColdState Config {layout = layout, hooks = hooks} =
       )
       mempty
       hooks
+      mempty
 
 -- | Marshals the serializable parts of the state to a string. This happens
 -- during a hot-reload where some state must be saved to persist across hot
@@ -181,21 +234,66 @@ marshalState :: State -> String
 marshalState
   ( State
       { mapped = mapped,
-        allWindows = allWindows
+        allWindows = allWindows,
+        extensibleState = xs
       }
     ) =
-    show
-      ( mapLayout serializeWindowLayout mapped,
+    show $
+      MarshalledState
+        (mapLayout serializeWindowLayout mapped)
         allWindows
-      )
+        (mapMaybe (\(k, v) -> (k,) <$> doMarshalEx v) (M.toList xs))
+    where
+      doMarshalEx (Left s) = Just s
+      doMarshalEx (Right (StateExtension a)) = marshalExtension a
+
+-- | Puts a state extension.
+xput :: forall a m. (ExtensionClass a, Wlike m) => a -> m ()
+xput val = liftW $ do
+  modify' $ \state@State {extensibleState = extensibleState} ->
+    state
+      { extensibleState =
+          M.insert
+            ( show $ someTypeRep (Proxy :: Proxy a)
+            )
+            (Right $ StateExtension val)
+            extensibleState
+      }
+
+-- | Modifies a state extension.
+xmodify :: forall a m. (ExtensionClass a, Wlike m) => (a -> a) -> m ()
+xmodify fn = xput . fn =<< xget
+
+-- | Modifies a state extension in the monadic context.
+xmodifyM :: forall a m. (ExtensionClass a, Wlike m) => (a -> m a) -> m ()
+xmodifyM fn = (xput <=< fn) =<< xget
+
+-- | Gets a state extension.
+xget :: forall a m. (ExtensionClass a, Wlike m) => m a
+xget = do
+  xs <- liftW $ gets extensibleState
+  case M.lookup (show $ someTypeRep (Proxy :: Proxy a)) xs of
+    Just (Right (StateExtension a)) -> return (fromMaybe initialValue (cast a))
+    Just (Left str) ->
+      let v = fromMaybe initialValue (demarshalExtension str)
+       in xput v >> return v
+    Nothing ->
+      xput (initialValue :: a) >> return initialValue
+
+xgets :: forall a b m. (ExtensionClass a, Wlike m) => (a -> b) -> m b
+xgets fn = fn <$> xget
 
 -- | Demarshals the string from "marshalState" into a state. Uses the provided
 -- config to fill out non-persistent parts of the state.
 demarshalState :: Config WindowLayout -> String -> State
 demarshalState Config {hooks = hooks, layout = layout} str =
-  State mapped allWindows hooks
+  State mapped allWindows hooks xs
   where
-    (mapLayout (readWindowLayout layout) -> mapped, allWindows) = read str
+    ( MarshalledState
+        (mapLayout (readWindowLayout layout) -> mapped)
+        allWindows
+        (fmap Left . M.fromList -> xs)
+      ) = read str
 
 -- | This is _the_ main monad used for Wetterhorn operations. Contains
 -- everything required to operate. Contains the state, configuration and
@@ -231,7 +329,7 @@ wio :: IO a -> W a
 wio = liftIO
 
 -- | Type class to lift an arbitrary 'W' computation into another monad.
-class Wlike m where
+class (Monad m) => Wlike m where
   liftW :: W a -> m a
 
 -- | Trivial instance of W for Wlike.
diff --git a/src/Wetterhorn/Keys/Macros.hs b/src/Wetterhorn/Keys/Macros.hs
new file mode 100644
index 0000000..9af9827
--- /dev/null
+++ b/src/Wetterhorn/Keys/Macros.hs
@@ -0,0 +1,114 @@
+module Wetterhorn.Keys.Macros where
+
+import Control.Monad.IO.Class (MonadIO (liftIO))
+import Data.Default.Class
+import Data.IORef (newIORef, readIORef, writeIORef)
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Maybe (fromMaybe)
+import Data.Word
+import Foreign (Ptr)
+import Text.Printf
+import Wetterhorn.Core.KeyEvent
+import Wetterhorn.Core.Keys
+import Wetterhorn.Core.W
+import Wetterhorn.Foreign.WlRoots (WlrInputDevice)
+
+data RecordedKey = RecordedKey
+  { recordedKeycode :: Word32,
+    recordedState :: KeyState,
+    recordedModifiers :: Word32,
+    recordedKeysym :: Word32,
+    recordedCodepoint :: Char
+  }
+  deriving (Read, Show)
+
+data MacrosState = MacrosState
+  { macros :: Map Char [RecordedKey],
+    currentlyRecording :: Maybe Char
+  }
+  deriving (Read, Show)
+
+instance Default MacrosState where
+  def = MacrosState mempty def
+
+instance ExtensionClass MacrosState
+
+startMacroRecording :: (Wlike m) => Char -> m ()
+startMacroRecording ch =
+  xmodify
+    ( \m@MacrosState {macros = macros} ->
+        m
+          { macros = Map.delete ch macros,
+            currentlyRecording = Just ch
+          }
+    )
+
+stopMacroRecording :: (Wlike m) => m ()
+stopMacroRecording = xmodify (\m -> m {currentlyRecording = Nothing})
+
+currentlyRecordingMacro :: (Wlike m) => m (Maybe Char)
+currentlyRecordingMacro = xgets currentlyRecording
+
+macroKeyBind :: KeysM ()
+macroKeyBind = do
+  curRec <- currentlyRecordingMacro
+  case curRec of
+    Just ch -> do
+      liftIO $ printf "End Recording %s\n" (show ch)
+      stopMacroRecording
+    Nothing -> do
+      kp' <- nextKeyPress
+      liftIO $ printf "Recording %s" (show $ codepoint kp')
+      startMacroRecording (codepoint kp')
+
+replayMacroKeybind :: KeysM ()
+replayMacroKeybind =
+  replayMacro' =<< nextKeyPress
+
+replayMacro' :: KeyEvent -> KeysM ()
+replayMacro' ke = do
+  top <- getTop
+  liftW $ replayMacro top (device ke) (codepoint ke)
+
+replayMacro :: KeysM () -> Ptr WlrInputDevice -> Char -> W ()
+replayMacro top inputDevice ch = do
+  macro <- xgets $ (fromMaybe [] . Map.lookup ch) . macros
+  case macro of
+    [] -> return ()
+    ks -> do
+      ioref <- wio $ newIORef (reverse ks)
+      useKeysWithContinuation
+        ( \fn -> do
+            ks' <- wio $ readIORef ioref
+            case ks' of
+              (k : ks'') -> do
+                wio $ writeIORef ioref ks''
+                fn (toKeyEvent k)
+              _ -> return ()
+        )
+        top
+  where
+    toKeyEvent (RecordedKey kc st mo sym cp) =
+      KeyEvent 0 kc st mo sym cp inputDevice
+
+pushMacroKey :: (Wlike m) => KeyEvent -> m ()
+pushMacroKey ke = do
+  cur <- xgets currentlyRecording
+  whenJust cur $ \ch ->
+    let recordedKey = toRecordedKey ke
+     in xmodify $ \m@MacrosState {macros = macros} ->
+          m {macros = Map.insertWith (++) ch [recordedKey] macros}
+  where
+    whenJust (Just a) fn = fn a
+    whenJust _ _ = return ()
+
+    toRecordedKey (KeyEvent _ c s m sym cp _) = RecordedKey c s m sym cp
+
+recordMacroContinuation :: (KeyEvent -> W ()) -> W ()
+recordMacroContinuation cont =
+  putKeyHandler
+    ( \ke -> do
+        pushMacroKey ke
+        cont ke
+    )