Introduce ghci command wrapper

Introduce ghci command wrapper, which can be used to cutomize ghci:
* process additionals actions before/after the command
* handle particular exceptions in given ways
* logging stats

We also split the timing and printing part of `timeIt` into different
functions.

2 files changed, 76 insertions, 33 deletions

diff --git a/ghc/GHCi/UI.hs b/ghc/GHCi/UI.hs
index 0c09844c1b..10ca51124e 100644
--- a/ghc/GHCi/UI.hs
+++ b/ghc/GHCi/UI.hs
@@ -488,6 +488,7 @@ interactiveUI config srcs maybe_exprs = do
                    ghci_commands      = availableCommands config,
                    ghci_macros        = [],
                    last_command       = Nothing,
+                   cmd_wrapper        = (cmdSuccess =<<),
                    cmdqueue           = [],
                    remembered_ctx     = [],
                    transient_ctx      = [],
@@ -973,9 +974,11 @@ runOneCommand eh gCmd = do
   mb_cmd1 <- maybe (noSpace gCmd) (return . Just) mb_cmd0
   case mb_cmd1 of
     Nothing -> return Nothing
-    Just c  -> ghciHandle (\e -> lift $ eh e >>= return . Just) $
-             handleSourceError printErrorAndFail
-               (doCommand c)
+    Just c  -> do
+      st <- getGHCiState
+      ghciHandle (\e -> lift $ eh e >>= return . Just) $
+        handleSourceError printErrorAndFail $
+          cmd_wrapper st $ doCommand c
                -- source error's are handled by runStmt
                -- is the handler necessary here?
   where
@@ -1014,14 +1017,14 @@ runOneCommand eh gCmd = do
     collectError = userError "unterminated multiline command :{ .. :}"
 
     -- | Handle a line of input
-    doCommand :: String -> InputT GHCi (Maybe Bool)
+    doCommand :: String -> InputT GHCi CommandResult
 
     -- command
-    doCommand stmt | (':' : cmd) <- removeSpaces stmt = do
-      result <- specialCommand cmd
-      case result of
-        True -> return Nothing
-        _    -> return $ Just True
+    doCommand stmt | stmt'@(':' : cmd) <- removeSpaces stmt = do
+      (stats, result) <- runWithStats (const Nothing) $ specialCommand cmd
+      let processResult True = Nothing
+          processResult False = Just True
+      return $ CommandComplete stmt' (processResult <$> result) stats
 
     -- haskell
     doCommand stmt = do
@@ -1033,12 +1036,13 @@ runOneCommand eh gCmd = do
           fst_line_num <- line_number <$> getGHCiState
           mb_stmt <- checkInputForLayout stmt gCmd
           case mb_stmt of
-            Nothing      -> return $ Just True
+            Nothing -> return CommandIncomplete
             Just ml_stmt -> do
               -- temporarily compensate line-number for multi-line input
-              result <- timeIt runAllocs $ lift $
+              (stats, result) <- runAndPrintStats runAllocs $ lift $
                 runStmtWithLineNum fst_line_num ml_stmt GHC.RunToCompletion
-              return $ Just (runSuccess result)
+              return $
+                CommandComplete ml_stmt (Just . runSuccess <$> result) stats
         else do -- single line input and :{ - multiline input
           last_line_num <- line_number <$> getGHCiState
           -- reconstruct first line num from last line num and stmt
@@ -1047,9 +1051,9 @@ runOneCommand eh gCmd = do
               stmt_nl_cnt2 = length [ () | '\n' <- stmt' ]
               stmt' = dropLeadingWhiteLines stmt -- runStmt doesn't like leading empty lines
           -- temporarily compensate line-number for multi-line input
-          result <- timeIt runAllocs $ lift $
+          (stats, result) <- runAndPrintStats runAllocs $ lift $
             runStmtWithLineNum fst_line_num stmt' GHC.RunToCompletion
-          return $ Just (runSuccess result)
+          return $ CommandComplete stmt' (Just . runSuccess <$> result) stats
 
     -- runStmt wrapper for temporarily overridden line-number
     runStmtWithLineNum :: Int -> String -> SingleStep
@@ -1745,7 +1749,9 @@ wrapDeferTypeErrors load =
     (\_ -> load)
 
 loadModule :: [(FilePath, Maybe Phase)] -> InputT GHCi SuccessFlag
-loadModule fs = timeIt (const Nothing) (loadModule' fs)
+loadModule fs = do
+  (_, result) <- runAndPrintStats (const Nothing) (loadModule' fs)
+  either (liftIO . Exception.throwIO) return result
 
 -- | @:load@ command
 loadModule_ :: [FilePath] -> InputT GHCi ()
diff --git a/ghc/GHCi/UI/Monad.hs b/ghc/GHCi/UI/Monad.hs
index cbf527e623..8f60dfbb7e 100644
--- a/ghc/GHCi/UI/Monad.hs
+++ b/ghc/GHCi/UI/Monad.hs
@@ -14,13 +14,14 @@ module GHCi.UI.Monad (
         GHCi(..), startGHCi,
         GHCiState(..), setGHCiState, getGHCiState, modifyGHCiState,
         GHCiOption(..), isOptionSet, setOption, unsetOption,
-        Command(..),
+        Command(..), CommandResult(..), cmdSuccess,
         PromptFunction,
         BreakLocation(..),
         TickArray,
         getDynFlags,
 
-        runStmt, runDecls, runDecls', resume, timeIt, recordBreak, revertCAFs,
+        runStmt, runDecls, runDecls', resume, recordBreak, revertCAFs,
+        ActionStats(..), runAndPrintStats, runWithStats, printStats,
 
         printForUserNeverQualify, printForUserModInfo,
         printForUser, printForUserPartWay, prettyLocations,
@@ -93,6 +94,10 @@ data GHCiState = GHCiState
         last_command   :: Maybe Command,
             -- ^ @:@ at the GHCi prompt repeats the last command, so we
             -- remember it here
+        cmd_wrapper    :: InputT GHCi CommandResult -> InputT GHCi (Maybe Bool),
+            -- ^ The command wrapper is run for each command or statement.
+            -- The 'Bool' value denotes whether the command is successful and
+            -- 'Nothing' means to exit GHCi.
         cmdqueue       :: [String],
 
         remembered_ctx :: [InteractiveImport],
@@ -164,6 +169,21 @@ data Command
      -- ^ 'CompletionFunc' for arguments
    }
 
+data CommandResult
+   = CommandComplete
+   { cmdInput :: String
+   , cmdResult :: Either SomeException (Maybe Bool)
+   , cmdStats :: ActionStats
+   }
+   | CommandIncomplete
+     -- ^ Unterminated multiline command
+   deriving Show
+
+cmdSuccess :: Haskeline.MonadException m => CommandResult -> m (Maybe Bool)
+cmdSuccess CommandComplete{ cmdResult = Left e } = liftIO $ throwIO e
+cmdSuccess CommandComplete{ cmdResult = Right r } = return r
+cmdSuccess CommandIncomplete = return $ Just True
+
 type PromptFunction = [String]
                    -> Int
                    -> GHCi SDoc
@@ -386,22 +406,39 @@ resume canLogSpan step = do
 -- --------------------------------------------------------------------------
 -- timing & statistics
 
-timeIt :: (a -> Maybe Integer) -> InputT GHCi a -> InputT GHCi a
-timeIt getAllocs action
-  = do b <- lift $ isOptionSet ShowTiming
-       if not b
-          then action
-          else do time1   <- liftIO $ getCurrentTime
-                  a <- action
-                  let allocs = getAllocs a
-                  time2   <- liftIO $ getCurrentTime
-                  dflags  <- getDynFlags
-                  let period = time2 `diffUTCTime` time1
-                  liftIO $ printTimes dflags allocs (realToFrac period)
-                  return a
-
-printTimes :: DynFlags -> Maybe Integer -> Double -> IO ()
-printTimes dflags mallocs secs
+data ActionStats = ActionStats
+  { actionAllocs :: Maybe Integer
+  , actionElapsedTime :: Double
+  } deriving Show
+
+runAndPrintStats
+  :: (a -> Maybe Integer)
+  -> InputT GHCi a
+  -> InputT GHCi (ActionStats, Either SomeException a)
+runAndPrintStats getAllocs action = do
+  result <- runWithStats getAllocs action
+  case result of
+    (stats, Right{}) -> do
+      showTiming <- lift $ isOptionSet ShowTiming
+      when showTiming $ do
+        dflags  <- getDynFlags
+        liftIO $ printStats dflags stats
+    _ -> return ()
+  return result
+
+runWithStats
+  :: ExceptionMonad m
+  => (a -> Maybe Integer) -> m a -> m (ActionStats, Either SomeException a)
+runWithStats getAllocs action = do
+  t0 <- liftIO getCurrentTime
+  result <- gtry action
+  let allocs = either (const Nothing) getAllocs result
+  t1 <- liftIO getCurrentTime
+  let elapsedTime = realToFrac $ t1 `diffUTCTime` t0
+  return (ActionStats allocs elapsedTime, result)
+
+printStats :: DynFlags -> ActionStats -> IO ()
+printStats dflags ActionStats{actionAllocs = mallocs, actionElapsedTime = secs}
    = do let secs_str = showFFloat (Just 2) secs
         putStrLn (showSDoc dflags (
                  parens (text (secs_str "") <+> text "secs" <> comma <+>