trampolined state

State.scala

package cats
package data

import free.Free.Trampoline
import free.Trampoline

// Quick port of http://blog.higher-order.com/assets/trampolines.pdf for cats
sealed abstract class State[S, A] {

  import State._

  def trampolined(initial: S): Trampoline[(S, A)]

  def apply(initial: S): (S, A) = trampolined(initial).run

  def flatMap[B](f: A => State[S, B]): State[S, B] = State.trampolined[S, B] { s =>
    Trampoline.suspend {
      trampolined(s).flatMap {
        case (s1, a) => Trampoline.suspend { f(a).trampolined(s1) }
      }
    }
  }

  def map[B](f: A => B): State[S, B] = flatMap(f andThen { b => state(b) })
}

object State {
  def apply[S, A](f: S => (S, A)): State[S, A] = new State[S, A] {
    def trampolined(s: S) = Trampoline.delay(f(s))
  }

  def trampolined[S, A](f: S => Trampoline[(S, A)]): State[S, A] = new State[S, A] {
    def trampolined(initial: S): Trampoline[(S, A)] = f(initial)
  }

  def state[S, A](a: => A): State[S, A] = apply(s => (s, a))
  def get[S]: State[S, S] = apply(s => (s, s))
  def set[S](s: S): State[S, Unit] = apply(_ => (s, ()))
  def modify[S](f: S => S): State[S, Unit] = apply(s => (f(s), ()))

  private implicit val function0Instance: Bimonad[Function0] =
    new Bimonad[Function0] {
      def extract[A](x: () => A): A = x()

      def coflatMap[A, B](fa: () => A)(f: (() => A) => B): () => B =
        () => f(fa)

      def pure[A](x: A): () => A = () => x

      def flatMap[A, B](fa: () => A)(f: A => () => B): () => B =
        () => f(fa())()
    }
}