a ghci-like repl with State effects in addition to IO effects

StateRepl.hs

-- In response to https://www.reddit.com/r/haskell/comments/lu14nu/is_it_possible_to_get_a_repl_in_a_state_monad/
--
-- The challenge is to write a repl in which we can run State effects in
-- addition to ghci's basic features of binding variables, inspecting values,
-- and running IO effects.
--
-- Well, the original request was whether it was possible to somehow add those
-- features to ghci. But being the maintainer of the hint library, I thought I
-- should demonstrate how easy it is to write your own ghci alternative which
-- has those features built-in!
--
-- Note that ghci has a lot of features, whereas this repl only has the
-- features I bothered to implement:
--
-- * executing State actions and IO actions
-- * evaluating expressions
-- * printing the result of such actions and expressions
-- * binding the result of such actions and expressions to variables
--
-- Here's what interacting with the repl implemented in this file looks like:
--
--     state = 0
--     > modify (+2)
--     state = 2
--     > get
--     2
--     state = 2
--     > x <- get
--     state = 2
--     > print x
--     2
--     state = 2
--     > printt x
--     <interactive>:2:7: error:
--         • Variable not in scope: printt :: Int -> a0
--         • Perhaps you meant ‘print’ (imported from Prelude)
--     state = 2

{-# LANGUAGE ScopedTypeVariables, ViewPatterns #-}
module Main where

import Control.Exception (SomeException(..), throwIO)
import Control.Monad.Catch (try)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.State (State, StateT)
import Data.Foldable (for_)
import Data.Function (fix)
import Data.Functor.Identity (runIdentity)
import Data.IORef (IORef, writeIORef)
import Data.Typeable (Typeable, cast, typeOf)
import Language.Haskell.Interpreter (InterpreterT)
import System.Console.Haskeline (InputT)
import System.IO (hPutStrLn, stderr)
import qualified Control.Monad.Trans.State as State
import qualified Language.Haskell.Interpreter as Hint
import qualified System.Console.Haskeline as Haskeline


-- The trick is to run our repl in a monad stack which contains both StateT (so
-- we can execute State effects) and InterpreterT (so we can interpret code
-- typed by the user). That's pretty much it, really.
type M = StateT Int (InputT (InterpreterT IO))

runM :: Int -> M a -> IO (Either Hint.InterpreterError (a, Int))
runM s
  = Hint.runInterpreter
  . Haskeline.runInputT Haskeline.defaultSettings
  . flip State.runStateT s

liftState
  :: State Int a
  -> M a
liftState = State.mapStateT (pure . runIdentity)

liftHaskeline
  :: InputT (InterpreterT IO) a
  -> M a
liftHaskeline = lift

liftHint
  :: InterpreterT IO a
  -> M a
liftHint = lift . lift


-- As promised, we can now evaluate State actions typed by the user, e.g.
-- "modify (+2)". Our Hint effect allows us to interpret the string as a type of
-- our choice, so we choose the type @State Int ()@, and we run that State
-- action in our monad stack.

evalStateAction_ :: String -> M ()
evalStateAction_ s = do
  action <- liftHint
          $ Hint.interpret
              s
              (Hint.as :: State Int ())
  liftState action

evalIOAction_ :: String -> M ()
evalIOAction_ s = do
  action <- liftHint
          $ Hint.interpret
              s
              (Hint.as :: IO ())
  liftIO action


-- Evaluating a State action like "get" and printing its result is a bit
-- trickier, because we don't know the type of that result. We thus cannot
-- simply interpret the string as a value of type @State Int Int@. The solution
-- is to tweak the string we interpret so that whatever the type of the original
-- expression, the type of the tweaked expression is @State Int String@.

evalStateAction :: String -> M ()
evalStateAction s = do
  action <- liftHint $ Hint.interpret
    ("fmap show (" ++ s ++ ")")
    (Hint.as :: State Int String)
  shown <- liftState action
  liftIO $ putStrLn shown

evalIOAction :: String -> M ()
evalIOAction s = do
  action <- liftHint $ Hint.interpret
    ("fmap show (" ++ s ++ ")")
    (Hint.as :: IO String)
  shown <- liftIO action
  liftIO $ putStrLn shown

evalValue :: String -> M ()
evalValue s = do
  shown <- liftHint $ Hint.interpret
    ("show (" ++ s ++ ")")
    (Hint.as :: String)
  liftIO $ putStrLn shown


-- For binding the result of an IO action or a pure expression to a variable,
-- like in "x <- pure 42" and "let x = 42", Hint provides 'runStmt'. It works
-- just like ghci.

evalIOBind :: String -> String -> M ()
evalIOBind var expr = do
  liftHint $ Hint.runStmt
    (var ++ " <- " ++ expr)

evalLet :: String -> String -> M ()
evalLet var expr = do
  liftHint $ Hint.runStmt
    ("let " ++ var ++ " = " ++ expr)


-- Binding the result of a State action to a variable, as in "x <- get", is
-- another story. Since 'runStmt' can only bind the result of IO actions and
-- pure expressions, we need to jump through hoops and tweak the string we
-- interpret quite a lot in order to turn it into a supported type.

evalStateBind :: String -> String -> M ()
evalStateBind var expr = do
  tmpState <- liftState State.get
  bind "tmpState" tmpState
  liftHint $ Hint.runStmt
    ("let (" ++ var ++ ", tmpState') = runState (" ++ expr ++ ") tmpState")
  tmpState' <- liftHint $ Hint.interpret
    "tmpState'"
    (Hint.as :: Int)
  liftState $ State.put tmpState'

-- Bind a variable inside the interpreter to a value from outside the
-- interpreter.
--
-- If 'value' has a Show instance, we can simply use
--
--     runStmt ("let " ++ var ++ " = " ++ show value)
--
-- By using the following IORef trick, we can support any value with a Typeable
-- instance! This seems pretty useful, I should add it to the hint library.
bind :: forall a. Typeable a => String -> a -> M ()
bind var value = do
  let t = typeOf (undefined :: a)
  liftHint $ Hint.runStmt
    ("tmpIORef <- newIORef (undefined :: " ++ show t ++ ")")
  tmpIORef <- liftHint $ Hint.interpret
    "tmpIORef"
    (Hint.as :: IORef a)
  liftIO $ writeIORef tmpIORef value
  liftHint $ Hint.runStmt
    (var ++ " <- readIORef tmpIORef")


-- When a user types an expression which doesn't typecheck or which throws an
-- exception, we don't want to abort the program, we just want to print the
-- error message. Each error message is a list of strings because the message
-- can span multiple lines.

type Error = [String]

printError :: Error -> M ()
printError errorLines = do
  for_ errorLines $ \errorLine -> do
    liftIO $ hPutStrLn stderr errorLine

catchError :: M a -> M (Either Error a)
catchError body = do
  r <- try body
  case r of
    Left (SomeException (cast -> Just (Hint.WontCompile ghcErrors))) -> do
      -- print ghc's multi-line error messages on multiple lines
      pure $ Left $ fmap Hint.errMsg $ ghcErrors
    Left e -> do
      -- for other errors, call Show and hope for the best
      pure $ Left $ lines $ show e
    Right a -> do
      pure $ Right a


-- When the user types "x <- expr", we don't know whether they intend "expr" to
-- be a State action or an IO action. So we simply try both and see which one
-- typechecks.

eval :: String -> M (Either Error ())
eval (words -> "let" : var : "=" : (unwords -> expr)) = do
  catchError $ evalLet var expr
eval (words -> var : "=" : (unwords -> expr)) = do
  catchError $ evalLet var expr
eval (words -> var : "<-" : (unwords -> expr)) = do
  keepFirstSuccess
    [ evalStateBind var expr
    , evalIOBind var expr
    ] 
eval s = do
  keepFirstSuccess
    [ evalStateAction_ s
    , evalStateAction s
    , evalIOAction_ s
    , evalIOAction s
    , evalValue s
    ] 

-- Similar to 'asum', but keep the last error instead of concatenating them. The
-- input list must be non-empty.
keepFirstSuccess :: [M a] -> M (Either Error a)
keepFirstSuccess [] = do
  error "keepFirstSuccess: empty list"
keepFirstSuccess [action] = do
  catchError action
keepFirstSuccess (action:actions) = do
  r <- catchError action
  case r of
    Left _ -> do
      keepFirstSuccess actions
    Right a -> do
      pure $ Right a
  

-- Finally, here's the repl loop: we initialize the monad stack, then repeatedly
-- read and evaluate lines of input.

main :: IO ()
main = do
  r <- runM 0 $ do
    liftHint $ Hint.setImports
      [ "Prelude"
      , "Control.Monad.Trans.State"
      , "Data.Functor.Identity"
      , "Data.IORef"
      ]
    fix $ \loop -> do
      state <- liftState State.get
      liftIO $ putStrLn $ "state = " ++ show state
      maybeS <- liftHaskeline $ Haskeline.getInputLine "> "
      case maybeS of
        Nothing -> do
          -- exit on Ctrl-D (and on Ctrl-C too via an exception)
          pure ()
        Just s -> do
          r <- eval s
          case r of
            Left err -> do
              printError err
            Right () -> do
              pure ()
          loop
  case r of
    Left err -> do
      throwIO err
    Right _ -> do
      pure ()