1 files changed, 39 insertions, 1 deletions
diff --git a/src/Wetterhorn/Core/Keys.hs b/src/Wetterhorn/Core/Keys.hs
index d82ac4c..54d7125 100644
--- a/src/Wetterhorn/Core/Keys.hs
+++ b/src/Wetterhorn/Core/Keys.hs
@@ -3,12 +3,14 @@ module Wetterhorn.Core.Keys where
 import Control.Monad (forever, void, when)
 import Control.Monad.Cont.Class
 import Control.Monad.IO.Class
-import Control.Monad.State (MonadState (get, put), MonadTrans (lift), StateT, evalStateT, gets)
+import Control.Monad.State (MonadState (get, put), MonadTrans (lift), StateT, evalStateT, gets, modify)
 import Control.Monad.Trans.Cont
 import Data.Bits
 import Data.Word
+import Wetterhorn.Core.ButtonEvent (ButtonEvent)
 import Wetterhorn.Core.KeyEvent
 import qualified Wetterhorn.Core.KeyEvent as KeyEvent
+import qualified Wetterhorn.Core.ButtonEvent as ButtonEvent
 import Wetterhorn.Core.W
 import Wetterhorn.Foreign.WlRoots (wlrSeatKeyboardNotifyKey, wlrSeatSetKeyboard)
 
@@ -114,6 +116,42 @@ putKeyHandler handler = do
             }
       }
 
+nextButtonEvent :: KeysM ButtonEvent
+nextButtonEvent = do
+  st <- KeysM get
+  KeysM $
+    shiftT $ \h ->
+      lift $ lift $ putButtonHandler (\ev -> evalStateT (h ev) st)
+  where
+    putButtonHandler h = do
+      modify $ \st -> st {currentHooks = (currentHooks st) {buttonHook = h}}
+
+nextButtonOrKeyEvent :: KeysM (Either ButtonEvent KeyEvent)
+nextButtonOrKeyEvent = do
+  st <- KeysM get
+  KeysM $
+    shiftT $ \rest -> 
+      lift $ lift $ do
+        putButtonHandler (\ev -> evalStateT (rest (Left ev)) st)
+        handleContinuation st (\ev -> evalStateT (rest (Right ev)) st)
+
+  where
+    putButtonHandler h = do
+      modify $ \st -> st {currentHooks = (currentHooks st) {buttonHook = h}}
+
+nextButtonOrKeyPress :: KeysM (Either ButtonEvent KeyEvent)
+nextButtonOrKeyPress = do
+  ev <- nextButtonOrKeyEvent
+  case ev of
+    Left bev | ButtonEvent.state bev == ButtonEvent.ButtonPressed -> return ev
+    Left bev -> forwardButtonEvent bev >> nextButtonOrKeyPress
+    Right kev | KeyEvent.state kev == KeyEvent.KeyPressed -> return ev
+    Right kev -> forwardEvent kev >> nextButtonOrKeyPress
+
+ where
+   forwardButtonEvent _ = return ()
+  
+
 -- | Returns the next KeyPressed event. This is likely what 90% of use cases
 -- want rather than nextKeyEvent.
 nextKeyPress :: KeysM KeyEvent