Replace Mutable Variables with the State Monad

Summary

You have a mutable variable (a var in Scala) that is both read from and written to outside of a tightly-scoped block.

Remodel the block as functions that take an initial value of the variable, and return both the final value of the variable and the original return value of the expression.

var x = 6
println(s"x = ${x}") // x = 6

x = x * 7
println(s"x = ${x}") // x = 42

We use the following form for our state monad:

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

We can use flatMap to implement other helpful methods:

implicit class StateExtras[A,S](s: State[A,S]) {
  def andThen[B](x: State[B,S]): State[B,S] =
    s.flatMap { _ => x }
  def map[B](f: A => B): State[B,S] =
    s.flatMap { a =>
      State { s =>
        (f(a),s)
      }
    }
}
val    six: State[Unit,Int] = State { _ => (                  (),  6) }
val  print: State[Unit,Int] = State { x => (println(s"x = ${x}"),  x) }
val times7: State[Unit,Int] = State { x => (                  (),x*7) }

val sixBy7: State[Unit,Int] = six andThen print andThen times7 andThen print

sixBy7 run -999 // x = 6
                // x = 42

Motivation

Mutable variables tend to preclude referential transparency, and can lead to all kind of bugs related to timing, threading, parallelism, lack of idempotency, evaluation order, lack of equational reasoning, etc.

The state monad models a mutable state change in a referentially transparent way. With it we can represent data mutation as a functional data structure.

Mechanics

When encountering a mutable variable, note where it's read from and written to:

var greeting: String = ""

val user = System.getenv("USER")
greeting = s"Greetings, ${user}!"

val osName = System.getProperty("os.name")
greeting = s"${greeting}  ${osName}, eh?  Solid."

println(greeting)

Modularize these reads and writes into discrete functions.

Each function takes an incoming version of the variable, and returns some value along with an outgoing version of the variable.

def getUser(greeting: String): (String,String) = {
  val user = System.getenv("USER")
  (user, s"Greetings, ${user}!")
}

def getOSName(greeting: String): (String,String) = {
  val osName = System.getProperty("os.name")
  (osName, s"${greeting}  ${osName}, eh?  Solid.")
}

def getGreeting(greeting: String): (String,String) =
  (greeting, greeting)

To make these explicitly composable, wrap each using the state monad:

val getUserS: State[String,String] =
  State { getUser }

val getOSNameS: State[String,String] =
  State { getOSName }

val getGreetingS: State[String,String] =
  State { getGreeting }

getUserS andThen getOSNameS andThen getGreetingS map println run ""

Example: stateful references

When it's preferable to code in terms of mutable references, we can use a data structure that wraps a mutable variable and represents its access and mutation using the state monad.

This data structure is frequently called STRef, or StateRef:

class StateRef[A,S](_a: A) {

  private var a: A = _a

  def read: State[A,S] =
    State(s => (a,s))

  def write(a2: A): State[Unit,S] =
    State({s =>
      a = a2
      ((),s)
    })

}

object StateRef {
  def apply[S,A](a: A): State[StateRef[A,S],S] =
    State(s => (new StateRef(a),s))
}

We can use one or more StateRefs together to code in terms of mutable references, while keeping the desired characteristics of the state monad:

val nameS: State[String,Unit] =
  for {
    firstSR  <- StateRef("James")
    middleSR <- StateRef("Earl")
    lastSR   <- StateRef("Douglas")
    nameSR   <- StateRef("")
    first    <- firstSR.read
    middle   <- middleSR.read
    _        <- middleSR.write("Buster") // does not overwrite `middle`
    last     <- lastSR.read
    _        <- nameSR.write(s"${first} ${middle} ${last}")
    name     <- nameSR.read
  } yield name

nameS map println run ()

Demo

Greetings, james!  Linux, eh?  Solid.
Greetings, james!  Linux, eh?  Solid.
James Earl Douglas