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.
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 ()