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

@haroldcarr The class MonadState s m defines get :: m s for reading the state in m; however, the class doesn't specify a way of eliminating m and returning an s. In other words, MonadState s m does not specify a function exec :: s -> m a -> s. It's up to the particular m to define such a function.

For State s a, of course we have execState :: State s a -> (s -> s) for eliminating State and returning an s. Since program :: MonadState Int m => m () is polymorphic over m, we can "extract the final state" (so to say) of program using State Int in place of m and by evaluating execState (program :: State Int ()) initialState (for whatever initialState :: Int you please), as in test above.

If you want to instantiate program at App, then how does one "extract the final state"? Well, what's the eliminator for App a? It's runApp :: App a -> (IORef Int -> IO a). So, you turn your App a into a function IORef Int -> IO a, you apply that function to a ref :: IORef Int you've got lying around, and then you read the value out of ref, as in main above.

Of course, the real answer here is that "state" is a metaphor. It's functions all the way down. Hopefully that clears it up, but if those aren't working for you, then I must be missing what you mean by "extract the final state." Let me know if this explanation helps, or if we need to pin down your question some more. Cheers!

Very well explained. Thanks. You did answer my question by saying to eliminate App to get the ref.

Can MonadState be used with a Writer and also return a value? program here doesn't return a meaningful value. To be specific, I'm thinking about this problem.

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.