Simple Encoding of a purely functional Finite State Machine using Scala and scalaz

FSM.scala

package example.statemachine

import scalaz.{State, Scalaz}, Scalaz._

object FSM {
  def apply[I, S](f: PartialFunction[(I, S), S]): FSM[I, S] =
    new FSM((i, s) => f.applyOrElse((i, s), (_: (I, S)) => s))

  private def states[S, O](xs: List[State[S, O]]): State[S, List[O]] =
    xs.sequence[({type λ[α]=State[S, α]})#λ, O]

  private def modify[I, S](f: (I, S) => S): I => State[S, Unit] =
    i => State.modify[S](s => f(i, s))
}

final class FSM[I, S] private (f: (I, S) => S) {
  def apply(is: List[I]): State[S, S] =
    FSM.states(is.map(FSM.modify(f))).flatMap(_ => State.get[S])

  def run(is: List[I]): State[S, S] =
    apply(is)
}

SimpleCandyDispenser.scala

package example.statemachine

/**
 * A candy vending machine FSM
 *
 * The Machine has two types of input
 *  - you can insert a coin
 *  - you can turn the knob
 *
 * The Machine can be in one of two states
 *  - locked
 *  - unlocked
 *
 * It also tracks how many candies are left and how many coins it contains
 *
 * The rules are as follows:
 *
 * - Inserting a coin into a locked machine will cause it to unlock if there's any candy left.
 * - Turning the knob on an unlocked machine will cause it to dispense candy and become locked.
 * - Turning the knob on a locked machine or inserting a coin into an unlocked machine does nothing.
 * - A machine that’s out of candy ignores all inputs.
 *
 * This is taken from [Functional Programming in Scala](http://www.manning.com/bjarnason/)
 */
object SimpleCandyDispenser {

  sealed trait Input
  case object Coin extends Input
  case object Turn extends Input

  sealed trait Machine {
    def candies: Int
    def coins: Int
  }
  object Machine {
    def apply(candies: Int, coins: Int): Machine = LockedMachine(candies, coins)
  }
  case class LockedMachine(candies: Int, coins: Int) extends Machine
  case class UnlockedMachine(candies: Int, coins: Int) extends Machine

  val fsm =
    FSM[Input, Machine] {
      case (Coin, LockedMachine(candies, coins)) if candies > 0 =>
        UnlockedMachine(candies, coins + 1)

      case (Turn, UnlockedMachine(candies, coins)) if candies > 0 =>
        LockedMachine(candies - 1, coins)
    }
}

SimpleCandyDispenserSpec.scala

package example.statemachine

import example.statemachine.SimpleCandyDispenser.{Turn, Coin, Input, Machine}
import org.scalatest.FunSpec

class SimpleCandyDispenserSpec extends FunSpec {

  case class Result(candiesDispensed: Int, coinsAccepted: Int)

  def simulate(inputs: List[Input])(machine: Machine): (Machine, Result) = {
    val newMachine = SimpleCandyDispenser.fsm.run(inputs).exec(machine)
    (newMachine, Result(machine.candies - newMachine.candies, newMachine.coins - machine.coins))
  }

  describe("The candy machine") {
    val machine = Machine(candies = 5, coins = 10)

    it("should accept coins and dispense candies") {
      val inputs = List(
        Coin, // + 1 coin
        Turn  // - 1 candy
      )

      val (newMachine, result) = simulate(inputs)(machine)

      assert(machine.candies == 5, "The original machine remains untouched")
      assert(machine.coins == 10, "The original machine remains untouched")

      assert(newMachine.candies == 4, "The machine dispensed one candy")
      assert(newMachine.coins == 11, "The machine accepted one coin")

      assert(result.candiesDispensed == 1, "The machine dispensed one candy")
      assert(result.coinsAccepted == 1, "The machine accepted one coin")
    }

    it("should ignore inputs for the wrong state") {
      val inputs = List(
        Coin, // + 1 coin
        Turn, // - 1 candy
        Coin, // + 2 coins
        Coin, // ignored
        Turn, // - 2 candies
        Turn, // ignored
        Turn, // ignored
        Turn, // ignored
        Coin, // + 3 coins
        Coin, // ignored
        Turn  // - 3 candies
      )

      val (newMachine, result) = simulate(inputs)(machine)

      assert(machine.candies == 5, "The original machine remains untouched")
      assert(machine.coins == 10, "The original machine remains untouched")

      assert(newMachine.candies == 2, "The machine dispensed four candies")
      assert(newMachine.coins == 13, "The machine accepted four coins")

      assert(result.candiesDispensed == 3, "The machine dispensed four candies")
      assert(result.coinsAccepted == 3, "The machine accepted four coins")
    }

    it("should ignore inputs when empty") {
      // get 6 candies
      val inputs = List.fill(6)(List(Coin, Turn)).flatten

      val (newMachine, result) = simulate(inputs)(machine)

      assert(machine.candies == 5, "The original machine remains untouched")
      assert(machine.coins == 10, "The original machine remains untouched")

      assert(newMachine.candies == 0, "The machine dispensed five candies")
      assert(newMachine.coins == 15, "The machine accepted five coins")

      assert(result.candiesDispensed == 5, "The machine dispensed only five candies")
      assert(result.coinsAccepted == 5, "The machine accepted only five coins")
    }
  }
}