Finalizers with monadic regions

gistfile1.scala

object SafeIO {

  trait Brace[M[_]] extends Monad[M] {
    def brace[A,B,C](acquire: M[A])(release: A => M[B], go: A => M[C]): M[C]
    def snag[A](m: M[A], f: Throwable => M[A]): M[A]
    def lift[A](t: Task[A]): M[A]
  }

  object Brace {
    def apply[M[_]:Brace]: Brace[M] = implicitly[Brace[M]]

    implicit val taskBrace: Brace[Task] = new Brace[Task] {
      def brace[A,B,C](acquire: Task[A])(
                       release: A => Task[B],
                       go: A => Task[C]) = for {
        a <- acquire
        // NB: the type of `onFinish` is not as general as it should be
        r <- go(a).onFinish(_ => release(a) flatMap (_ => Task.now(())))
      } yield r
      def snag[A](m: Task[A], f: Throwable => Task[A]) =
        m handleWith PartialFunction(f)
      def lift[A](t: Task[A]) = t
      def point[A](a: => A) = Task.delay(a)
      def bind[A,B](a: Task[A])(f: A => Task[B]) = a flatMap f
    }

    implicit def readerBrace[M[_]:Brace,R]: Brace[Kleisli[M,R,?]] =
      new Brace[Kleisli[M,R,?]] {
        def brace[A,B,C](acquire: Kleisli[M,R,A])(
                         release: A => Kleisli[M,R,B],
                         go: A => Kleisli[M,R,C]) =
          Kleisli { r =>
            Brace[M].brace(acquire(r))(release(_)(r), go(_)(r))
          }
        def snag[A](m: Kleisli[M,R,A], f: Throwable => Kleisli[M,R,A]) =
          Kleisli { r =>
            Brace[M].snag(m(r), f(_)(r))
          }
        def lift[A](t: Task[A]) = Brace[M].lift(t).liftKleisli
        def point[A](a: => A) = Kleisli.pure(a)
        def bind[A,B](a: Kleisli[M,R,A])(f: A => Kleisli[M,R,B]) = a flatMap f
      }

    implicit def regionBrace[M[_]:Brace,R]: Brace[Region[M,R,?]] =
      new Brace[Region[M,R,?]] {
        val B = readerBrace[M, IORef[List[Finalizer]]]
        def brace[A,B,C](acquire: Region[M,R,A])(
                         release: A => Region[M,R,B],
                         go: A => Region[M,R,C]) =
          Region(B.brace(acquire.run)(release(_).run, go(_).run))
        def snag[A](m: Region[M,R,A], f: Throwable => Region[M,R,A]) =
          Region(B.snag(m.run, f(_).run))
        def lift[A](t: Task[A]) = Region(B.lift(t))
        def point[A](a: => A) = Region(B.point(a))
        def bind[A,B](a: Region[M,R,A])(f: A => Region[M,R,B]) = a flatMap f
      }

    implicit def processBrace[M[_]:Brace]: Brace[Process[M,?]] =
      new Brace[Process[M,?]] {
        val B = Brace[M]
        def brace[A,B,C](acquire: Process[M,A])(
                         release: A => Process[M,B],
                         go: A => Process[M,R,C]) = {
      }
  }

  // A resource belonging to the monad M
  // (might be tempted to use Codensity here,
  //  but we explicitly don't want to allow lowering back to `M`)
  class Resource[M[_], R](r: R)
    def apply[A](f: R => M[A]): M[A] = f(r)
  }

  type Finalizer = Task[Unit]

  // Region monad M with safe I/O resources
  // To run all finalizers at the end of the computation, we maintain a mutable list
  // of finalizers and update it when a new resource is acquired.
  case class Region[M[_],R,A](run: Kleisli[M, IORef[List[Finalizer]], A]) {
    def map[B](f: A => B): Region[M,R,B] = Region(run map f)
    def flatMap[B](f: A => Region[M,R,B]): Region[M,R,B] = Region(run flatMap f)
  }

  // IO monad with safe resources
  type SIO[S, A] = Region[Task, S, A]

  // Acquire a resource, register a finalizer, and perform some action
  def acquire[M[_]:Brace,S](resource: M[R])(release: R => Finalizer[M])
    : Region[M,S,Resource[Region[M,S,?],R]] = {
    val B = readerBrace[M, List[Finalizer]]
    Region(for {
      r <- acquire.liftKleisli
      finalizers <- ask
      _ <- lift(finalizers.modify(release(r) :: _))
    } yield Resource(r))
  }

  type BracketedProcess[M[_],A] = Forall[Region[Process[M,?],?,A]]

  // Acquires a resource and registers a finalizer for when it goes our of scope
  def bracket[M[_]:Brace,R,A](resource: M[R])(
                              release: R => Process[M,Unit],
                              go: R => Process[M,A]): BracketedProcess[M,A] =
      new BracketedProcess[M,A] {
        def apply[S] =
          Region(Process.eval(acquire(resource)(release(_).run)).run
                 .mapK(mr => Process.eval(mr).flatMap(_(go))))
      }

  // Opens a safe region that tracks finalizers dynamically
  // and runs them as soon as they're out of scope
  // Ensures that all available resources are not yet finalized
  // and that all finalizers are run exactly once.
  def region[M[_]:Brace,A](m: Forall[Region[M,?,A]]): M[A] = {
    val B = Brace[M]
    def cleanup(finalizers: IORef[List[Finalizer]]) =
      B.lift(finalizers.read.flatMap(_.sequence_))
    B.brace(B.lift(IORef.create(List[Finalizer]())))(
            cleanup,
            m.apply.run.apply(_))
  }

  def runBracket[M[_]:Brace,A](p: BracketedProcess[M,A]): Process[M,A] = {
    val B = Brace[M]
    def cleanup(finalizers: IORef[List[Finalizer]]) =
      B.lift(finalizers.read.flatMap(_.sequence_))
    B.brace(B.lift(IORef.create(List[Finalizer]())))(
            cleanup,
            p.apply.run.apply(_))
  }
}