Is a state monad with two state variable types (in and out) still a monad?

Question

Haskell's state monad State s a forces me to keep the same type of s during the whole do block. But since the state monad is really just a function, what if I define it as State i o a = State (i -> (o, a))?. The return and bind functions would look exactly the same as in the standard state monad, but with the types changed:

return :: a -> State st st a
bind :: (State i o a) -> (a -> (State o o' b)) -> (State i o' b)

I don't think is possible to implement Monad in Haskell using this definition because it expects a single State i o type in bind (only a can change). But this question is not about Haskell but about whether this would technically be a monad or not. Or if not, would it be some kind of superset of a monad (such that all the monad laws still apply but has some extra features)?

This is something I found would be useful in another language I'm working on, which is based on lambda calculus, so I'm using Haskell as a reference. I just don't want this to break other stuff later on where I'd expect the monad laws to apply.

Answer 1

What you are looking for is an indexed Monad. See e.g. the definition in category-extras:

The definition of an indexed Monad:

class IxApplicative m => IxMonad m where
  ibind :: (a -> m j k b) -> m i j a -> m i k b

The State indexed Monad:

class IxMonad m => IxMonadState m where
  iget :: m i i i
  iput :: j -> m i j ()