MonadState Example

monadstate-example.hs

-- monadstate-example.hs
--
-- Load this program in GHCi:
--
--     stack repl \
--         --resolver nightly \
--         --package transformers \
--         --package mtl \
--         monadstate-example.hs
--
-- Then try `test` and `main`.
--
--     GHCi> test
--     ...
--     GHCi> main
--     ...
--
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE MultiWayIf #-}

import Control.Monad.Reader
import Control.Monad.State
import Data.IORef
import Text.Printf


-- Imagine that this is the entry point of a large program.
-- Notice that `m` is abstract. This program works with any `m`
-- as long as `MonadState Int m` is implemented.
program :: MonadState Int m => m ()
program = do
  n <- get
  if | n == 10   -> return ()
     | even n    -> put (n `div` 2 + 1) >> program
     | otherwise -> put (3*n + 1)       >> program


-- In our test, we'll go ahead and run our `program` using
-- `State Int` for `m`. This test doesn't do any I/O to use
-- `program`; it's only using `IO` in order to print the
-- results and signal failure.
test :: IO ()
test =
  let
    initialState = 0
    expectedResult = 10

    -- execState :: State s a -> s -> s
    --
    -- `execState m s` will run the action `m` by
    -- supplying `s` as the initial state and will
    -- return the ending state.
    actualResult = execState (program :: State Int ()) initialState
  in
    if actualResult == expectedResult then
      putStrLn "Test Passed!"
    else
      -- Using printf in a test: fine, whatever.
      -- Using printf in production: OMG! are you crazy?!
      error $ printf
        "Test Failed: expectedResult = %d, actualResult = %d"
        expectedResult
        actualResult


-- We need to define a type that `program` can use in our
-- `main`. While `State` is fine for testing, it's not memory-
-- safe. Our program can run in bounded memory if we keep the
-- state in an `IORef` instead.
newtype App a = App { runApp :: IORef Int -> IO a }
  deriving (
    Functor, Applicative, Monad,

    -- `MonadIO` gives us `liftIO :: IO a -> App a`.
    -- Having `liftIO` essentially means our custom `App` type
    -- gets to inherit all of the built-in `IO` operations.
    MonadIO,

    -- `MonadReader (IORef Int)` gives us `ask :: App (IORef Int)`
    -- `ask` allows us to get an `IORef Int` inside an `App` do block
    -- any time we want, deterministically (i.e., we'll get the same
    -- one every time we ask).
    MonadReader (IORef Int)

    -- We're able to derive all of these wonderful instances because
    -- these instances already exist for `ReaderT (IORef Int) IO` and
    -- because `App a` and `ReaderT (IORef Int) IO a` have identical
    -- underlying implementations, namely `IORef Int -> IO a`.
    ) via ReaderT (IORef Int) IO


-- We need to implement `MonadState Int App` so that we can
-- use our `program` using `App` in place of the abstract `m`.
instance MonadState Int App where
  get :: App Int
  get = do
    stateRef <- ask                             -- ask for the state ref
    currentState <- liftIO (readIORef stateRef) -- read the current state
    return currentState                         -- return the current state

  put :: Int -> App ()
  put newState = do
    stateRef <- ask                       -- ask for the state ref
    liftIO (writeIORef stateRef newState) -- write the new state
    return ()                             -- return nothin'


-- Haskell won't let us write `main :: App ()`. What would it
-- even mean if we could? An `App ()` is really a function
-- `IORef Int -> IO ()`. To run an `App ()`, we first need to
-- create an `IORef Int`, then we can plug it into the function
-- and get the `IO ()` out.
main :: IO ()
main = do
  stateRef <- newIORef 0

  -- runApp :: App a -> IORef Int -> IO a
  --
  -- We use `runApp` to turn our `App ()` into a
  -- function `IORef Int -> IO ()` so that we can
  -- evaluate it by plugging in `stateRef`.
  runApp (program :: App ()) stateRef

  endState <- readIORef stateRef
  print endState

runState :: State s a -> s -> (s, a) (c.f. https://hackage.haskell.org/package/containers-0.7/docs/Data-Sequence-Internal.html#v:runState). That gives you a tuple with the final state and a payload.

@asarkar you want to see the intermediate states, right? You can do something like this

type StepsT :: Type -> (Type -> Type) -> Type -> Type
newtype StepsT s m a = StepsT {un :: StateT ([s], s) m a}
  deriving (Functor, Applicative, Monad) via StateT ([s], s) m

instance Monad m => MonadState s (StepsT s m) where
  get = StepsT $ fmap snd get
  put s = StepsT $ modify $ \(ss, s') -> (s' : ss, s)

runStepsT :: Functor m => StepsT s m a -> s -> m (a, [s])
runStepsT (StepsT m) s = fmap (\(a, (ss, s')) -> (a, reverse (s' : ss))) (runStateT m ([], s))

then runStepsT preserves all the intermediate states.

> runStepsT program 0
((),[0,1,4,3,10])

No modification to program is necessary.