Haskell: attempting to desugar simple State monad get and put

Question

To study the details of the State monad, I'm trying to create myself a complete desugared version of a simple state monad function, completing the thought started in How does 'get' actually /get/ the initial state in Haskell?, in the answer by J Cooper.

The example state monad function simply swaps the state and the input value, so that (conceptually) if the input is (v, s) then the output is (s, v). I show three translations, first from do notation to desugared >>= and >>, then placing those operators in function position, and finally attempting to replace them and get/put with their definitions.

The 'do' version and the first two translations work, but the final translation does not. Problems:

1. Upon loading the module, GHCi reports that z1 is not in scope.
2. I've not figured out exactly how to represent omitting the argument passing in the >> translation.

How should these be fixed?

FWIW, current Haskell platform (GHC 7.4.2).

Thanks!

-- simpleswap

import Control.Monad.State

-- =============================================
-- 'Do' version
simpleswap1 :: String -> State String String
simpleswap1 inp = do
    z1 <- get
    put inp
    return z1

-- =============================================
-- Desugared to >>= and >>
simpleswap2 :: String -> State String String
simpleswap2 inp = 
    get >>= 
    \z1 -> put inp >>
    return z1 

-- =============================================
-- >>= and >> changed to function position
simpleswap3 :: String -> State String String
simpleswap3 inp = 
    (>>=) get 
    (\z1 -> (>>) (put inp)  (return z1) )


-- =============================================
-- Attempt to translate >>=, >>, get and put

simpleswap4 :: String -> State String String
simpleswap4 inp = 
    state $ \s1 -> 
        -- (>>=) 
        let (a2, s2) = runState ( {- get -}  state $ \sg -> (sg,sg) ) s1
        in  runState (rhs1 a2) s2
        where 
            rhs1 a2 = \z1 -> 
            -- (>>)
                state $ \s3 -> 
                    let (a4, s4) = runState ( {- put inp -}  state $ \_ -> (inp, ()) ) s3
                    in runState (rhs2 a4) s4
                    where
                        rhs2 a4 = return z1

-- =============================================
main = do
    putStrLn "version 1004"
    let v = "vvv"
    let s  = "sss"
    putStrLn ("Before val: " ++ v ++ "  state: " ++ s)    
    let (v2, s2) = runState (simpleswap4 v) s
    putStrLn ("After val: " ++ v2 ++ "  state: " ++ s2)

-- =============================================

Answer 1

There are a few minor mistakes in simpleswap4. Here is a corrected version:

simpleswap4 :: String -> State String String
simpleswap4 inp = 
    state $ \s1 -> 
        -- (>>=) 
        let (z1, s2) = runState ( {- get -}  state $ \sg -> (sg,sg) ) s1
        in  runState (rhs1 z1) s2
        where 
            rhs1 z1 = 
            -- (>>)
                state $ \s3 -> 
                    let (_, s4) = runState ( {- put inp -}  state $ \_ -> ((), inp) ) s3
                    in runState rhs2 s4
                    where
                        rhs2 = return z1

I've renamed a2 to z1 (in lines 5 and 6). This doesn't change the semantics, but emphasized that the first component of the pair returned by the desugared get call is actually the result that gets bound to z1 in the previous versions of simpleswap.

The type of rhs1 should be String -> State String String. In your version, it gets an additional lambda-bound variable. It's unclear what the difference between a2 and z1 should be in your version. Removing the lambda (in line 8) also has the advantage of fixing your scoping problem. You are using z1 in the nested where-clause, but the where can only see variables bound on the left hand side of the declaration it is attached to.

In line 11, I've replaced a4 with _. This is to emphasize that (>>) does discard the result of the first action. As a consequence, rhs2 is not parameterized over this result either.