I'm implementing my own FSM Actors because akka.actor.FSM is disappointing. Part of the implementation is a new way of dealing with data mutation on events. I want Vars, but I want to enforce the scope in which reads or writes can happen. Following is an implementation for scoped, thread-safe Refs that also preserve their previous state before t…

Ref.scala

package fsm

import scala.annotation.implicitNotFound
import scala.util.control._
import RefContext._
import Ref._

/** 
 * Ref is a scoped type-safe mutable reference.
 *
 * ==Overview==
 *
 * It's scope is to protect developers from accidentally mutating data
 * in the wrong context:
 *
 *  - the type system insures that reads only happen in a proper read
 *    context and that writes only happen in a proper write context
 * 
 *  - at runtime, if the read/write evidence escapes the current lexical
 *    scope, then an IllegalStateException is thrown
 *
 *  - also at runtime, if the read/write evidence escapes the thread in 
 *    which the read/write context was started, then an IllegalStateException
 *    is thrown
 *
 * The usefulness in the context of a FSM Actor is to only allow data mutations
 * in a single lexical, thread-local scope and to only allow reads-only
 * accesses in `onTransition` methods.
 *
 * {{{
 * scala> val ref = Ref(0)
 * ref: fsm.Ref[Int] = Ref(0)
 *
 * scala> ref.get 
 * <console>:12: error: You cannot read from this variable (an implicit ReadPermission is required)
 * 
 * scala> RefContext.reads { implicit ctx =>
 *   ref.get
 * }
 *
 * res: Int = 0
 * 
 * scala> ref := 12
 * <console>:12: error: You cannot write in this variable (an implicit WritePermission is required)
 * 
 * scala> Ref.Context.writes { implicit ctx =>
 *   ref := 12
 * }
 * }}}
 *
 * ==Transitions / Transactions==
 *
 * The write context behaves like a dumb transaction - so say you're in a write 
 * context and you're making changes to Refs, if an exception gets
 * triggered, then all data changes to those Refs will be rolled-back to their 
 * previous values. When in a read context, you also have access to the old value.
 * Again, in the context of a FSM, this is useful for inspecting previous changes
 * that help to make decisions in `onTransition` functions.
 *
 * Example:
 * {{{
 * scala> val ref = Ref("old-value")
 * ref: fsm.Ref[String] = Ref(old-value)
 *
 * scala> RefContext.writes { implicit ctx => ref := "new-value" }
 *
 * scala> RefContext.reads { implicit ctx => s"OLD: ${ref.old} ; NEW: ${ref.get}" }
 * res: String = OLD: Some(old-value) ; NEW: new-value
 * }}}
 *
 * If an exception is thrown, then the changes are not committed:
 * {{{
 * scala> val ref = Ref("old-value")
 * ref: fsm.Ref[String] = Ref(old-value)
 *
 * scala> RefContext.writes { implicit ctx => ref := "changed-again?"; throw new RuntimeException("dummy") }
 * java.lang.RuntimeException: dummy
 *   at $anonfun$1.apply(<console>:12)
 *   ...
 *
 * scala> RefContext.reads { implicit ctx => s"OLD: ${ref.old} ; NEW: ${ref.get}" }
 * res: String = OLD: None ; NEW: old-value
 * }}}
 */
@specialized final class Ref[T] private (initialValue: T) {
  /** 
   * Returns current value
   */
  def apply()(implicit ev: ReadPermission): T = {
    checkPermission(ev)
    _tmpValue getOrElse _currentValue
  }

  /** 
   * Alias for apply(), returns current value
   */
  def get(implicit ev: ReadPermission): T =
    apply()(ev)

  /** 
   * Updates current value.
   */
  def update(value: T)(implicit ev: WritePermission): Unit = {
    checkPermission(ev)
    ev.register(this)
    _tmpValue = Some(value);
  }

  /** 
   * Updates current value, a nice alias for update()
   */
  def `:=`(value: T)(implicit ev: WritePermission): Unit = 
    update(value)

  /** 
   * Returns old value (before the write context happened and committed)
   */
  def old(implicit ev: ReadPermission): Option[T] = {
    checkPermission(ev)
    _oldValue
  }

  /** 
   * Addition/increment operator for numbers, updates current value.
   */
  def `+=`(value: T)(implicit ev: WritePermission, num: Numeric[T]): Unit =
    update(num.plus(get, value))

  /** 
   * Subtraction/decrement operator for numbers, updates current value.
   */
  def `-=`(value: T)(implicit ev: WritePermission, num: Numeric[T]): Unit =
    update(num.minus(get, value))

  /** 
   * Multiplication operator for numbers, updates current value.
   */
  def `*=`(value: T)(implicit ev: WritePermission, num: Numeric[T]): Unit =
    update(num.times(get, value))

  /** 
   * Quotation operator for numbers, updates current value.
   */
  def `/=`(value: T)(implicit ev: WritePermission, num: Integral[T]): Unit =
    update(num.quot(get, value))

  override def toString = 
    s"Ref(${_tmpValue getOrElse _currentValue})"

  /** 
   * Commits the current changes - Internal API
   */
  private[fsm] def commit(): Unit = 
    for (tmp <- _tmpValue) {
      _oldValue = Some(_currentValue)
      _currentValue = tmp
      _tmpValue = None
    }

  /** 
   * Rollsback the current changes - Internal API
   */
  private[fsm] def rollback(): Unit = 
    _tmpValue = None

  private[this] var _currentValue: T = initialValue;
  private[this] var _tmpValue = Option.empty[T]
  private[this] var _oldValue = Option.empty[T]
}

object Ref {
  /** 
   * Ref constructor
   */
  def apply[T](initialValue: T): Ref[T] = 
    new Ref(initialValue)

  /** 
   * ReadPermission/WritePermission need an `isValid` check
   * that gets called at runtime (in case it fails, the check throws an
   * IllegalStateException - which btw, kills Actors and possibly the JVM
   * as Akka treats IllegalStateExceptions very seriously)
   */
  sealed trait Permission {
    private[fsm] def isValid: Boolean
  }

  /** 
   * Required evidence in case you want to do reads.
   */
  @implicitNotFound("You cannot read from this variable (an implicit ReadPermission is required)")
  sealed trait ReadPermission extends Permission

  /** 
   * Required evidence in case you want to do writes.
   */
  @implicitNotFound("You cannot write in this variable (an implicit WritePermission is required)")
  sealed trait WritePermission extends ReadPermission {
    /** 
     * Part of the internal API, register this var's usage
     * within the active context, used for transactionality.
     */
    private[fsm] def register(ref: Ref[_]): Unit
  }
}

object RefContext {
  import Ref._

  /** 
   * Creates a context meant for reading
   */
  def reads[T](cb: ReadPermission => T): T = {
    val ev = new LocalReadPermission()
    try { cb(ev) } finally { ev.invalidate() }
  }

  /** 
   * Creates a context meant for writing
   */
  def writes[T](cb: WritePermission => T): T = {
    val ev = new LocalWritePermission()
    try { 
      val ret = cb(ev)
      ev.commit() 
      ret
    }
    catch { 
      case NonFatal(ex) => 
        ev.rollback()
        throw ex 
    }
    finally { 
      ev.invalidate() 
    }
  }

  /** 
   * INTERNAL API ...
   */

  private[fsm] class LocalReadPermission extends ReadPermission with Context
  private[fsm] class LocalWritePermission extends WritePermission with ReadPermission with Context

  private[fsm] trait Context {
    @volatile
    private[this] var registered = Set.empty[Ref[_]]

    val _isValid = {
      val th = new ThreadLocal[Boolean] { override val initialValue = false }
      th.set(true)
      th
    }

    def isValid: Boolean = _isValid.get()
    def invalidate(): Unit = _isValid.set(false)

    def register(ref: Ref[_]): Unit = 
      registered = registered + ref
    
    def commit(): Unit = 
      for (ref <- registered) {
        ref.commit()
      }

    def rollback(): Unit =
      for (ref <- registered) {
        ref.rollback()
      }
  }

  @inline
  private[fsm] def checkPermission(ev: Permission): Unit = 
    if (!ev.isValid) ev match {
      case _:WritePermission =>
        throw new IllegalStateException("WritePermission for writing variable has escaped its scope")
      case _:ReadPermission =>
        throw new IllegalStateException("ReadPermission for reading variable has escaped its scope")
    }
}