Scala State monad - combining different state types

Question

I'm wrapping my head around State monad. Trivial examples are easy to understand. I'm now moving to a real world case where the domain objects are composite. For example, with the following domain objects (they don't make much sense, just sheer example):

case class Master(workers: Map[String, Worker])
case class Worker(elapsed: Long, result: Vector[String])
case class Message(workerId: String, work: String, elapsed: Long)

Considering Worker as S types in State[S, +A] monad it's quite easy to write a few combinators like these:

type WorkerState[+A] = State[Worker, A]
def update(message: Message): WorkerState[Unit] = State.modify { w =>
    w.copy(elapsed = w.elapsed + message.elapsed,
           result = w.result :+ message.work)
}
def getWork: WorkerState[Vector[String]] = State { w => (w.result, w) }
def getElapsed: WorkerState[Long] = State { w => (w.elapsed, w) }
def updateAndGetElapsed(message: Message): WorkerState[Long] = for {
    _ <- update(message)
    elapsed <- getElapsed
} yield elapsed
// etc.

What is the idiomatic way to combine these with the Master state combinators? e.g.

type MasterState[+A] = State[Master, A]
def updateAndGetElapsedTime(message: Message): MasterState[Option[Long]]

I can implement this like so:

def updateAndGetElapsedTime(message: Message): MasterState[Option[Long]] =   
    State { m =>
        m.workers.get(message.workerId) match {
            case None => (None, m)
            case Some(w) =>
                val (t, newW) = updateAndGetElapsed(message).run(w)
                (Some(t), m.copy(m.workers.updated(message.workerId, newW))
        }
    }

What I don't like is that I have to manually run the State monad inside the last transformer. My real world example is a bit more involved. With this approach it quickly gets messy.

Is there more idiomatic way to run this sort of incremental updates?

Nice question! Are you referring to some concrete `State` implementation like `scalaz`? — Odomontois, Jun 02 '15 at 06:43
It's definitely look's like nice example for `LensT` usage, can't wait to see some expert's answer. — Odomontois, Jun 02 '15 at 06:49

score 8 · Answer 1 · answered Jun 02 '15 at 07:07

It's possible to do this pretty nicely by combining lenses and the state monad. First for the setup (I've edited yours lightly to get it to compile with Scalaz 7.1):

case class Master(workers: Map[String, Worker])
case class Worker(elapsed: Long, result: Vector[String])
case class Message(workerId: String, work: String, elapsed: Long)

import scalaz._, Scalaz._

type WorkerState[A] = State[Worker, A]

def update(message: Message): WorkerState[Unit] = State.modify { w =>
  w.copy(
    elapsed = w.elapsed + message.elapsed,
    result = w.result :+ message.work
  )
}

def getWork: WorkerState[Vector[String]] = State.gets(_.result)
def getElapsed: WorkerState[Long] = State.gets(_.elapsed)
def updateAndGetElapsed(message: Message): WorkerState[Long] = for {
  _ <- update(message)
  elapsed <- getElapsed
} yield elapsed

And now for a couple of general purpose lenses that allow us to look inside a Master:

val workersLens: Lens[Master, Map[String, Worker]] = Lens.lensu(
  (m, ws) => m.copy(workers = ws),
  _.workers
)

def workerLens(workerId: String): PLens[Master, Worker] =
  workersLens.partial andThen PLens.mapVPLens(workerId)

And then we're basically done:

def updateAndGetElapsedTime(message: Message): State[Master, Option[Long]] =
  workerLens(message.workerId) %%= updateAndGetElapsed(message)

Here the %%= just tells us what state operation to perform once we've zoomed in to the appropriate worker via our lens.

I'm running my own poor man's monad implementations. Do I understand it correnctly that to get this working I need an implementation of Lens and integration of State with Lens (in particular that %%= operation)? — ak., Jun 02 '15 at 07:15
Also what's your take on using Scalaz vs handwritten implementations of monads? — ak., Jun 02 '15 at 07:15
If you're doing this kind of thing I'd strongly suggest using Scalaz or cats instead of rolling your own. — Travis Brown, Jun 02 '15 at 07:16
@ak As i can see, main difference between your own `State` and `scalaz.State` is the order in `runState`'s result tuple. I assume you could at least glue those implementations via implicit conversion. — Odomontois, Jun 02 '15 at 07:20
This seems like a good solution to a common problem. It also works well you have multiple components of different types with their own states. I'm wondering if this combination of State + Lens has been adopted by any typelevel library yet? Perhaps there would also be a typeclass to help with forwarding calls to components. — James McCabe, Oct 24 '15 at 19:54

Scala State monad - combining different state types

1 Answers1