State Monad - HASKELL

instance Monad ST where
   --return :: a -> ST a
   return x = S (\s -> (x,s))

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

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

data Tree a = Leaf a | Node (Tree a) (Tree a) deriving (Show)

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

mlabel :: Tree a -> ST (Tree Int)
mlabel (Leaf _) = fresh >>= (\n -> return (Leaf n))
mlabel (Node l r) = mlabel l >>= (\l ->
                    mlabel r >>= (\r ->
                    return (Node l r)))

fresh :: ST Int
fresh = S (\n -> (n , n +1))

Hi, so this my code and I want to ensure that my understanding of the expansion of the mlabel function is correct. And I am not using any additional imports.

So suppose mlabel gets a input of Leaf 'a'

fresh >>== (\n -> return (Leaf n))
S (\n -> (n, n+1) >>== (\n -> return (Leaf n))
= S (\s -> let (x, s') = app (S (\n -> (n, n+1)) s
               (x, s') = (s, s+1)
               in app ((\n -> return (Leaf n) x) s'
                = app (S (\x -> (Leaf x, x+1)) s'
                = (\x -> (Leaf x, x+1) s'
                = (Leaf s+1, (s+1)+1)

Solution

You haven't included the definitions of your >>= and return operations for this monad, but I assume you have something like:

return x = S (\s -> (x, s))

a >>= f = S $ \s ->
  let (x, s') = app a s
  in app (f x) s'

If so, there's a problem with your expansion here:

    app ((\n -> return (Leaf n) x) s'
=   app (S (\x -> (Leaf x, x+1)) s'

You've got a missing close parenthesis in the first line, and I think you skipped too many steps and got yourself turned around.

Anyway, this should look more like this:

    app ((\n -> return (Leaf n)) x) s'
=   app (return (Leaf x)) s'                -- apply lambda
=   app (S (\s -> (Leaf x, s))) s'          -- definition of `return`
=   (Leaf x, s')                            -- evaluate `app`

Now, when we substitute in the values of x and s' from let (x, s') = (s, s+1) in ..., we get:

=   (Leaf s, s+1)

and not (Leaf s+1, (s+1)+1).

It's probably safer to rewrite the whole let xxx in yyy statement instead of trying to rewrite the xxx and yyy parts separately, so:

    S $ \s -> let (x, s') = app (S (\n -> (n, n+1))) s
              in app ((\n -> return (Leaf n)) x) s'
-- apply lambda
=   S $ \s -> let (x, s') = app (S (\n -> (n, n+1))) s
              in app (return (Leaf x)) s'
-- definition of `return`
=   S $ \s -> let (x, s') = app (S (\n -> (n, n+1))) s
              in app (S (\s -> (Leaf x, s))) s'
-- expand both `app` calls:
=   S $ \s -> let (x, s') = (s, s+1)
              in (Leaf x, s')
-- substitute `let` definitions:
=   S $ \s -> (Leaf s, s+1)