How to derive a state monad from first principles?

I am trying to come up with an implementation of State Monad derived from examples of function composition. Here I what I came up with:

First deriving the concept of Monad:

data Maybe' a = Nothing' | Just' a deriving Show

sqrt' :: (Floating a, Ord a) => a -> Maybe' a
sqrt' x = if x < 0 then Nothing' else Just' (sqrt x)

inv' :: (Floating a, Ord a) => a -> Maybe' a
inv' x = if x == 0 then Nothing' else Just' (1/x)

log' :: (Floating a, Ord a) => a -> Maybe' a
log' x = if x == 0 then Nothing' else Just' (log x)

We can have function that composes these functions as follows:

sqrtInvLog' :: (Floating a, Ord a) => a -> Maybe' a
sqrtInvLog' x = case (sqrt' x) of
                  Nothing' -> Nothing'
                  (Just' y) -> case (inv' y) of
                                Nothing' -> Nothing'
                                (Just' z) -> log' z

This could be simplified by factoring out the case statement and function application:

fMaybe' :: (Maybe' a) -> (a -> Maybe' b) -> Maybe' b
fMaybe' Nothing' _ = Nothing'
fMaybe' (Just' x) f = f x

-- Applying fMaybe' =>
sqrtInvLog'' :: (Floating a, Ord a) => a -> Maybe' a
sqrtInvLog'' x = (sqrt' x) `fMaybe'` (inv') `fMaybe'` (log')`

Now we can generalize the concept to any type, instead of just Maybe' by defining a Monad =>

class Monad' m where
  bind' :: m a -> (a -> m b) -> m b
  return' :: a -> m a

instance Monad' Maybe' where
  bind' Nothing' _ = Nothing'
  bind' (Just' x) f = f x
  return' x = Just' x

Using Monad' implementation, sqrtInvLog'' can be written as:

sqrtInvLog''' :: (Floating a, Ord a) => a -> Maybe' a
sqrtInvLog''' x = (sqrt' x) \bind'` (inv') `bind'` (log')`

Trying to apply the concept to maintain state, I defined something as shown below:

data St a s = St (a,s) deriving Show

sqrtLogInvSt' :: (Floating a, Ord a) => St a a -> St (Maybe' a) a
sqrtLogInvSt' (St (x,s)) = case (sqrt' x) of
                             Nothing' -> St (Nothing', s)
                             (Just' y) -> case (log' y) of
                                            Nothing' -> St (Nothing', s+y)
                                            (Just' z) -> St (inv' z, s+y+z)

It is not possible to define a monad using the above definition as bind needs to be defined as taking in a single type "m a".

Second attempt based on Haskell's definition of State Monad:

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

First attempt to define function that is built using composed functions and maintains state:

fex1 :: Int->State Int Int
fex1 x = State { runState = \s->(r,(s+r)) } where r = x `mod` 2`

fex2 :: Int->State Int Int
fex2 x = State { runState = \s-> (r,s+r)} where r = x * 5

A composed function:

fex3 x = (runState (fex2 y)) st where (st, y) = (runState (fex1 x)) 0

But the definition newtype State s a = State { runState :: s -> (a, s) } does not fit the pattern of m a -> (a -> m b) -> m b of bind

An attempt could be made as follows:

instance Monad' (State s) where
   bind' st f = undefined
   return' x = State { runState = \s -> (x,s) }

bind' is undefined above becuase I did not know how I would implement it.

I could derive why monads are useful and apply it the first example (Maybe') but cannot seem to apply it to State. It will be useful to understand how I could derive the State Moand using concepts defined above.

Note that I have asked a similar question earlier: Haskell - Unable to define a State monad like function using a Monad like definition but I have expanded here and added more details.

Solution

Your composed function fex3 has the wrong type:

fex3 :: Int -> (Int, Int)

Unlike with your sqrtInvLog' example for Maybe', State does not appear in the type of fex3.

We could define it as

fex3 :: Int -> State Int Int
fex3 x = State { runState = \s ->
    let (y, st) = runState (fex1 x) s in
        runState (fex2 y) st }

The main difference to your definition is that instead of hardcoding 0 as the initial state, we pass on our own state s.

What if (like in your Maybe example) we wanted to compose three functions? Here I'll just reuse fex2 instead of introducing another intermediate function:

fex4 :: Int -> State Int Int
fex4 x = State { runState = \s ->
        let (y, st) = runState (fex1 x) s in
            let (z, st') = runState (fex2 y) st in
                runState (fex2 z) st' }

SPOILERS:

The generalized version bindState can be extracted as follows:

bindState m f = State { runState = \s ->
    let (x, st) = runState m s in
    runState (f x) st }

fex3' x = fex1 x `bindState` fex2
fex4' x = fex1 x `bindState` fex2 `bindState` fex2

We can also start with Monad' and types.

The m in the definition of Monad' is applied to one type argument (m a, m b). We can't set m = State because State requires two arguments. On the other hand, partial application is perfectly valid for types: State s a really means (State s) a, so we can set m = State s:

instance Monad' (State s) where
   -- return' :: a -> m a (where m = State s)
   -- return' :: a -> State s a
   return' x = State { runState = \s -> (x,s) }

   -- bind' :: m a -> (a -> m b) -> m b (where m = State s)
   -- bind' :: State s a -> (a -> State s b) -> State s b
   bind' st f =
   -- Good so far: we have two arguments
   --   st :: State s a
   --   f :: a -> State s b
   -- We also need a result
   --   ... :: State s b
   -- It must be a State, so we can start with:
       State { runState = \s ->
   -- Now we also have
   --   s :: s
   -- That means we can run st:
           let (x, s') = runState st s in
   --   runState :: State s a -> s -> (a, s)
   --   st :: State s a
   --   s :: s
   --   x :: a
   --   s' :: s
   -- Now we have a value of type 'a' that we can pass to f:
   --   f x :: State s b
   -- We are already in a State { ... } context, so we need
   -- to return a (value, state) tuple. We can get that from
   -- 'State s b' by using runState again:
           runState (f x) s'
       }