Stateful computations in Scala using monads - scala

Question

I've been trying to learn functional programming in Scala, and I finally managed to understand how to use the for comprehension to work with state:
#!/usr/bin/env scala
case class State[A,S](run: S => (A,S)) {
def map[B](f: A => B): State[B,S] =
State(s => {
val (a, s1) = run(s)
(f(a), s1)
})
def flatMap[B](f: A => State[B,S]): State[B,S] =
State(s => {
val (a,s1) = run(s)
f(a).run(s1)
})
}
val increment = State[Unit,Int] {
x => ((),x+1)
}
val read = State[Int,Int] {
x => (x,x)
}
def prog = for {
_ <- increment
x <- read
_ <- increment
y <- read
} yield (x,y)
val ans = prog.run(0)._1
println(ans)
Although this runs fine, I did not manage to do something similar using a state monad, it is more complicated than, e.g., Option, because it takes an extra type. How do I do something similar to this code with a state monad?
EDIT: Apparently, my question was not clear. I want to run this using a monad trait, like this one, which I took from "Functional programming in Scala":
def stateMonad[S] = new Monad[({type lambda[x] = State[S,x]})#lambda] {
def unit[A](a: => A): State[S,A] = State(s => (a, s))
def flatMap[A,B](st: State[S,A])(f: A => State[S,B]): State[S,B] =
st flatMap f
}
And then perform the computation by instantiating this with something like val M = stateMonad[Int].

Answer 1

After trying around, I managed to get it working. So, I think I'll end up answering my own question. The solution is
trait Monad[M[_]] {
def unit[A](a: => A): M[A]
def flatMap[A,B](ma: M[A])(f: A => M[B]): M[B]
}
class StateMonad[S] extends Monad[({type lambda[x] = State[x,S]})#lambda] {
def unit[A](a: => A): State[A,S] = State(s => (a, s))
def flatMap[A,B](st: State[A,S])(f: A => State[B,S]): State[B,S] =
st flatMap f
def increment: State[Unit,Int] = State(x => ((),x+1))
def read: State[Int,Int] = State(x => (x,x))
}
val m = new StateMonad[Int]
def prog = for {
_ <- m.increment
x <- m.read
_ <- m.increment
y <- m.read
} yield (x,y)
The idea is to make the StateMonad class inherit from Monad, and include all the functions that manipulate state as methods of the StateMonad class. As was pointed out, my previous code could already be considered a monad, but I think doing it this way is better.