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' (mimicking 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))