# Functional State in Scala

Given a stateful function `run` and a class `State` that wraps it:

``````case class State[A,S](run: S => (A,S)) {

def map[B](g: A => B): State[B,S] = ???

def flatMap[B](g: A => State[B,S]): State[B,S] = ???
}``````

Implement `map` and `flatMap` so that we can use `State` in a `for`-comprehension:

``````def put(k: String, v: Int): State[Int, Map[String,Int]] =
State(m => (v, m + (k -> v)))

def get(k: String): State[Int, Map[String,Int]] =
State(m => (m(k), m))

def getAndDouble(k: String): State[Int, Map[String,Int]] =
State({ m =>
val v = m(k)
(v, m + (k -> v * 2))
})

val resultS: State[Tuple5[Int,Int,Int,Int,Int], Map[String,Int]] =
for {
a <- put("foo", 21)      // a = 21, state = Map(foo -> 21)
b <- get("foo")          // b = 21, state = Map(foo -> 21)
c <- getAndDouble("foo") // c = 21, state = Map(foo -> 42)
d <- getAndDouble("foo") // d = 42, state = Map(foo -> 84)
e <- get("foo")          // e = 84, state = Map(foo -> 84)
} yield (a,b,c,d,e)

println(resultS.run(Map.empty)) // ((0,21,21,42,84),Map(foo -> 84))``````

# Solution

Here is one possible solution.
``````case class State[A,S](run: S => (A,S)) {

def map[B](g: A => B): State[B,S] =
State({ s =>
val (a,s2) = run(s)
(g(a),s2)
})

def flatMap[B](g: A => State[B,S]): State[B,S] =
State({ s =>
val (a,s2) = run(s)
g(a).run(s2)
})
}``````

## Demo

This file is literate Scala, and can be run using Codedown:

``````\$ curl https://earldouglas.com/posts/exercises/scala-state.md |
codedown scala |
xargs -0 scala -nc -e
((21,21,21,42,84),Map(foo -> 84))``````