AtomicReferenceW.scala

  object ReferenceDirective extends Enumeration {
    val Replace, DoNotReplace = Value

    def replace: Value = Replace
    def doNotReplace: Value = DoNotReplace
  }
  import java.util.concurrent.atomic._
  import scalaz._
  
  /**
   * the point of this code is to be able to write a program which looks something like this:
   * 
   * {{{
   * final val ref = new AtomicReference[MyState](init)
   * 
   * val someAction: StateT[IO, MyState, (Directive, MyResult)] = ???
   * 
   * val action = ref.modifyAndMaybeApplySTM(someAction)
   * }}}
   * `action` is of type `IO[(Either[MyState, MyState], MyResult)]` - the `Either` is a 
   * left if the state is not modified and a Right if it was.
   * 
   * As an IO action it represents the atomic modification of the AtomicReference in such 
   * a way as to be composable with other IO actions
   */
  implicit class AtomicReferenceW[S](self: AtomicReference[S]) {
    final def modifyS[A](state: State[S, A]): A = testAndApplyS(_ => true, state).get //Cannot prove non-emptiness

    @annotation.tailrec final def testAndApplyS[A](p: S => Boolean, state: State[S, A]): Option[A] = {
      val s = self.get
      if (p(s)) {
        val (t, a) = state.run(s)
        if (self.compareAndSet(s, t))
          Some(a)
        else
          testAndApplyS(p, state)
      }
      else None
    }


    @annotation.tailrec final def testAndMaybeApplyS[A](p: S => Boolean, state: State[S, (ReferenceDirective.Value, A)]): (Option[S], Option[A]) = {
      val s = self.get
      if (p(s)) {
        val (t, (d, a)) = state.run(s)
        if (d == ReferenceDirective.Replace)
          if (self.compareAndSet(s, t))
            Some(t) -> Some(a)
          else
            testAndMaybeApplyS(p, state)
        else
          None -> Some(a)
      }
      else
        None -> None
    }

    final def modifyAndMaybeApplySTM[M[_]: Monad, A](stateT: StateT[M, S, (ReferenceDirective.Value, A)]): M[(Either[S, S], A)] = {
      import scalaz.syntax.monad._
      //Match is exhaustive because we will always perform the state action
      testAndMaybeApplySTM[M, A](s => Monad[M].point(true), stateT) map { case (e, Some(a)) => e -> a }
    }

    final def testAndMaybeApplySTM[M[_]: Monad, A](p: S => M[Boolean], stateT: StateT[M, S, (ReferenceDirective.Value, A)]): M[(Either[S, S], Option[A])] = {
      import scalaz.syntax.monad._
      import scalaz.std.option._
      val M = Monad[M]

      def rightS(s: S): Either[S, S] = Right(s)
      def leftS(s: S): Either[S, S] = Left(s)

      def run(s: S): M[(Either[S, S], Option[A])]
        = for {
            x <- stateT.run(s)
              (t, (d, a)) = x
            y <- if (d == ReferenceDirective.Replace)
                   M.point(self.compareAndSet(s, t)).ifM( M.point(rightS(t) -> some(a)), testAndMaybeApplySTM(p, stateT))
                 else
                   M.point(leftS(s) -> some(a))

          } yield y

      for {
        s <- M.point(self.get)
        t <- p(s).ifM(run(s), M.point(leftS(s) -> none[A]))
      } yield t
    }
  }