haskell state monad

state.hs

import Test.HUnit
import System.Random
import qualified Control.Monad.State as S

threeCoins :: StdGen -> (Bool, Bool, Bool)
threeCoins gen =
  let (firstCoin, newGen) = random gen
      (secondCoin, newGen1) = random newGen
      (thirdCoin, _)  = random newGen1
  in
   (firstCoin, secondCoin, thirdCoin)

type Stack = [Int]

-- pop :: Stack -> (Int, Stack)
-- pop (x : xs) = (x, xs)

-- push :: Int -> Stack -> ((), Stack)
-- push x stack = ((), x : stack)

-- stackManip :: Stack -> (Int, Stack)
-- stackManip stack = let
--   ((), newStack1) = push 3 stack
--   (a, newStack2) = pop newStack1
--   in pop newStack2

-- newtype State s a = State {runState :: s -> (a, s)}

-- instance Monad (State s) where
--   return x = State $ \s -> (x, s)
--   (State h) >>= f = State $ \s ->
--                              let (a, newState) = hs
--                                  (State g) = f a
--                              in g newState

pop :: S.State Stack Int
pop = S.state $ \ (x:xs) -> (x, xs)

push :: Int -> S.State Stack ()
push x = S.state $ \ s -> ((), x:s)

stackManip :: S.State Stack Int
stackManip = do
  push 3
  a <- pop
  pop

stackStuff :: S.State Stack ()
stackStuff = do
  a <- pop
  if a  == 5
    then push 5
    else do
      push 3
      push 8

moreStack :: S.State Stack ()
moreStack = do
  a <- stackManip
  if a == 100
    then stackStuff
    else return ()

stackyStack :: S.State Stack ()
stackyStack = do
  stackNow <- S.get
  if stackNow == [1,2,3]
    then S.put [8, 3, 1]
    else S.put [9, 2, 1]

randomSt :: (RandomGen g, Random a) => S.State g a
randomSt = S.state random

threeCoins' :: S.State StdGen (Bool, Bool, Bool)
threeCoins' = do
  a <- randomSt
  b <- randomSt
  c <- randomSt
  return (a, b, c)

tests = ["stackManip" ~: (S.runState stackManip [5, 8, 2, 1]) ~?= (5, [8,2,1]),
         "stackStuff" ~: (S.runState stackStuff [9,0,2,1,0]) ~?= ((), [8,3,0,2,1,0]),
         "moreStack" ~: (S.runState moreStack [1,2,3]) ~?= ((), [2,3]),
         "stackyStack1" ~: (S.runState stackyStack [1,2,3]) ~?=  ((), [8,3,1]),
         "stackyStack2" ~: (S.runState stackyStack [1,2,4]) ~?=  ((), [9,2,1])
         ]

runTests = do
  runTestTT $ TestList tests