Experimenting with hand-rolled interpreters for free structures

FreeState.scala

package freestate

sealed abstract class FreeState[S, A] { self =>
  def tag: Int

  final def map[B](f: A => B): FreeState[S, B] = self match {
    case FreeState.Pure(a) => FreeState.Pure(f(a))
    case _ =>
      FreeState.FlatMap(self, (a: A) => FreeState.Pure(f(a)))
  }

  final def flatMap[B](f: A => FreeState[S, B]): FreeState[S, B] =
    FreeState.FlatMap(self, f)
}

object FreeState {
  final def pure[S, A](a: A): FreeState[S, A] = Pure(a)

  final def get[S]: FreeState[S, S] = FreeState.Get()

  final def set[S](s: S): FreeState[S, Unit] = FreeState.Set(s)

  object Tags {
    final val Pure = 0
    final val FlatMap = 1
    final val Get = 2
    final val Set = 3
  }

  final case class Pure[S, A](a: A) extends FreeState[S, A] {
    override final def tag = Tags.Pure
  }

  final case class FlatMap[S, A, B](fa: FreeState[S, A], f: A => FreeState[S, B]) 
      extends FreeState[S, B] {
    override final def tag = Tags.FlatMap
  }

  final case class Get[S]() extends FreeState[S, S] {
    override final def tag = Tags.Get
  }

  final case class Set[S](s: S) extends FreeState[S, Unit] {
    override final def tag = Tags.Set
  }
}

object Interpreter {
  def runOptimized[S, A](init: S)(fa: FreeState[S, A]): (S, A) = {
    var currOp: FreeState[S, Any] = fa.asInstanceOf[FreeState[S, Any]]
    var done: Boolean = false

    val conts: java.util.Stack[Any => FreeState[S, Any]] = new java.util.Stack

    var state: S = init
    var res: Any = null

    do {
      currOp.tag match {
        case FreeState.Tags.Pure =>
          res = currOp.asInstanceOf[FreeState.Pure[S, A]].a

          if (conts.empty())
            done = true
          else
            currOp = conts.pop()(res)

        case FreeState.Tags.Get =>
          res = state

          if (conts.empty())
            done = true
          else
            currOp = conts.pop()(res)

        case FreeState.Tags.Set =>
          val op = currOp.asInstanceOf[FreeState.Set[S]]

          state = op.s

          if (conts.empty())
            done = true
          else
            currOp = conts.pop()(res)

        case FreeState.Tags.FlatMap =>
          val op = currOp.asInstanceOf[FreeState.FlatMap[S, Any, Any]]

          currOp = op.fa
          conts.push(op.f)
      }
    } while (!done)

    (state,
      // We assume that res == null means we should return ()
      if (res != null) res.asInstanceOf[A]
      else ().asInstanceOf[A]
    )
  }

  def runIdiomatic[S, A](init: S)(fa: FreeState[S, A]): (S, A) = fa match {
    case FreeState.Pure(a) =>
      (init, a)

    case FreeState.FlatMap(innerFA, f) =>
      val (nextState, result) = runIdiomatic(init)(innerFA)

      runIdiomatic(nextState)(f(result))

    case FreeState.Get() =>
      (init, init.asInstanceOf[A])

    case FreeState.Set(s) =>
      (s, ().asInstanceOf[A])
  }

  def test(init: Int): FreeState[Int, Int] = for {
    _ <- FreeState.set(init)
    fst <- FreeState.get
    _ <- FreeState.set(fst + 1)
    snd <- FreeState.get
  } yield snd
}