ghci: Add support for prompt functions

This is an updated version of @jlengyel's original patch adding support
for prompt functions.

2 files changed, 73 insertions, 50 deletions

diff --git a/ghc/GhciMonad.hs b/ghc/GhciMonad.hs
index 6d068be485..c09b61d153 100644
--- a/ghc/GhciMonad.hs
+++ b/ghc/GhciMonad.hs
@@ -15,6 +15,7 @@ module GhciMonad (
         GHCiState(..), setGHCiState, getGHCiState, modifyGHCiState,
         GHCiOption(..), isOptionSet, setOption, unsetOption,
         Command,
+        PromptFunction,
         BreakLocation(..),
         TickArray,
         getDynFlags,
@@ -66,15 +67,22 @@ import Control.Applicative (Applicative(..))
 -----------------------------------------------------------------------------
 -- GHCi monad
 
--- the Bool means: True = we should exit GHCi (:quit)
+-- | A GHCi command
+--
+-- the @Bool@ means: @True@ = we should exit GHCi (@:quit@)
 type Command = (String, String -> InputT GHCi Bool, CompletionFunc GHCi)
 
+-- | A function to generate the GHCi prompt.
+type PromptFunction = [String]  -- ^ names of modules in scope
+                   -> Int       -- ^ current line number
+                   -> IO String -- ^ an action returning a prompt string
+
 data GHCiState = GHCiState
      {
         progname       :: String,
         args           :: [String],
-        prompt         :: String,
-        prompt2        :: String,
+        prompt         :: PromptFunction,
+        prompt2        :: PromptFunction,
         editor         :: String,
         stop           :: String,
         options        :: [GHCiOption],
diff --git a/ghc/InteractiveUI.hs b/ghc/InteractiveUI.hs
index e5c4e11dea..026d6ea681 100644
--- a/ghc/InteractiveUI.hs
+++ b/ghc/InteractiveUI.hs
@@ -116,9 +116,7 @@ import GHC.TopHandler ( topHandler )
 data GhciSettings = GhciSettings {
         availableCommands :: [Command],
         shortHelpText     :: String,
-        fullHelpText      :: String,
-        defPrompt         :: String,
-        defPrompt2        :: String
+        fullHelpText      :: String
     }
 
 defaultGhciSettings :: GhciSettings
@@ -126,9 +124,7 @@ defaultGhciSettings =
     GhciSettings {
         availableCommands = ghciCommands,
         shortHelpText     = defShortHelpText,
-        fullHelpText      = defFullHelpText,
-        defPrompt         = default_prompt,
-        defPrompt2        = default_prompt2
+        fullHelpText      = defFullHelpText
     }
 
 ghciWelcomeMsg :: String
@@ -302,7 +298,13 @@ defFullHelpText =
   "   :set args <arg> ...         set the arguments returned by System.getArgs\n" ++
   "   :set prog <progname>        set the value returned by System.getProgName\n" ++
   "   :set prompt <prompt>        set the prompt used in GHCi\n" ++
+  "   :set prompt-function <expr> set the function used to create the GHCi prompt\n" ++
+  "                                 of type [String] -> Int -> IO String\n" ++
+  "                                 which will be passed the current list of\n" ++
+  "                                 imported modules and the current line number\n" ++
   "   :set prompt2 <prompt>       set the continuation prompt used in GHCi\n" ++
+  "   :set prompt2-function       set the function used to create the GHCi\n" ++
+  "     <expr>                       continuation prompt. See :set prompt-function\n" ++
   "   :set editor <cmd>           set the command used for :edit\n" ++
   "   :set stop [<n>] <cmd>       set the command to run when a breakpoint is hit\n" ++
   "   :unset <option> ...         unset options\n" ++
@@ -345,10 +347,8 @@ findEditor = do
         return ""
 #endif
 
-default_progname, default_prompt, default_prompt2, default_stop :: String
+default_progname, default_stop :: String
 default_progname = "<interactive>"
-default_prompt = "%s> "
-default_prompt2 = "%s| "
 default_stop = ""
 
 default_args :: [String]
@@ -409,9 +409,11 @@ interactiveUI config srcs maybe_exprs = do
    startGHCi (runGHCi srcs maybe_exprs)
         GHCiState{ progname           = default_progname,
                    GhciMonad.args     = default_args,
-                   prompt             = defPrompt config,
-                   prompt2            = defPrompt2 config,
                    stop               = default_stop,
+                   prompt             = (\xs _ -> return $
+                                          intercalate " " xs ++ "> "),
+                   prompt2            = (\xs _ -> return $
+                                          intercalate " " xs ++ "| "),
                    editor             = default_editor,
                    options            = [],
                    -- We initialize line number as 0, not 1, because we use
@@ -656,6 +658,7 @@ mkPrompt = do
   st <- getGHCiState
   imports <- GHC.getContext
   resumes <- GHC.getResumeContext
+  dflags <- getDynFlags
 
   context_bit <-
         case resumes of
@@ -674,25 +677,28 @@ mkPrompt = do
              | otherwise = empty
 
         rev_imports = reverse imports -- rightmost are the most recent
-        modules_bit =
-             hsep [ char '*' <> ppr m | IIModule m <- rev_imports ] <+>
-             hsep (map ppr [ myIdeclName d | IIDecl d <- rev_imports ])
+        module_list = [char '*' <> ppr m | IIModule m <- rev_imports] ++
+                      map ppr [myIdeclName d | IIDecl d <- rev_imports]
+        module_string_list = map (showSDoc dflags) module_list
+        deflt_prompt = dots <> context_bit <> hsep module_list
 
          --  use the 'as' name if there is one
         myIdeclName d | Just m <- ideclAs d = m
                       | otherwise           = unLoc (ideclName d)
 
-        deflt_prompt = dots <> context_bit <> modules_bit
+        line_no = 1 + line_number st
+
+  promptString <- liftIO $ (prompt st) module_string_list line_no
 
-        f ('%':'l':xs) = ppr (1 + line_number st) <> f xs
+  let   f ('%':'l':xs) = ppr (1 + line_number st) <> f xs
         f ('%':'s':xs) = deflt_prompt <> f xs
         f ('%':'%':xs) = char '%' <> f xs
         f (x:xs) = char x <> f xs
         f [] = empty
 
-  dflags <- getDynFlags
-  return (showSDoc dflags (f (prompt st)))
+        promptDoc = dots <> context_bit <> (f promptString)
 
+  return (showSDoc dflags promptDoc)
 
 queryQueue :: GHCi (Maybe String)
 queryQueue = do
@@ -2055,14 +2061,30 @@ setCmd str
         case toArgs rest of
             Right [prog] -> setProg prog
             _ -> liftIO (hPutStrLn stderr "syntax: :set prog <progname>")
-    Right ("prompt",  rest) -> setPrompt  $ dropWhile isSpace rest
-    Right ("prompt2", rest) -> setPrompt2 $ dropWhile isSpace rest
+    Right ("prompt-function",  rest) ->
+        setPromptFunc setPrompt $ dropWhile isSpace rest
+    Right ("prompt",          rest) ->
+        setPromptString setPrompt (dropWhile isSpace rest) "syntax: :set prompt <string>"
+    Right ("prompt2-function", rest) ->
+        setPromptFunc setPrompt2 $ dropWhile isSpace rest
+    Right ("prompt2",         rest) ->
+        setPromptString setPrompt2 (dropWhile isSpace rest) "syntax: :set prompt2 <string>"
     Right ("editor",  rest) -> setEditor  $ dropWhile isSpace rest
     Right ("stop",    rest) -> setStop    $ dropWhile isSpace rest
     _ -> case toArgs str of
          Left err -> liftIO (hPutStrLn stderr err)
          Right wds -> setOptions wds
 
+setPromptFunc :: (PromptFunction -> GHCi ()) -> String -> GHCi ()
+setPromptFunc f s = do
+    -- We explicitly annotate the type of the expression to ensure
+    -- that unsafeCoerce# is passed the exact type necessary rather
+    -- than a more general one
+    let exprStr = "(" ++ s ++ ") :: [String] -> Int -> IO String"
+    (HValue funValue) <- GHC.compileExpr exprStr
+    f (unsafeCoerce# funValue)
+
+
 setiCmd :: String -> GHCi ()
 setiCmd ""   = GHC.getInteractiveDynFlags >>= liftIO . showDynFlags False
 setiCmd "-a" = GHC.getInteractiveDynFlags >>= liftIO . showDynFlags True
@@ -2155,30 +2177,23 @@ setStop cmd = do
   st <- getGHCiState
   setGHCiState st{ stop = cmd }
 
-setPrompt :: String -> GHCi ()
-setPrompt = setPrompt_ f err
-  where
-    f v st = st { prompt = v }
-    err st = "syntax: :set prompt <prompt>, currently \"" ++ prompt st ++ "\""
+setPrompt :: PromptFunction -> GHCi ()
+setPrompt v = modifyGHCiState (\st -> st { prompt = v})
 
-setPrompt2 :: String -> GHCi ()
-setPrompt2 = setPrompt_ f err
-  where
-    f v st = st { prompt2 = v }
-    err st = "syntax: :set prompt2 <prompt>, currently \"" ++ prompt2 st ++ "\""
+setPrompt2 :: PromptFunction -> GHCi ()
+setPrompt2 v = modifyGHCiState (\st -> st {prompt2 = v})
 
-setPrompt_ :: (String -> GHCiState -> GHCiState) -> (GHCiState -> String) -> String -> GHCi ()
-setPrompt_ f err value = do
-  st <- getGHCiState
-  if null value
-      then liftIO $ hPutStrLn stderr $ err st
-      else case value of
-           '\"' : _ -> case reads value of
-                       [(value', xs)] | all isSpace xs ->
-                           setGHCiState $ f value' st
-                       _ ->
-                           liftIO $ hPutStrLn stderr "Can't parse prompt string. Use Haskell syntax."
-           _ -> setGHCiState $ f value st
+setPromptString :: (PromptFunction -> GHCi ()) -> String -> String -> GHCi ()
+setPromptString f value err = do
+    if null value
+        then liftIO $ hPutStrLn stderr $ err
+        else case value of
+            '\"' : _ -> case reads value of
+                        [(value', xs)] | all isSpace xs ->
+                            f (\_ _ -> return value')
+                        _ ->
+                          liftIO $ hPutStrLn stderr "Can't parse prompt string. Use Haskell syntax."
+            _ -> f (\_ _ -> return value)
 
 setOptions wds =
    do -- first, deal with the GHCi opts (+s, +t, etc.)
@@ -2255,8 +2270,10 @@ unsetOptions str
          defaulters =
            [ ("args"   , setArgs default_args)
            , ("prog"   , setProg default_progname)
-           , ("prompt" , setPrompt default_prompt)
-           , ("prompt2", setPrompt2 default_prompt2)
+           , ("prompt" , setPrompt (\xs _ -> return $
+                           intercalate " " xs ++ "> "))
+           , ("prompt2", setPrompt2 (\xs _ -> return $
+                           intercalate " " xs ++ "| "))
            , ("editor" , liftIO findEditor >>= setEditor)
            , ("stop"   , setStop default_stop)
            ]
@@ -2320,8 +2337,6 @@ showCmd str = do
   case words str of
         ["args"]     -> liftIO $ putStrLn (show (GhciMonad.args st))
         ["prog"]     -> liftIO $ putStrLn (show (progname st))
-        ["prompt"]   -> liftIO $ putStrLn (show (prompt st))
-        ["prompt2"]  -> liftIO $ putStrLn (show (prompt2 st))
         ["editor"]   -> liftIO $ putStrLn (show (editor st))
         ["stop"]     -> liftIO $ putStrLn (show (stop st))
         ["imports"]  -> showImports
@@ -2337,7 +2352,7 @@ showCmd str = do
         ["languages"] -> showLanguages -- backwards compat
         ["language"]  -> showLanguages
         ["lang"]      -> showLanguages -- useful abbreviation
-        _ -> throwGhcException (CmdLineError ("syntax:  :show [ args | prog | prompt | prompt2 | editor | stop | modules\n" ++
+        _ -> throwGhcException (CmdLineError ("syntax:  :show [ args | prog | editor | stop | modules\n" ++
                                               "               | bindings | breaks | context | packages | language ]"))
 
 showiCmd :: String -> GHCi ()