State Monad with multiple state values

Question

Consider the following:

do
  x1 <- new 2
  set x1 3
  x2 <- get x1
  y1 <- new 10
  set y1 20
  y2 <- get y1
  return (x2 + y2)

I want this to result in 23. Is there a way to implement something like this in pure Haskell, and if so how? I understand STRef does something like this, but I just want to do it in ordinary Haskell (not worried about efficiency at the moment). I presume I'll have to make a data type and make it an instance of Monad, but I'm not sure of the details, so a working example would be helpful.

Answer 1

This allows more than one value, but it's hairier :) This is nicely simplified with Daniel's suggestion of Dynamic.

import Data.Dynamic
import Data.Maybe
import Control.Monad.State
import Data.Map as M

newtype Ref a = Ref {ref :: Int}

type MutState = State (Int, Map Int Dynamic)

val :: Typeable a => Ref a -> MutState a
val r = snd `fmap` get >>= 
        return . fromJust . (>>= fromDynamic) .  M.lookup (ref r)

new :: Typeable a => a -> MutState (Ref a)
new a = do
  (curr, binds) <- get
  put (curr + 1, M.insert (curr + 1) (toDyn a) binds)
  return . Ref $ curr + 1

set :: Typeable a => Ref a -> a -> MutState ()
set (Ref i) a = do
  (c, m) <- get
  put (c, M.insert i (toDyn a) m)

runMut :: MutState a -> a
runMut = flip evalState (0, M.fromList [])

Then to use it

default (Int) -- too lazy for signatures :)
test :: Int
test = runMut $ do
  x1 <- new 2
  set x1 3
  x2 <- val x1
  y1 <- new 10
  set y1 20
  y2 <- val y1
  return (x2 + y2)

Refs are basically Ints with some type information attached and val will look up the appropriate Dynamic and attempt to force it into the correct type.

If this was real code, you should hide the implementations of Ref and MutState. For convenience, I've fromJusted the return of val bur if you want a safe implementation I suppose you could layer State and Maybe monads to deal with unbound variables.

And in case you are worried about the typeable constraints, as shown above they are trivially derived.

Answer 2

There is an implementation already in Control.Monad.State, but it is cumbersome for generality sake: one complication comes from MonadState class, and another from the fact that plain State is implemented in terms of more general StateT.

Here is an example of your task using that implementation. No mutability was used. Note that your example was pasted as is, just adding x prefix:

import Control.Monad.State
import qualified Data.Map as M

type MyMap a = M.Map Int a
type MyState a b = State (MyMap a) b
type MyRef = Int

xrun :: MyState a b -> b
xrun x = evalState x (M.empty)

mget :: MyState a (MyMap a)
mget = get

mput :: MyMap a -> MyState a ()
mput = put

mmodify :: (MyMap a -> MyMap a) -> MyState a ()
mmodify x = modify x

xnew :: s -> MyState s MyRef
xnew val = do
    s <- mget
    let newRef = if M.null s then 0 else fst (M.findMax s) + 1
    mput $ M.insert newRef val s
    return newRef

xset :: MyRef -> a -> MyState a () 
xset ref val = modify $ M.insert ref val

xget :: MyRef -> MyState a a
xget ref = fmap (\s -> case M.lookup ref s of Just v -> v) get

test :: MyState Int Int
test = do
  x1 <- xnew 2
  xset x1 3
  x2 <- xget x1
  y1 <- xnew 10
  xset y1 20
  y2 <- xget y1
  return (x2 + y2)

main = print $ xrun test

It is possible to implement all the functions in the module and >>=/return without using stock implementations from Control.Monad preserving the signatures.

Here it is:

module MyState (State, get, put, modify, evalState) where

newtype State s a = State (s -> (a, s))

evalState :: State s a -> s -> a
evalState (State f) = fst . f

instance Monad (State s) where
    return a = State $ \s -> (a, s)
    State f >>= g = State $ \s -> 
        case f s of 
            (a', s') -> case g a' of 
                State h -> h s'

instance Functor (State s) where
    fmap f (State g) = State $ 
        \s -> case g s of (a, s) -> (f a, s) 

get :: State s s
get = State (\s -> (s, s))

put :: s -> State s ()
put s = State $ \_ -> ((), s)

modify :: (s -> s) -> State s ()
modify f = get >>= put . f

Save it to MyState.hs and replace import Control.Monad.State with import MyState.