Pure profunctor lenses in Scala

Lenses.scala

import scala.language.higherKinds

object Lenses {

  trait Profunctor[P[_, _]] {
    def dimap[A, B, C, D](f: C => A, g: B => D)(p: P[A, B]): P[C, D]
  }

  object Profunctor {

    def function2IsProfunctor: Profunctor[? => ?] = new Profunctor[? => ?] {
      def dimap[A, B, C, D](f: C => A, g: B => D)(p: A => B) =
        f andThen p andThen g
    }

  }

  trait Strong[P[_, _]] {
    def first[X, A, B](p: P[A, B]): P[(X, A), (X, B)]
    val profunctor: Profunctor[P]
  }

  object Strong {

    def function2IsStrong: Strong[? => ?] = new Strong[? => ?] {

      def first[X, A, B](p: (A) => B) = {
        case (x, a) => (x, p(a))
      }

      val profunctor = Profunctor.function2IsProfunctor
    }

  }

  trait Choice[P[_, _]] {
    def left[X, A, B](p: P[A, B]): P[Either[X, A], Either[X, B]]
    val profunctor: Profunctor[P]
  }

  object Choice {

    def function2IsChoice: Choice[? => ?] = new Choice[? => ?] {

      def left[X, A, B](p: A => B) = {
        case Left(x) => Left(x)
        case Right(a) => Right(p(a))
      }

      val profunctor = Profunctor.function2IsProfunctor
    }

  }

  trait PrismLike[S, T, A, B] {
    def review(b: B): T
    def under(s: S): Either[T, A]
    def preview(s: S): Option[A] = under(s).toOption
    def view(s: S)(implicit M: Monoid[A]): A = preview(s).getOrElse(M.zero)
  }

  trait LensLike[S, T, A, B] {
    def get(s: S): A
    def over(f: A => B)(s: S): T
    def set(b: B)(s: S): T = over(_ => b)(s)
  }

  trait Monoid[A] {
    def zero: A
    def combine(l: A, r: A): A
  }

  final case class ConstK[K, A, B](run: A => K) extends (A => K) {
    def apply(a: A): K = run(a)
  }

  object ConstK {
    def identity[A, B]: ConstK[A, A, B] = ConstK(Predef.identity)

    implicit def constKIsProfunctor[K]: Profunctor[ConstK[K, ?, ?]] =
      new Profunctor[ConstK[K, ?, ?]] {
        def dimap[A, B, C, D](f: C => A, g: B => D)(p: ConstK[K, A, B]) =
          ConstK[K, C, D](f andThen p)
      }

    implicit def constKIsStrong[K]: Strong[ConstK[K, ?, ?]] =
      new Strong[ConstK[K, ?, ?]] {
        def first[X, A, B](p: ConstK[K, A, B]) =
          ConstK[K, (X, A), (X, B)] {
            case (x, a) => p(a)
          }

        val profunctor = constKIsProfunctor
      }

    implicit def constKIsChoice[K](
      implicit M: Monoid[K]): Choice[ConstK[K, ?, ?]] =
      new Choice[ConstK[K, ?, ?]] {
        def left[X, A, B](p: ConstK[K, A, B]) =
          ConstK[K, Either[X, A], Either[X, B]] {
            case Left(_) => M.zero
            case Right(a) => p(a)
          }

        val profunctor = constKIsProfunctor
      }
  }

  final case class Always[A, B](stored: B) extends (A => B) {
    def apply(a: A): B = stored
  }

  object Always {
    implicit val alwaysIsProfunctor: Profunctor[Always] =
      new Profunctor[Always] {
        def dimap[A, B, C, D](f: C => A, g: B => D)(p: Always[A, B]) =
          Always(g(p.stored))
      }

    implicit val alwaysIsChoice: Choice[Always] = new Choice[Always] {
      def left[X, A, B](p: Always[A, B]) =
        Always(Right(p.stored))

      val profunctor = alwaysIsProfunctor
    }
  }

  final case class Partial[K, A, B](run: A => Either[B, K])
      extends (A => Either[B, K]) {
    def apply(a: A): Either[B, K] = run(a)
  }

  object Partial {
    def identity[A, B]: Partial[A, A, B] = Partial(Right(_))

    implicit def partialIsProfunctor[K]: Profunctor[Partial[K, ?, ?]] =
      new Profunctor[Partial[K, ?, ?]] {
        def dimap[A, B, C, D](f: C => A, g: B => D)(p: Partial[K, A, B]) =
          Partial[K, C, D] { c =>
            p(f(c)).left.map(g)
          }
      }

    implicit def partialIsChoice[K]: Choice[Partial[K, ?, ?]] =
      new Choice[Partial[K, ?, ?]] {
        def left[X, A, B](
          p: Partial[K, A, B]): Partial[K, Either[X, A], Either[X, B]] =
          Partial[K, Either[X, A], Either[X, B]] {
            case Left(x) => Left(Left(x))
            case Right(a) =>
              p.run(a) match {
                case Left(b) => Left(Right(b))
                case Right(k) => Right(k)
              }
          }

        val profunctor = partialIsProfunctor
      }
  }

  trait Optic[C[_[_, _]], S, T, A, B] {
    def apply[P[_, _]](p: P[A, B])(implicit ev: C[P]): P[S, T]

    // NOTE: Dunno how to get variance right down here :(
    //
    //    def andThen[CC[_[_, _]] <: C, X, Y](
    //      other: Optic[CC, X, Y, S, T]): Optic[CC, X, Y, A, B] = {
    //
    //      val outer: Optic[C, S, T, A, B] = this
    //
    //      new Optic[CC, X, Y, A, B] {
    //
    //        def apply[P[_, _]](p: P[A, B])(implicit ev: CC[P]): P[X, Y] = {
    //          val pp: P[S, T] = outer[P](p)(ev)
    //          other[P](pp)(ev)
    //        }
    //
    //      }
    //    }
  }

  trait Lens[S, T, A, B]
      extends Optic[Strong, S, T, A, B]
      with LensLike[S, T, A, B] {

    def apply[P[_, _]](p: P[A, B])(implicit ev: Strong[P]): P[S, T] = {
      val pp: P[(S, A), (S, B)] = ev.first(p)
      ev.profunctor.dimap[(S, A), (S, B), S, T](
        s => (s, get(s)),
        { case (s, b) => set(b)(s) }
      )(pp)
    }

    def get(s: S): A =
      apply[ConstK[A, ?, ?]](ConstK.identity).apply(s)

    def over(f: A => B)(s: S): T =
      apply[? => ?](f)(Strong.function2IsStrong)(s)

  }

  object Lens {

    def apply[S, T, A, B](get0: S => A, set0: (B, S) => T): Lens[S, T, A, B] =
      new Lens[S, T, A, B] {
        override def get(s: S) = get0(s)
        override def set(b: B)(s: S) = set0(b, s)
      }

    def _1[A, B, X]: Lens[(A, X), (B, X), A, B] =
      apply[(A, X), (B, X), A, B](_._1, {
        case (b, (a, x)) => (b, x)
      })

    def _2[A, B, X]: Lens[(X, A), (X, B), A, B] =
      apply[(X, A), (X, B), A, B](_._2, {
        case (b, (x, a)) => (x, b)
      })

  }

  trait Prism[S, T, A, B]
      extends Optic[Choice, S, T, A, B]
      with PrismLike[S, T, A, B] {

    def apply[P[_, _]](p: P[A, B])(implicit ev: Choice[P]): P[S, T] = {
      val pp: P[Either[T, A], Either[T, B]] = ev.left(p)
      ev.profunctor.dimap[Either[T, A], Either[T, B], S, T](
        s => under(s), {
          case Left(t) => t
          case Right(b) => review(b)
        }
      )(pp)
    }

    def review(b: B): T =
      apply[Always](Always(b)).stored

    def under(s: S): Either[T, A] = {
      apply[Partial[A, ?, ?]](Partial.identity).run(s)
    }

    // NOTE: This one is cool. "Same" definition as Lens. Don't know how to
    // make Scala see that clearly, however.
    override def view(s: S)(implicit M: Monoid[A]): A =
      apply[ConstK[A, ?, ?]](ConstK.identity).apply(s)
  }

  object Prism {

    def apply[S, T, A, B](review0: B => T,
                          under0: S => Either[T, A]): Prism[S, T, A, B] =
      new Prism[S, T, A, B] {
        override def review(b: B): T = review0(b)
        override def under(s: S): Either[T, A] = under0(s)
      }

    def _Left[A, B, X]: Prism[Either[A, X], Either[B, X], A, B] =
      apply[Either[A, X], Either[B, X], A, B](Left(_), {
        case Left(a) => Right(a)
        case Right(x) => Left(Right(x))
      })

    def _Right[A, B, X]: Prism[Either[X, A], Either[X, B], A, B] =
      apply[Either[X, A], Either[X, B], A, B](Right(_), {
        case Left(x) => Left(Left(x))
        case Right(a) => Right(a)
      })

  }

}