Add search mode for text based searching with ocr

7 files changed, 125 insertions, 71 deletions

diff --git a/src/Chelleport.hs b/src/Chelleport.hs
index 96c9fb6..4b44dd1 100644
--- a/src/Chelleport.hs
+++ b/src/Chelleport.hs
@@ -14,23 +14,18 @@ import Chelleport.Control
   )
 import Chelleport.Draw (MonadDraw (windowPosition), cellSize)
 import Chelleport.KeySequence (findMatchPosition, generateGrid, isKeycodeDigit, isValidKey, keycodeToInt, nextChars, toKeyChar)
+import Chelleport.OCR (MonadOCR, getWordsOnScreen)
 import Chelleport.Types
-import Chelleport.Utils (intToCInt)
+import Chelleport.Utils (cIntToInt, intToCInt, isEmpty, isNotEmpty)
 import qualified Chelleport.View
 import Control.Monad (forM_)
 import Control.Monad.IO.Class (MonadIO)
 import Control.Monad.Reader (ReaderT (runReaderT))
+import Data.Char (toLower)
+import Data.List (isInfixOf)
 import Data.Maybe (fromMaybe, isJust)
 import qualified SDL
 
--- run :: IO ()
--- run = do
---   ctx <- initializeContext
---   benchmark "ocr" $ do
---     res <- (`runReaderT` ctx) . runAppM $ getWordsOnScreen
---     print $ "---" ++ show (length res)
---   pure ()
-
 run :: IO ()
 run = do
   ctx <- initializeContext
@@ -45,19 +40,10 @@ run = do
     runAppWithCtx :: (MonadIO m) => DrawContext -> AppM m x -> m x
     runAppWithCtx ctx = (`runReaderT` ctx) . runAppM
 
-initialState :: (Monad m) => m State
+initialState :: (Monad m) => m (State, Maybe AppAction)
 initialState = do
   let cells = fromMaybe (pure undefined) $ generateGrid 0 (rows, columns) hintKeys
-  pure $
-    State
-      { stateGrid = cells,
-        stateKeySequence = [],
-        stateIsMatched = False,
-        stateIsShiftPressed = False,
-        stateIsDragging = False,
-        stateRepetition = 1,
-        stateMode = ModeSearch
-      }
+  pure (defaultAppState {stateGrid = cells}, Just $ SetMode defaultHintsMode)
   where
     rows = 9
     columns = 16
@@ -79,6 +65,12 @@ eventHandler event =
       -- 0-9
       | isKeycodeDigit (eventToKeycode ev) ->
           Just $ UpdateRepetition (fromMaybe 0 $ keycodeToInt $ eventToKeycode ev)
+      -- Enable search mode
+      | withCtrl ev && isKeyPressWith ev SDL.KeycodeS ->
+          Just $ SetMode defaultSearchMode
+      -- Enable hints mode
+      | withCtrl ev && isKeyPressWith ev SDL.KeycodeH ->
+          Just $ SetMode defaultHintsMode
       -- Space / Shift+Space
       | isKeyPressWith ev SDL.KeycodeSpace ->
           if withShift ev
@@ -101,9 +93,18 @@ eventHandler event =
           Just $ UpdateShiftState False
     _ -> Nothing
 
-update :: (MonadAppShell m, MonadDraw m, MonadControl m) => State -> AppAction -> m (State, Maybe AppAction)
--- Act on key inputs
-update state (HandleKeyInput key) = do
+update :: (MonadAppShell m, MonadDraw m, MonadControl m, MonadOCR m) => State -> AppAction -> m (State, Maybe AppAction)
+-- Set mode
+update state (SetMode mode) = do
+  case mode of
+    ModeHints -> pure (state {stateMode = mode}, Nothing)
+    ModeSearch {} -> do
+      wordsOnScreen <- getWordsOnScreen
+      let updatedMode = mode {searchWords = wordsOnScreen, searchFilteredWords = wordsOnScreen}
+      pure (state {stateMode = updatedMode}, Nothing)
+
+-- HINTS MODE: Act on key inputs
+update state@(State {stateMode = ModeHints}) (HandleKeyInput key) = do
   case (toKeyChar key, validChars) of
     (Just keyChar, Just validChars')
       | stateIsMatched state && keyChar `elem` ("HJKL" :: String) -> do
@@ -114,35 +115,51 @@ update state (HandleKeyInput key) = do
           let newKeySequence = stateKeySequence state ++ [keyChar]
           let matchPosition = findMatchPosition newKeySequence $ stateGrid state
           let state' = state {stateKeySequence = newKeySequence, stateIsMatched = isJust matchPosition}
-          pure (state', MoveMousePosition <$> matchPosition)
+          action <- maybe (pure Nothing) (fmap (Just . MoveMousePosition) . getPosition) matchPosition
+          pure (state', action)
     _ -> pure (state, Nothing)
   where
     validChars = nextChars (stateKeySequence state) (stateGrid state)
+    getPosition (row, col) = do
+      (wcell, hcell) <- cellSize state
+      let x = (wcell `div` 2) + wcell * intToCInt col
+      let y = (hcell `div` 2) + hcell * intToCInt row
+      (winx, winy) <- windowPosition
+      pure (cIntToInt $ winx + x, cIntToInt $ winy + y)
     incrementValue = do
       (wcell, hcell) <- cellSize state
       if stateIsShiftPressed state
         then pure (wcell `div` 4, hcell `div` 4)
         else pure (wcell `div` 16, hcell `div` 16)
 
+-- SEARCH MODE: Act on key inputs
+update state@(State {stateMode = ModeSearch {searchWords, searchInputText}}) (HandleKeyInput key) = do
+  case toKeyChar key of
+    Just keyChar -> do
+      let searchText = searchInputText ++ [toLower keyChar]
+      let matches = filterMatches searchText
+      let highlightedWord = if isNotEmpty matches then Just $ head matches else Nothing
+      let updatedMode = (stateMode state) {searchInputText = searchText, searchFilteredWords = matches}
+      pure (state {stateMode = updatedMode}, MoveMousePosition . wordPosition <$> highlightedWord)
+    _ -> do
+      pure (state, Nothing)
+  where
+    wordPosition w = (cIntToInt $ matchStartX w, cIntToInt $ matchStartY w)
+    filterMatches text
+      | isEmpty text = searchWords
+      | otherwise = filter (isInfixOf text . map toLower . matchText) searchWords
+
 -- Move mouse incrementally
 update state (IncrementMouseCursor (incX, incY)) = do
   (curX, curY) <- getMousePointerPosition
-  let count = intToCInt $ case stateRepetition state of 0 -> 1; n -> n
-  moveMousePointer (curX + count * intToCInt incX) (curY + count * intToCInt incY)
-  pure (state {stateRepetition = 1}, Nothing)
+  let count = case stateRepetition state of 0 -> 1; n -> n
+  let pos = (cIntToInt curX + count * incX, cIntToInt curY + count * incY)
+  pure (state {stateRepetition = 1}, Just $ MoveMousePosition pos)
 
 -- Move mouse to given position
-update state (MoveMousePosition (row, col)) = do
-  (x, y) <- getPosition
-  moveMousePointer x y
+update state (MoveMousePosition (x, y)) = do
+  moveMousePointer (intToCInt x) (intToCInt y)
   pure (state, Nothing)
-  where
-    getPosition = do
-      (wcell, hcell) <- cellSize state
-      let x = (wcell `div` 2) + wcell * intToCInt col
-      let y = (hcell `div` 2) + hcell * intToCInt row
-      (winx, winy) <- windowPosition
-      pure (winx + x, winy + y)
 
 -- Reset entered key sequence and state
 update state ResetKeys = do
diff --git a/src/Chelleport/AppShell.hs b/src/Chelleport/AppShell.hs
index ba31982..2e95a85 100644
--- a/src/Chelleport/AppShell.hs
+++ b/src/Chelleport/AppShell.hs
@@ -35,18 +35,19 @@ type EventHandler state appAction = state -> SDL.Event -> Maybe appAction
 
 type View m state = state -> m ()
 
-type Initializer m state = m state
+type Initializer m state appAction = m (state, Maybe appAction)
 
 setupAppShell ::
   (MonadIO m) =>
   DrawContext ->
-  Initializer m state ->
+  Initializer m state appAction ->
   Update m state appAction ->
   EventHandler state appAction ->
   View m state ->
   m ()
-setupAppShell (DrawContext {ctxRenderer = renderer}) getInitState update eventHandler draw =
-  getInitState >>= appLoop
+setupAppShell (DrawContext {ctxRenderer = renderer}) getInitState update eventHandler draw = do
+  state <- getInitState >>= evalUpdateResult
+  appLoop state
   where
     appLoop currentState = do
       SDL.rendererDrawColor renderer $= colorBackground
diff --git a/src/Chelleport/Control.hs b/src/Chelleport/Control.hs
index 80be6f8..f7a9c66 100644
--- a/src/Chelleport/Control.hs
+++ b/src/Chelleport/Control.hs
@@ -64,10 +64,10 @@ instance (MonadIO m) => MonadControl (AppM m) where
       X11.sync display False
 
 isKeyPressed :: SDL.KeyboardEventData -> Bool
-isKeyPressed = (== SDL.Pressed) . SDL.keyboardEventKeyMotion
+isKeyPressed = (SDL.Pressed ==) . SDL.keyboardEventKeyMotion
 
 isKeyRelease :: SDL.KeyboardEventData -> Bool
-isKeyRelease = (== SDL.Released) . SDL.keyboardEventKeyMotion
+isKeyRelease = (SDL.Released ==) . SDL.keyboardEventKeyMotion
 
 eventToKeycode :: SDL.KeyboardEventData -> SDL.Keycode
 eventToKeycode = SDL.keysymKeycode . SDL.keyboardEventKeysym
diff --git a/src/Chelleport/Draw.hs b/src/Chelleport/Draw.hs
index 3f67848..3720ba0 100644
--- a/src/Chelleport/Draw.hs
+++ b/src/Chelleport/Draw.hs
@@ -14,6 +14,7 @@ class (Monad m) => MonadDraw m where
   drawLine :: (CInt, CInt) -> (CInt, CInt) -> m ()
   drawText :: (CInt, CInt) -> Color -> Text -> m (CInt, CInt)
   drawCircle :: Int -> (CInt, CInt) -> m ()
+  fillRect :: (CInt, CInt) -> (CInt, CInt) -> m ()
   setDrawColor :: Color -> m ()
   windowSize :: m (CInt, CInt)
   windowPosition :: m (CInt, CInt)
@@ -27,6 +28,11 @@ instance (MonadIO m) => MonadDraw (AppM m) where
     renderer <- asks ctxRenderer
     SDL.rendererDrawColor renderer $= color
 
+  fillRect (x, y) (w, h) = do
+    renderer <- asks ctxRenderer
+    let rect = SDL.Rectangle (SDL.P $ SDL.V2 x y) (SDL.V2 w h)
+    SDL.fillRect renderer (Just rect)
+
   drawText (x, y) color text = do
     DrawContext {ctxRenderer = renderer, ctxFont = font} <- ask
     surface <- TTF.blended font color text
@@ -65,6 +71,9 @@ instance (MonadIO m) => MonadDraw (AppM m) where
     SDL.V2 x y <- asks ctxWindow >>= SDL.getWindowAbsolutePosition
     pure (x, y)
 
+fillRectVertices :: (MonadDraw m) => (CInt, CInt) -> (CInt, CInt) -> m ()
+fillRectVertices (x1, y1) (x2, y2) = fillRect (x1, y1) (x2 - x1, y2 - y1)
+
 cellSize :: (MonadDraw m) => State -> m (CInt, CInt)
 cellSize (State {stateGrid}) = do
   (width, height) <- windowSize
diff --git a/src/Chelleport/OCR.hs b/src/Chelleport/OCR.hs
index 496b6a0..f6bc5b9 100644
--- a/src/Chelleport/OCR.hs
+++ b/src/Chelleport/OCR.hs
@@ -1,8 +1,8 @@
-module Chelleport.OCR (getWordsOnScreen) where
+module Chelleport.OCR (MonadOCR (..)) where
 
 import Chelleport.Types
 import Control.Monad.IO.Class (MonadIO (liftIO))
-import Control.Monad.RWS (asks)
+import Control.Monad.RWS (MonadReader (ask))
 import qualified Data.ByteString as BS
 import Foreign (Bits (shiftR), Ptr, Storable (peek, pokeByteOff), alloca, allocaBytes, peekArray, (.&.))
 import Foreign.C (CInt, CString, newCString)
@@ -15,48 +15,39 @@ import System.IO (hPutStrLn)
 import System.IO.Temp (emptySystemTempFile)
 
 foreign import ccall unsafe "libchelleport.h findWordCoordinates"
-  c_findWordCoordinates :: CString -> Ptr CInt -> IO (Ptr OCRMatch)
+  c_getAllWordCoordinates :: CString -> Ptr CInt -> IO (Ptr OCRMatch)
 
 class (Monad m) => MonadOCR m where
   getWordsOnScreen :: m [OCRMatch]
 
 instance (MonadIO m) => MonadOCR (AppM m) where
   getWordsOnScreen = do
-    SDL.V2 width height <- asks ctxWindow >>= SDL.get . SDL.windowSize
-    SDL.V2 x y <- asks ctxWindow >>= SDL.getWindowAbsolutePosition
+    ctx <- ask
+    SDL.V2 width height <- SDL.get . SDL.windowSize . ctxWindow $ ctx
+    SDL.V2 x y <- SDL.getWindowAbsolutePosition . ctxWindow $ ctx
     liftIO $ do
-      imgFilePath <- liftIO $ createTemporaryScreenshot (x, y) (width, height)
+      imgFilePath <- liftIO $ createTemporaryScreenshot ctx (x, y) (width, height)
       findWordCoordinates imgFilePath <* removeFile imgFilePath
 
 findWordCoordinates :: String -> IO [OCRMatch]
 findWordCoordinates imgPath = alloca $ \sizePtr -> do
   imgPathC <- newCString imgPath
-  arrayPtr <- c_findWordCoordinates imgPathC sizePtr
+  arrayPtr <- c_getAllWordCoordinates imgPathC sizePtr
 
   size <- peek sizePtr
   peekArray (fromIntegral size) arrayPtr
 
-createTemporaryScreenshot :: (CInt, CInt) -> (CInt, CInt) -> IO String
-createTemporaryScreenshot offset size = do
+createTemporaryScreenshot :: DrawContext -> (CInt, CInt) -> (CInt, CInt) -> IO String
+createTemporaryScreenshot ctx offset size = do
   tmpFilePath <- emptySystemTempFile "chelleport-screenshot.png"
-  screenshot tmpFilePath offset size
+  screenshot ctx tmpFilePath offset size
   pure tmpFilePath
 
-screenshot :: String -> (CInt, CInt) -> (CInt, CInt) -> IO ()
-screenshot filename (offsetX, offsetY) (width, height) = do
-  dpy <- X11.openDisplay ""
-  root <- X11.rootWindow dpy (X11.defaultScreen dpy)
+screenshot :: DrawContext -> String -> (CInt, CInt) -> (CInt, CInt) -> IO ()
+screenshot (DrawContext {ctxX11Display = display}) filename (offsetX, offsetY) (width, height) = do
+  root <- X11.rootWindow display (X11.defaultScreen display)
 
-  image <-
-    X11.getImage
-      dpy
-      root
-      offsetX
-      offsetY
-      (fromIntegral width)
-      (fromIntegral height)
-      (fromIntegral X11.allPlanes_aux)
-      X11.zPixmap
+  image <- X11.getImage display root offsetX offsetY (fromIntegral width) (fromIntegral height) (fromIntegral X11.allPlanes_aux) X11.zPixmap
 
   allocaBytes (fromIntegral $ width * height * 3) $ \ptr -> do
     let getPixel :: CInt -> CInt -> IO ()
@@ -74,7 +65,6 @@ screenshot filename (offsetX, offsetY) (width, height) = do
     savePPMFile filename (fromIntegral width) (fromIntegral height) rgbData
 
   X11.destroyImage image
-  X11.closeDisplay dpy
 
 savePPMFile :: FilePath -> Int -> Int -> BS.ByteString -> IO ()
 savePPMFile path width height rgbData = withFile path WriteMode $ \h -> do
diff --git a/src/Chelleport/Types.hs b/src/Chelleport/Types.hs
index cb580e0..9894e54 100644
--- a/src/Chelleport/Types.hs
+++ b/src/Chelleport/Types.hs
@@ -17,9 +17,21 @@ type KeySequence = [Char]
 
 type KeyGrid = [[Cell]]
 
-data Mode = ModeHints | ModeSearch
+data Mode
+  = ModeHints
+  | ModeSearch
+      { searchWords :: [OCRMatch],
+        searchFilteredWords :: [OCRMatch],
+        searchInputText :: String
+      }
   deriving (Show, Eq)
 
+defaultSearchMode :: Mode
+defaultSearchMode = ModeSearch {searchWords = [], searchFilteredWords = [], searchInputText = ""}
+
+defaultHintsMode :: Mode
+defaultHintsMode = ModeHints
+
 data State = State
   { stateGrid :: KeyGrid,
     stateKeySequence :: KeySequence,
@@ -31,6 +43,18 @@ data State = State
   }
   deriving (Show, Eq)
 
+defaultAppState :: State
+defaultAppState =
+  State
+    { stateGrid = [],
+      stateKeySequence = "",
+      stateIsMatched = False,
+      stateIsShiftPressed = False,
+      stateIsDragging = False,
+      stateRepetition = 1,
+      stateMode = ModeHints
+    }
+
 data AppAction
   = ChainMouseClick MouseButtonType
   | HandleKeyInput SDL.Keycode
@@ -44,6 +68,7 @@ data AppAction
   | TriggerMouseClick MouseButtonType
   | UpdateShiftState Bool
   | UpdateRepetition Int
+  | SetMode Mode
   deriving (Show, Eq)
 
 data DrawContext = DrawContext
@@ -66,7 +91,7 @@ data OCRMatch = OCRMatch
     matchEndY :: !CInt,
     matchText :: !String
   }
-  deriving (Show)
+  deriving (Show, Eq)
 
 instance Storable OCRMatch where
   sizeOf _ = 4 * sizeOf (undefined :: CInt) + sizeOf (undefined :: Ptr CChar)
diff --git a/src/Chelleport/View.hs b/src/Chelleport/View.hs
index fe06fd7..1a4f1f8 100644
--- a/src/Chelleport/View.hs
+++ b/src/Chelleport/View.hs
@@ -10,7 +10,19 @@ import qualified Data.Text as Text
 import Foreign.C (CInt)
 
 render :: (MonadDraw m) => State -> m ()
-render state = do
+render state = case stateMode state of
+  ModeHints -> renderHintsView state
+  ModeSearch {searchFilteredWords} -> renderSearchView state searchFilteredWords
+
+renderSearchView :: (MonadDraw m) => State -> [OCRMatch] -> m ()
+renderSearchView state matches = do
+  renderGridLines state
+  setDrawColor colorWhite
+  forM_ matches $ \(OCRMatch {matchStartX, matchStartY, matchEndX, matchEndY}) -> do
+    fillRectVertices (matchStartX, matchStartY) (matchEndX, matchEndY)
+
+renderHintsView :: (MonadDraw m) => State -> m ()
+renderHintsView state = do
   renderGridLines state
 
   (wcell, hcell) <- cellSize state