module ST where

import MyState

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

unST (ST m) = m

instance Monad (ST s) where
  return a = ST (\ s -> (a, s))
  (ST m) >>= k = ST (\ s0 -> let { (a, s1) = m s0 } in unST (k a) s1)

instance Functor (ST s) where
  fmap f x = x >>= \ a -> return (f a)

instance Applicative (ST s) where
  pure = return
  f <*> x = f >>= \ g -> x >>= \ a -> return (g a)

extendST :: Pos s t -> ST t a -> ST s a
extendST p (ST m) = ST (\ s ->
                      let (t, rest) = p s
                          (a, t1)   = m t
                        in (a, rest t1))

shrinkST :: (s -> t) -> (t -> s) -> ST s a -> ST t a
shrinkST extract rest (ST m) = ST (\ t ->
                          let (a, s) = m (rest t)
                            in (a, extract s))

instance MyState ST where
  get p   = ST (\ s -> (fst (p s), s))
  set p v = ST (\ s -> ((), snd (p s) v))
  extend  = extendST
  shrink  = shrinkST

evalST st s = fst (unST st s)