Scalaz state monad examples

I haven't seen many examples of the scalaz state monad. There is this example but it is hard to understand and there is only one other question on stack overflow it seems.

I'm going to post a few examples I've played with but I would welcome additional ones. Also if somebody can provide example on why init, modify, put and gets are used for that would be great.

Edit: here is an awesome 2 hours presentation on the state monad.

I assume, scalaz 7.0.x and the following imports (look at answer history for scalaz 6.x):

import scalaz._
import Scalaz._

The state type is defined as State[S, A] where S is type of the state and A is the type of the value being decorated. The basic syntax to create a state value makes use of the State[S, A] function:

// Create a state computation incrementing the state and returning the "str" value
val s = State[Int, String](i => (i + 1, "str")) 

To run the state computation on a initial value:

// start with state of 1, pass it to s
s.eval(1)
// returns result value "str"

// same but only retrieve the state
s.exec(1)
// 2

// get both state and value
s(1) // or s.run(1)
// (2, "str")

The state can be threaded through function calls. To do this instead of Function[A, B], define Function[A, State[S, B]]]. Use the State function...

import java.util.Random
def dice() = State[Random, Int](r => (r, r.nextInt(6) + 1))

Then the for/yield syntax can be used to compose functions:

def TwoDice() = for {
  r1 <- dice()
  r2 <- dice()
} yield (r1, r2)

// start with a known seed 
TwoDice().eval(new Random(1L))
// resulting value is (Int, Int) = (4,5)

Here is another example. Fill a list with TwoDice() state computations.

val list = List.fill(10)(TwoDice())
// List[scalaz.IndexedStateT[scalaz.Id.Id,Random,Random,(Int, Int)]]

Use sequence to get a State[Random, List[(Int,Int)]]. We can provide a type alias.

type StateRandom[x] = State[Random,x]
val list2 = list.sequence[StateRandom, (Int,Int)]
// list2: StateRandom[List[(Int, Int)]] = ...
// run this computation starting with state new Random(1L)
val tenDoubleThrows2 = list2.eval(new Random(1L))
// tenDoubleThrows2  : scalaz.Id.Id[List[(Int, Int)]] =
//   List((4,5), (2,4), (3,5), (3,5), (5,5), (2,2), (2,4), (1,5), (3,1), (1,6))

Or we can use sequenceU which will infer the types:

val list3 = list.sequenceU
val tenDoubleThrows3 = list3.eval(new Random(1L))
// tenDoubleThrows3  : scalaz.Id.Id[List[(Int, Int)]] = 
//   List((4,5), (2,4), (3,5), (3,5), (5,5), (2,2), (2,4), (1,5), (3,1), (1,6))

Another example with State[Map[Int, Int], Int] to compute frequency of sums on the list above. freqSum computes the sum of the throws and counts frequencies.

def freqSum(dice: (Int, Int)) = State[Map[Int,Int], Int]{ freq =>
  val s = dice._1 + dice._2
  val tuple = s -> (freq.getOrElse(s, 0) + 1)
  (freq + tuple, s)
}

Now use traverse to apply freqSum over tenDoubleThrows. traverse is equivalent to map(freqSum).sequence.

type StateFreq[x] = State[Map[Int,Int],x]
// only get the state
tenDoubleThrows2.copoint.traverse[StateFreq, Int](freqSum).exec(Map[Int,Int]())
// Map(10 -> 1, 6 -> 3, 9 -> 1, 7 -> 1, 8 -> 2, 4 -> 2) : scalaz.Id.Id[Map[Int,Int]]

Or more succinctly by using traverseU to infer the types:

tenDoubleThrows2.copoint.traverseU(freqSum).exec(Map[Int,Int]())
// Map(10 -> 1, 6 -> 3, 9 -> 1, 7 -> 1, 8 -> 2, 4 -> 2) : scalaz.Id.Id[Map[Int,Int]]

Note that because State[S, A] is a type alias for StateT[Id, S, A], tenDoubleThrows2 ends up being typed as Id. I use copoint to turn it back into a List type.

In short, it seems the key to use state is to have functions returning a function modifying the state and the actual result value desired... Disclaimer: I have never used state in production code, just trying to get a feel for it.

Additional info on @ziggystar comment

I gave up on trying using stateT may be someone else can show if StateFreq or StateRandom can be augmented to perform the combined computation. What I found instead is that the composition of the two state transformers can be combined like this:

def stateBicompose[S, T, A, B](
      f: State[S, A],
      g: (A) => State[T, B]) = State[(S,T), B]{ case (s, t) =>
  val (newS, a) = f(s)
  val (newT, b) = g(a) apply t
  (newS, newT) -> b
}

It's predicated on g being a one parameter function taking the result of the first state transformer and returning a state transformer. Then the following would work:

def diceAndFreqSum = stateBicompose(TwoDice, freqSum)
type St2[x] = State[(Random, Map[Int,Int]), x]
List.fill(10)(diceAndFreqSum).sequence[St2, Int].exec((new Random(1L), Map[Int,Int]()))

I stumbled on an interesting blog post Grok Haskell Monad Transformers from sigfp that has an example of applying two state monads through a monad transformer. Here is a scalaz translation.

The first example shows a State[Int, _] monad:

val test1 = for {
  a <- init[Int] 
  _ <- modify[Int](_ + 1)
  b <- init[Int]
} yield (a, b)

val go1 = test1 ! 0
// (Int, Int) = (0,1)

So I have here an example of using init and modify. After playing with it a bit, init[S] turns out to be really convenient to generate a State[S,S] value, but the other thing it allows is to access the state inside the for comprehension. modify[S] is a convenient way to transform the state inside the for comprehension. So the example above can be read as:

• a <- init[Int]: start with an Int state, set it as the value wrapped by the State[Int, _] monad and bind it to a
• _ <- modify[Int](_ + 1): increment the Int state
• b <- init[Int]: take the Int state and bind it to b (same as for a but now the state is incremented)
• yield a State[Int, (Int, Int)] value using a and b.

The for comprehension syntax already makes it trivial to work on the A side in State[S, A]. init, modify, put and gets provide some tools to work on the S side in State[S, A].

The second example in the blog post translates to:

val test2 = for {
  a <- init[String]
  _ <- modify[String](_ + "1")
  b <- init[String]
} yield (a, b)

val go2 = test2 ! "0"
// (String, String) = ("0","01")

Very much the same explanation as test1.

The third example is more tricky and I hope there is something simpler that I have yet to discover.

type StateString[x] = State[String, x]

val test3 = {
  val stTrans = stateT[StateString, Int, String]{ i => 
    for {
      _ <- init[String]
      _ <- modify[String](_ + "1")
      s <- init[String]
    } yield (i+1, s)
  }
  val initT = stateT[StateString, Int, Int]{ s => (s,s).pure[StateString] }
  for {
    b <- stTrans
    a <- initT
  } yield (a, b)
}

val go3 = test3 ! 0 ! "0"
// (Int, String) = (1,"01")

In that code, stTrans takes care of the transformation of both states (increment and suffix with "1") as well as pulling out the String state. stateT allows us to add state transformation on an arbitrary monad M. In this case the state is an Int that is incremented. If we called stTrans ! 0 we would end up with M[String]. In our example, M is StateString, so we'll end up with StateString[String] which is State[String, String].

The tricky part here is that we want to pull out the Int state value out from stTrans. This is what initT is for. It just creates an object that gives access to the state in a way we can flatMap with stTrans.

Edit: Turns out all of that awkwardness can be avoided if we truly reused test1 and test2 which conveniently store the wanted states in the _2 element of their returned tuples:

// same as test3:
val test31 = stateT[StateString, Int, (Int, String)]{ i => 
  val (_, a) = test1 ! i
  for (t <- test2) yield (a, (a, t._2))
}

Here is a very small example on how State can be used:

Let's define a small "game" where some game units are fighting the boss (who is also a game unit).

case class GameUnit(health: Int)
case class Game(score: Int, boss: GameUnit, party: List[GameUnit])


object Game {
  val init = Game(0, GameUnit(100), List(GameUnit(20), GameUnit(10)))
}

When the play is on we want to keep track of the game state, so let's define our "actions" in terms of a state monad:

Let's hit the boss hard so he loses 10 from his health:

def strike : State[Game, Unit] = modify[Game] { s =>
  s.copy(
    boss = s.boss.copy(health = s.boss.health - 10)
  )
}

And the boss can strike back! When he does everyone in a party loses 5 health.

def fireBreath : State[Game, Unit] = modify[Game] { s =>
  val us = s.party
    .map(u => u.copy(health = u.health - 5))
    .filter(_.health > 0)

  s.copy(party = us)
}

Now we can compose these actions into play:

def play = for {
  _ <- strike
  _ <- fireBreath
  _ <- fireBreath
  _ <- strike
} yield ()

Of course in the real life the play will be more dynamic, but it is food enough for my small example :)

We can run it now to see the final state of the game:

val res = play.exec(Game.init)
println(res)

>> Game(0,GameUnit(80),List(GameUnit(10)))

So we barely hit the boss and one of the units have died, RIP.

The point here is the composition. State (which is just a function S => (A, S)) allows you to define actions that produce results and as well manipulate some state without knowing too much where the state is coming from. The Monad part gives you composition so your actions can be composed:

 A => State[S, B] 
 B => State[S, C]
------------------
 A => State[S, C]

and so on.

P.S. As for differences between get, put and modify:

modify can be seen as get and put together:

def modify[S](f: S => S) : State[S, Unit] = for {
  s <- get
  _ <- put(f(s))
} yield ()

or simply

def modify[S](f: S => S) : State[S, Unit] = get[S].flatMap(s => put(f(s)))

So when you use modify you conceptually use get and put, or you can just use them alone.