-- Tree relabelling example from chapter 12 of Programming in Haskell,
-- Graham Hutton, Cambridge University Press, 2016.

-- The state monad

type State = Int

newtype ST a = S (State -> (a,State))

app :: ST a -> State -> (a,State)
app (S st) x = st x

instance Functor ST where
   -- fmap :: (a -> b) -> ST a -> ST b
   fmap g st = S (\s -> let (x,s') = app st s in (g x, s'))

instance Applicative ST where
   -- pure :: a -> ST a
   pure x = S (\s -> (x,s))

   -- (<*>) :: ST (a -> b) -> ST a -> ST b
   stf <*> stx = S (\s ->
      let (f,s')  = app stf s
          (x,s'') = app stx s' in (f x, s''))

instance Monad ST where
   -- (>>=) :: ST a -> (a -> ST b) -> ST b
   st >>= f = S (\s -> let (x,s') = app st s in app (f x) s')

-- Relabelling trees

inc :: ST String 
inc = S (\n -> ("", n+1))

double :: ST String
double = S (\n -> ("", 2*n))

report :: ST String
report = S (\n ->( show  n, n))

saveState :: ST State
saveState = S (\n -> (n,n))

resetState :: State -> ST State
resetState m = S (\n -> (0,m))

dothis1 = 
   report >>= (\str1 ->
   inc >>
   inc >>
   report >>= (\str2 ->
   inc >>
   inc >>
   return ("str1 = "++ str1 
      ++ "\nstr2 = " ++ str2)))

dothis2 = do 
   str1 <- report
   inc
   inc
   str2 <- report
   inc
   inc
   return ("str1 = "++ str1 
      ++ "\nstr2 = " ++ str2)