System/Console/Haskeline/Vi.hs

    #if __GLASGOW_HASKELL__ < 802
    {-# OPTIONS_GHC -Wno-redundant-constraints #-}
    #endif
    module System.Console.Haskeline.Vi where

    import System.Console.Haskeline.Command
    import System.Console.Haskeline.Monads
    import System.Console.Haskeline.Key
    import System.Console.Haskeline.Command.Completion
    import System.Console.Haskeline.Command.History
    import System.Console.Haskeline.Command.KillRing
    import System.Console.Haskeline.Command.Undo
    import System.Console.Haskeline.LineState
    import System.Console.Haskeline.InputT

    import Data.Char
    import Control.Monad(liftM)
    import Control.Monad.Catch (MonadMask)

    type EitherMode = Either CommandMode InsertMode

    type SavedCommand m = Command (ViT m) (ArgMode CommandMode) EitherMode

    data ViState m = ViState {
                lastCommand :: SavedCommand m,
                lastSearch :: [Grapheme]
             }

    emptyViState :: Monad m => ViState m
    emptyViState = ViState {
                lastCommand = return . Left . argState,
                lastSearch = []
            }

    type ViT m = StateT (ViState m) (InputCmdT m)

    type InputCmd s t = forall m . (MonadIO m, MonadMask m) => Command (ViT m) s t
    type InputKeyCmd s t = forall m . (MonadIO m, MonadMask m) => KeyCommand (ViT m) s t

    viKeyCommands :: InputKeyCmd InsertMode (Maybe String)
    viKeyCommands = choiceCmd [
                    simpleChar '\n' +> finish
                    , ctrlChar 'd' +> eofIfEmpty
                    , simpleInsertions >+> viCommands
                    , simpleChar '\ESC' +> change enterCommandMode
                        >|> viCommandActions
                    ]

    viCommands :: InputCmd InsertMode (Maybe String)
    viCommands = keyCommand viKeyCommands

    simpleInsertions :: InputKeyCmd InsertMode InsertMode
    simpleInsertions = choiceCmd
                    [  simpleKey LeftKey +> change goLeft
                       , simpleKey RightKey +> change goRight
                       , simpleKey Backspace +> change deletePrev
                       , simpleKey Delete +> change deleteNext
                       , simpleKey Home +> change moveToStart
                       , simpleKey End +> change moveToEnd
                       , insertChars
                       , ctrlChar 'l' +> clearScreenCmd
                       , simpleKey UpKey +> historyBack
                       , simpleKey DownKey +> historyForward
                       , simpleKey SearchReverse +> searchForPrefix Reverse
                       , simpleKey SearchForward +> searchForPrefix Forward
                       , searchHistory
                       , simpleKey KillLine +> killFromHelper (SimpleMove moveToStart)
                       , ctrlChar 'w' +> killFromHelper wordErase
                       , completionCmd (simpleChar '\t')
                       ]

    insertChars :: InputKeyCmd InsertMode InsertMode
    insertChars = useChar $ loop []
        where
            loop ds d = change (insertChar d) >|> keyChoiceCmd [
                            useChar $ loop (d:ds)
                            , withoutConsuming (storeCharInsertion (reverse ds))
                            ]
            storeCharInsertion s = storeLastCmd $ change (applyArg
                                                            $ withCommandMode $ insertString s)
                                                    >|> return . Left

    -- If we receive a ^D and the line is empty, return Nothing
    -- otherwise, act like '\n' (mimicing how Readline behaves)
    eofIfEmpty :: (Monad m, Save s, Result s) => Command m s (Maybe String)
    eofIfEmpty s
        | save s == emptyIM = return Nothing
        | otherwise = finish s

    viCommandActions :: InputCmd CommandMode (Maybe String)
    viCommandActions = keyChoiceCmd [
                        simpleChar '\n' +> finish
                        , ctrlChar 'd' +> eofIfEmpty
                        , simpleCmdActions >+> viCommandActions
                        , exitingCommands >+> viCommands
                        , repeatedCommands >+> chooseEitherMode
                        ]
        where
            chooseEitherMode :: InputCmd EitherMode (Maybe String)
            chooseEitherMode (Left cm) = viCommandActions cm
            chooseEitherMode (Right im) = viCommands im

    exitingCommands :: InputKeyCmd CommandMode InsertMode
    exitingCommands =  choiceCmd [
                          simpleChar 'i' +> change insertFromCommandMode
                        , simpleChar 'I' +> change (moveToStart . insertFromCommandMode)
                        , simpleKey Home +> change (moveToStart . insertFromCommandMode)
                        , simpleChar 'a' +> change appendFromCommandMode
                        , simpleChar 'A' +> change (moveToEnd . appendFromCommandMode)
                        , simpleKey End +> change (moveToStart  . insertFromCommandMode)
                        , simpleChar 's' +> change (insertFromCommandMode . deleteChar)
                        , simpleChar 'S' +> noArg >|> killAndStoreI killAll
                        , simpleChar 'C' +> noArg >|> killAndStoreI (SimpleMove moveToEnd)
                        ]

    simpleCmdActions :: InputKeyCmd CommandMode CommandMode
    simpleCmdActions = choiceCmd [
                        simpleChar '\ESC' +> change id -- helps break out of loops
                        , simpleChar 'r'   +> replaceOnce
                        , simpleChar 'R'   +> replaceLoop
                        , simpleChar 'D' +> noArg >|> killAndStoreCmd (SimpleMove moveToEnd)
                        , ctrlChar 'l' +> clearScreenCmd
                        , simpleChar 'u' +> commandUndo
                        , ctrlChar 'r' +> commandRedo
                        -- vi-mode quirk: history is put at the start of the line.
                        , simpleChar 'j' +> historyForward >|> change moveToStart
                        , simpleChar 'k' +> historyBack >|> change moveToStart
                        , simpleKey DownKey +> historyForward  >|> change moveToStart
                        , simpleKey UpKey +> historyBack >|> change moveToStart
                        , simpleChar '/' +> viEnterSearch '/' Reverse
                        , simpleChar '?' +> viEnterSearch '?' Forward
                        , simpleChar 'n' +> viSearchHist Reverse []
                        , simpleChar 'N' +> viSearchHist Forward []
                        , simpleKey KillLine +> noArg >|> killAndStoreCmd (SimpleMove moveToStart)
                        ]

    replaceOnce :: InputCmd CommandMode CommandMode
    replaceOnce = try $ changeFromChar replaceChar

    repeatedCommands :: InputKeyCmd CommandMode EitherMode
    repeatedCommands = choiceCmd [argumented, doBefore noArg repeatableCommands]
        where
            start = foreachDigit startArg ['1'..'9']
            addDigit = foreachDigit addNum ['0'..'9']
            argumented = start >+> loop
            loop = keyChoiceCmd [addDigit >+> loop
                                , repeatableCommands
                                -- if no match, bail out.
                                , withoutConsuming (change argState) >+> return . Left
                                ]

    pureMovements :: InputKeyCmd (ArgMode CommandMode) CommandMode
    pureMovements = choiceCmd $ charMovements ++ map mkSimpleCommand movements
        where
            charMovements = [ charMovement 'f' $ \c -> goRightUntil $ overChar (==c)
                            , charMovement 'F' $ \c -> goLeftUntil $ overChar (==c)
                            , charMovement 't' $ \c -> goRightUntil $ beforeChar (==c)
                            , charMovement 'T' $ \c -> goLeftUntil $ afterChar (==c)
                            ]
            mkSimpleCommand (k,move) = k +> change (applyCmdArg move)
            charMovement c move = simpleChar c +> keyChoiceCmd [
                                            useChar (change . applyCmdArg . move)
                                            , withoutConsuming (change argState)
                                            ]

    useMovementsForKill :: Command m s t -> (KillHelper -> Command m s t) -> KeyCommand m s t
    useMovementsForKill alternate useHelper = choiceCmd $
                specialCases
                ++ map (\(k,move) -> k +> useHelper (SimpleMove move)) movements
        where
            specialCases = [ simpleChar 'e' +> useHelper (SimpleMove goToWordDelEnd)
                           , simpleChar 'E' +> useHelper (SimpleMove goToBigWordDelEnd)
                           , simpleChar '%' +> useHelper (GenericKill deleteMatchingBrace)
                           -- Note 't' and 'f' behave differently than in pureMovements.
                           , charMovement 'f' $ \c -> goRightUntil $ afterChar (==c)
                           , charMovement 'F' $ \c -> goLeftUntil $ overChar (==c)
                           , charMovement 't' $ \c -> goRightUntil $ overChar (==c)
                           , charMovement 'T' $ \c -> goLeftUntil $ afterChar (==c)
                           ]
            charMovement c move = simpleChar c +> keyChoiceCmd [
                                        useChar (useHelper . SimpleMove . move)
                                        , withoutConsuming alternate]

    repeatableCommands :: InputKeyCmd (ArgMode CommandMode) EitherMode
    repeatableCommands = choiceCmd
                            [ repeatableCmdToIMode
                            , repeatableCmdMode >+> return . Left
                            , simpleChar '.' +> saveForUndo >|> runLastCommand
                            ]
        where
            runLastCommand s = liftM lastCommand get >>= ($ s)

    repeatableCmdMode :: InputKeyCmd (ArgMode CommandMode) CommandMode
    repeatableCmdMode = choiceCmd
                        [ simpleChar 'x' +> repeatableChange deleteChar
                        , simpleChar 'X' +> repeatableChange (withCommandMode deletePrev)
                        , simpleChar '~' +> repeatableChange (goRight . flipCase)
                        , simpleChar 'p' +> storedCmdAction (pasteCommand pasteGraphemesAfter)
                        , simpleChar 'P' +> storedCmdAction (pasteCommand pasteGraphemesBefore)
                        , simpleChar 'd' +> deletionCmd
                        , simpleChar 'y' +> yankCommand
                        , ctrlChar 'w' +> killAndStoreCmd wordErase
                        , pureMovements
                        ]
        where
            repeatableChange f = storedCmdAction (saveForUndo >|> change (applyArg f))

    flipCase :: CommandMode -> CommandMode
    flipCase CEmpty = CEmpty
    flipCase (CMode xs y zs) = CMode xs (modifyBaseChar flipCaseG y) zs
        where
            flipCaseG c | isLower c = toUpper c
                        | otherwise = toLower c

    repeatableCmdToIMode :: InputKeyCmd (ArgMode CommandMode) EitherMode
    repeatableCmdToIMode = simpleChar 'c' +> deletionToInsertCmd

    deletionCmd :: InputCmd (ArgMode CommandMode) CommandMode
    deletionCmd = keyChoiceCmd
                        [ reinputArg >+> deletionCmd
                        , simpleChar 'd' +> killAndStoreCmd killAll
                        , useMovementsForKill (change argState) killAndStoreCmd
                        , withoutConsuming (change argState)
                        ]

    deletionToInsertCmd :: InputCmd (ArgMode CommandMode) EitherMode
    deletionToInsertCmd = keyChoiceCmd
            [ reinputArg >+> deletionToInsertCmd
            , simpleChar 'c' +> killAndStoreIE killAll
            -- vim, for whatever reason, treats cw same as ce and cW same as cE.
            -- readline does this too, so we should also.
            , simpleChar 'w' +> killAndStoreIE (SimpleMove goToWordDelEnd)
            , simpleChar 'W' +> killAndStoreIE (SimpleMove goToBigWordDelEnd)
            , useMovementsForKill (liftM Left . change argState) killAndStoreIE
            , withoutConsuming (return . Left . argState)
            ]

    yankCommand :: InputCmd (ArgMode CommandMode) CommandMode
    yankCommand = keyChoiceCmd
                    [ reinputArg >+> yankCommand
                    , simpleChar 'y' +> copyAndStore killAll
                    , useMovementsForKill (change argState) copyAndStore
                    , withoutConsuming (change argState)
                    ]
        where
            copyAndStore = storedCmdAction . copyFromArgHelper

    reinputArg :: LineState s => InputKeyCmd (ArgMode s) (ArgMode s)
    reinputArg = foreachDigit restartArg ['1'..'9'] >+> loop
      where
        restartArg n = startArg n . argState
        loop = keyChoiceCmd
                [ foreachDigit addNum ['0'..'9'] >+> loop
                , withoutConsuming return
                ]

    goToWordDelEnd, goToBigWordDelEnd :: InsertMode -> InsertMode
    goToWordDelEnd = goRightUntil $ atStart (not . isWordChar)
                                        .||. atStart (not . isOtherChar)
    goToBigWordDelEnd = goRightUntil $ atStart (not . isBigWordChar)

    movements :: [(Key,InsertMode -> InsertMode)]
    movements = [ (simpleChar 'h', goLeft)
                , (simpleChar 'l', goRight)
                , (simpleChar ' ', goRight)
                , (simpleKey LeftKey, goLeft)
                , (simpleKey RightKey, goRight)
                , (simpleChar '0', moveToStart)
                , (simpleChar '$', moveToEnd)
                , (simpleChar '^', skipRight isSpace . moveToStart)
                , (simpleChar '%', findMatchingBrace)
                ------------------
                -- Word movements
                -- move to the start of the next word
                , (simpleChar 'w', goRightUntil $
                                    atStart isWordChar .||. atStart isOtherChar)
                , (simpleChar 'W', goRightUntil (atStart isBigWordChar))
                -- move to the beginning of the previous word
                , (simpleChar 'b', goLeftUntil $
                                    atStart isWordChar .||. atStart isOtherChar)
                , (simpleChar 'B', goLeftUntil (atStart isBigWordChar))
                -- move to the end of the current word
                , (simpleChar 'e', goRightUntil $
                                    atEnd isWordChar .||. atEnd isOtherChar)
                , (simpleChar 'E', goRightUntil (atEnd isBigWordChar))
                ]

    {-
    From IEEE 1003.1:
    A "bigword" consists of: a maximal sequence of non-blanks preceded and followed by blanks
    A "word" consists of either:
     - a maximal sequence of wordChars, delimited at both ends by non-wordchars
     - a maximal sequence of non-blank non-wordchars, delimited at both ends by either blanks
       or a wordchar.
    -}
    isBigWordChar, isWordChar, isOtherChar :: Char -> Bool
    isBigWordChar = not . isSpace
    isWordChar = isAlphaNum .||. (=='_')
    isOtherChar = not . (isSpace .||. isWordChar)

    (.||.) :: (a -> Bool) -> (a -> Bool) -> a -> Bool
    (f .||. g) x = f x || g x

    foreachDigit :: (Monad m, LineState t) => (Int -> s -> t) -> [Char]
                    -> KeyCommand m s t
    foreachDigit f ds = choiceCmd $ map digitCmd ds
        where digitCmd d = simpleChar d +> change (f (toDigit d))
              toDigit d = fromEnum d - fromEnum '0'

    -- This mimics the ctrl-w command in readline's vi mode, which corresponds to
    -- the tty's werase character.
    wordErase :: KillHelper
    wordErase = SimpleMove $ goLeftUntil $ atStart isBigWordChar

    ------------------
    -- Matching braces

    findMatchingBrace :: InsertMode -> InsertMode
    findMatchingBrace (IMode xs (y:ys))
        | Just b <- matchingRightBrace yc,
          Just ((b':bs),ys') <- scanBraces yc b ys = IMode (bs++[y]++xs) (b':ys')
        | Just b <- matchingLeftBrace yc,
          Just (bs,xs') <- scanBraces yc b xs = IMode xs' (bs ++ [y]++ys)
      where yc = baseChar y
    findMatchingBrace im = im

    deleteMatchingBrace :: InsertMode -> ([Grapheme],InsertMode)
    deleteMatchingBrace (IMode xs (y:ys))
        | Just b <- matchingRightBrace yc,
          Just (bs,ys') <- scanBraces yc b ys = (y : reverse bs, IMode xs ys')
        | Just b <- matchingLeftBrace yc,
          Just (bs,xs') <- scanBraces yc b xs = (bs ++ [y], IMode xs' ys)
      where yc = baseChar y
    deleteMatchingBrace im = ([],im)

    scanBraces :: Char -> Char -> [Grapheme] -> Maybe ([Grapheme],[Grapheme])
    scanBraces c d = scanBraces' (1::Int) []
        where
            scanBraces' 0 bs xs = Just (bs,xs)
            scanBraces' _ _ [] = Nothing
            scanBraces' n bs (x:xs) = scanBraces' m (x:bs) xs
                where m | baseChar x == c = n+1
                        | baseChar x == d = n-1
                        | otherwise = n

    matchingRightBrace, matchingLeftBrace :: Char -> Maybe Char
    matchingRightBrace = flip lookup braceList
    matchingLeftBrace = flip lookup (map (\(c,d) -> (d,c)) braceList)

    braceList :: [(Char,Char)]
    braceList = [('(',')'), ('[',']'), ('{','}')]

    ---------------
    -- Replace mode
    replaceLoop :: InputCmd CommandMode CommandMode
    replaceLoop = saveForUndo >|> change insertFromCommandMode >|> loop
                    >|> change enterCommandModeRight
        where
            loop = try (oneReplaceCmd >+> loop)
            oneReplaceCmd = choiceCmd [
                    simpleKey LeftKey +> change goLeft
                    , simpleKey RightKey +> change goRight
                    , changeFromChar replaceCharIM
                    ]

    ---------------------------
    -- Saving previous commands

    storeLastCmd :: Monad m => SavedCommand m -> Command (ViT m) s s
    storeLastCmd act = \s -> do
            modify $ \vs -> vs {lastCommand = act}
            return s

    storedAction :: Monad m => SavedCommand m -> SavedCommand m
    storedAction act = storeLastCmd act >|> act

    storedCmdAction :: Monad m => Command (ViT m) (ArgMode CommandMode) CommandMode
                                -> Command (ViT m) (ArgMode CommandMode) CommandMode
    storedCmdAction act = storeLastCmd (liftM Left . act) >|> act

    storedIAction :: Monad m => Command (ViT m) (ArgMode CommandMode) InsertMode
                            -> Command (ViT m) (ArgMode CommandMode) InsertMode
    storedIAction act = storeLastCmd (liftM Right . act) >|> act

    killAndStoreCmd :: MonadIO m => KillHelper -> Command (ViT m) (ArgMode CommandMode) CommandMode
    killAndStoreCmd = storedCmdAction . killFromArgHelper

    killAndStoreI :: MonadIO m => KillHelper -> Command (ViT m) (ArgMode CommandMode) InsertMode
    killAndStoreI = storedIAction . killFromArgHelper

    killAndStoreIE :: MonadIO m => KillHelper -> Command (ViT m) (ArgMode CommandMode) EitherMode
    killAndStoreIE helper = storedAction (killFromArgHelper helper >|> return . Right)

    noArg :: Monad m => Command m s (ArgMode s)
    noArg = return . startArg 1

    -------------------
    -- Vi-style searching

    data SearchEntry = SearchEntry {
                        entryState :: InsertMode,
                        searchChar :: Char
                        }

    searchText :: SearchEntry -> [Grapheme]
    searchText SearchEntry {entryState = IMode xs ys} = reverse xs ++ ys

    instance LineState SearchEntry where
        beforeCursor prefix se = beforeCursor (prefix ++ stringToGraphemes [searchChar se])
                                    (entryState se)
        afterCursor = afterCursor . entryState

    viEnterSearch :: Monad m => Char -> Direction
                        -> Command (ViT m) CommandMode CommandMode
    viEnterSearch c dir s = setState (SearchEntry emptyIM c) >>= loopEntry
        where
            modifySE f se = se {entryState = f (entryState se)}
            loopEntry = keyChoiceCmd [
                            editEntry >+> loopEntry
                            , simpleChar '\n' +> \se ->
                                viSearchHist dir (searchText se) s
                            , withoutConsuming (change (const s))
                            ]
            editEntry = choiceCmd [
                            useChar (change . modifySE . insertChar)
                            , simpleKey LeftKey +> change (modifySE goLeft)
                            , simpleKey RightKey +> change (modifySE goRight)
                            , simpleKey Backspace +> change (modifySE deletePrev)
                            , simpleKey Delete +> change (modifySE deleteNext)
                            ]

    viSearchHist :: forall m . Monad m
        => Direction -> [Grapheme] -> Command (ViT m) CommandMode CommandMode
    viSearchHist dir toSearch cm = do
        vstate :: ViState m <- get
        let toSearch' = if null toSearch
                            then lastSearch vstate
                            else toSearch
        result <- doSearch False SearchMode {
                                        searchTerm = toSearch',
                                        foundHistory = save cm, -- TODO: not needed
                                        direction = dir}
        case result of
            Left e -> effect e >> setState cm
            Right sm -> do
                put vstate {lastSearch = toSearch'}
                setState (restore (foundHistory sm))