existential optics

cats-eo.scala

package eo

import cats.{Applicative, Bifunctor, Functor, Traverse}
import cats.arrow._
import cats.implicits._

trait Accessor[F[_, _]] {
  def get[A]: [X] => F[X, A] => A
}

trait ReverseAccessor[F[_, _]] {
  def reverseGet[A]: [X] => A => F[X, A]
}

trait ForgetfulFunctor[F[_, _]] {
  def map[A, B]: [X] => F[X, A] => (A => B) => F[X, B]
}

object ForgetfulFunctor {
  given bifunctorFF[F[_, _]](using B: Bifunctor[F]): ForgetfulFunctor[F] with
    def map[A, B]: [X] => F[X, A] => (A => B) => F[X, B] =
      [X] =>
        (fa: F[X, A]) => (f: A => B) => B.bimap[X, A, X, B](fa)(identity, f)
}

trait ForgetfulContramap[F[_, _]] {
  def contramap[A, B]: [X] => F[X, A] => (B => A) => F[X, B]
}

trait ForgetfulTraverse[F[_, _], C[_[_]]] {
  def traverse[A, B, G[_]: C]: [X] => F[X, A] => (A => G[B]) => G[F[X, B]]
}

object ForgetfulTraverse {
  given tupleFTraverse: ForgetfulTraverse[Tuple2, Functor] with
      def traverse[A, B, G[_]: Functor]: [X] => ((X, A)) => (A => G[B]) => G[(X, B)] =
        [X] => (fa: (X, A)) => (f: A => G[B]) => f(fa._2).map(fa._1 -> _)

  given eitherFTraverse: ForgetfulTraverse[Either, Applicative] with
    def traverse[A, B, G[_]: Applicative]: [X] => Either[X, A] => (A => G[B]) => G[Either[X, B]] =
      [X] => (fa: Either[X, A]) => (f: A => G[B]) => fa.fold(_.asLeft[B].pure[G], f.andThen(_.map(_.asRight[X])))

  given affineFTraverse: ForgetfulTraverse[Affine, Applicative] with
    def traverse[A, B, G[_]: Applicative]: [X] => Affine[X, A] => (A => G[B]) => G[Affine[X, B]] =
          [X] => (fa: Affine[X, A]) => (f: A => G[B]) => fa match {
            case Left(x0) => x0.asLeft.asInstanceOf[Affine[X, B]].pure[G]
            case Right((x1, a)) => f(a.asInstanceOf[A]).map(b => (x1, b).asRight.asInstanceOf[Affine[X, B]])
          }

}

trait AssociativeFunctor[F[_, _], X, Y] {
  type Z
  def associateLeft[S, A, C]: (S, S => F[X, A], A => F[Y, C]) => F[Z, C]
  def associateRight[D, B, T]: (F[Z, D], F[Y, D] => B, F[X, B] => T) => T
}

object AssociativeFunctor {
  given tupleAssocF[X, Y]: AssociativeFunctor[Tuple2, X, Y] with
    type Z = (X, Y)
    def associateLeft[S, A, C]: (S, S => (X, A), A => (Y, C)) => (Z, C) =
      case (s, f, g) =>
        val (x, a) = f(s)
        val (y, c) = g(a)
        (x -> y, c)
    def associateRight[D, B, T]: ((Z, D), ((Y, D)) => B, ((X, B)) => T) => T =
      case (((x, y), d), g, f) => f(x, g(y, d))

  given eitherAssocF[X, Y]: AssociativeFunctor[Either, X, Y] with
    type Z = Either[X, Y]
    def associateLeft[S, A, C]
        : (S, S => Either[X, A], A => Either[Y, C]) => Either[Z, C] =
      case (s, f, g) =>
        f(s).fold(_.asLeft[Y].asLeft[C], g(_).leftMap(_.asRight[X]))
    def associateRight[D, B, T]
        : (Either[Z, D], Either[Y, D] => B, Either[X, B] => T) => T =
      case (Right(d), f, g)       => g(Right(f(Right(d))))
      case (Left(Right(y)), f, g) => g(Right(f(Left(y))))
      case (Left(Left(x)), _, g)  => g(Left(x))

  given affineAssocF[X0, X1, X <: Tuple2[X0, X1], Y0, Y1, Y <: Tuple2[Y0, Y1]]
      : AssociativeFunctor[Affine, X, Y] with
    type Z = (Either[X0, (X1, Y0)], (X1, Y1))

    def associateLeft[S, A, C]
        : (S, S => Affine[X, A], A => Affine[Y, C]) => Affine[Z, C] =
      case (s, f, g) =>
        f(s).fold(
          _.asLeft[(X1, Y0)].asLeft[((X1, Y1), C)],
          { case (x1, a) =>
            g(a).bimap(c => (x1 -> c).asRight[X0], p => (x1, p._1) -> p._2)
          }
        )

    def associateRight[D, B, T]
        : (Affine[Z, D], Affine[Y, D] => B, Affine[X, B] => T) => T =
      case (Right(((x1, y1), d)), f, g)  => g(Right(x1 -> f(Right(y1 -> d))))
      case (Left(Right((x1, y0))), f, g) => g(Right(x1 -> f(Left(y0))))
      case (Left(Left(x0)), _, g)        => g(Left(x0))

}

type Affine[A, B] = A match
  case (a1, a2) => Either[a1, (a2, B)]

type Forgetful[X, A] = A
type FId[A] = [X] =>> Forgetful[X, A]

object Forgetful {

  given functor: ForgetfulFunctor[Forgetful] with
    def map[A, B]: [X] => Forgetful[X, A] => (A => B) => Forgetful[X, B] =
      [X] => (a: FId[A][X]) => (f: A => B) => f(a)

  given accesor: Accessor[Forgetful] with
    def get[A]: [X] => Forgetful[X, A] => A =
      [X] => (fa: Forgetful[X, A]) => fa

  given revaccesor: ReverseAccessor[Forgetful] with
    def reverseGet[A]: [X] => A => Forgetful[X, A] =
      [X] => (a: A) => a
}

trait Optic[S, T, A, B, F[_, _]] { self =>
  type X

  def to: S => F[X, A]
  def from: F[X, B] => T

  def andThen[C, D](o: Optic[A, B, C, D, F])(using
      af: AssociativeFunctor[F, self.X, o.X]
  ): Optic[S, T, C, D, F] = new Optic {
    type X = af.Z
    def to: S => F[X, C] = s => af.associateLeft(s, self.to, o.to)
    def from: F[X, D] => T = xd => af.associateRight(xd, o.from, self.from)
  }

  def morph[G[_, _]](using cf: Composer[F, G]): Optic[S, T, A, B, G] =
    cf.to(self)

}

trait Composer[F[_, _], G[_, _]] {
  def to[S, T, A, B](o: Optic[S, T, A, B, F]): Optic[S, T, A, B, G]
}

object Composer {

  given chain[F[_, _], G[_, _], H[_, _]](using
      f2g: Composer[F, G],
      g2h: Composer[G, H]
  ): Composer[F, H] with
    def to[S, T, A, B](o: Optic[S, T, A, B, F]): Optic[S, T, A, B, H] =
      g2h.to(f2g.to(o))

  given forgetful2lens: Composer[Forgetful, Tuple2] with
    def to[S, T, A, B](
        o: Optic[S, T, A, B, Forgetful]
    ): Optic[S, T, A, B, Tuple2] = new Optic[S, T, A, B, Tuple2] {
      type X = Unit
      def to: S => Tuple2[X, A] = s => () -> o.to(s)
      def from: Tuple2[X, B] => T = xb => o.from(xb._2)
    }

  given lens2affine: Composer[Tuple2, Affine] with
    def to[S, T, A, B](
        o: Optic[S, T, A, B, Tuple2]
    ): Optic[S, T, A, B, Affine] = new Optic[S, T, A, B, Affine] {
      type X = (T, o.X)
      def to: S => Affine[X, A] = s => Right(o.to(s))
      def from: Affine[X, B] => T = {
        case Left(t)  => t
        case Right(p) => o.from(p)
      }
    }

  given prism2affine: Composer[Either, Affine] with
    def to[S, T, A, B](
        o: Optic[S, T, A, B, Either]
    ): Optic[S, T, A, B, Affine] = new Optic[S, T, A, B, Affine] {
      type X = (o.X, S)
      def to: S => Affine[X, A] = s => o.to(s).map(s -> _)
      def from: Affine[X, B] => T = xb => o.from(xb.map(_._2))
    }

}

object Optic {
  given id[F[_, _], A](using
      ra: ReverseAccessor[F],
      ac: Accessor[F]
  ): Optic[A, A, A, A, F] with
    type X = Unit
    def to: A => F[X, A] = ra.reverseGet[A][X]
    def from: F[X, A] => A = ac.get[A][X]

  given outerProfunctor[A, B, F[_, _]]
      : Profunctor[[S, T] =>> Optic[S, T, A, B, F]] with
    def dimap[S, T, R, U](
        o: Optic[S, T, A, B, F]
    )(f: R => S)(g: T => U): Optic[R, U, A, B, F] = new Optic[R, U, A, B, F] {
      type X = o.X
      def to: R => F[X, A] = f.andThen(o.to)
      def from: F[X, B] => U = o.from.andThen(g)
    }

  given innerProfunctor[S, T, F[_, _]](using
      F: ForgetfulFunctor[F]
  ): Profunctor[[B, A] =>> Optic[S, T, A, B, F]] with
    def dimap[B, A, D, C](
        o: Optic[S, T, A, B, F]
    )(f: D => B)(g: A => C): Optic[S, T, C, D, F] = new Optic[S, T, C, D, F] {
      type X = o.X
      def to: S => F[X, C] = s => F.map(o.to(s))(g)
      def from: F[X, D] => T = o.from.compose(F.map(_)(f))
    }

  def iso[S, T, A, B, F[_, _]: Accessor: ReverseAccessor](
      f: S => A,
      g: B => T
  ) = new Optic[S, T, A, B, F] {
    def to: S => F[X, A] = (s: S) =>
      summon[ReverseAccessor[F]].reverseGet[A][X](f(s))
    def from: F[X, B] => T = (fb: F[X, B]) => g(summon[Accessor[F]].get(fb))
  }

  def prism[S, T, A, B, F[_, _]](
      getOrModify: S => Either[T, A],
      reverseGet: B => T
  ) = new Optic[S, T, A, B, Either] {
    type X = T
    def to: S => Either[T, A] = getOrModify
    def from: Either[T, B] => T = identity[T] ||| reverseGet
  }

  def prism2[S, A, B](getOption: S => Option[A], reverseGet: B => S) =
    prism((s: S) => getOption(s).toRight(s), reverseGet)

  def lens[S, T, A, B](get: S => A, replace: ((S, B)) => T) =
    new Optic[S, T, A, B, Tuple2] {
      type X = S
      def to: S => Tuple2[S, A] = identity[S] &&& get
      def from: Tuple2[S, B] => T = replace
    }

  def lens2[S, T, A, B](get: S => A, replace: B => S => T) =
    lens(get, (t: (S, B)) => replace(t._2)(t._1))

  def first[A, B] = new Optic[(A, B), (A, B), A, A, Tuple2] {
    type X = B
    def to: ((A, B)) => (X, A) = _.swap
    def from: ((X, A)) => (A, B) = _.swap
  }

  def second[A, B] = new Optic[(A, B), (A, B), B, B, Tuple2] {
    type X = A
    def to: ((A, B)) => (X, B) = identity
    def from: ((X, B)) => (A, B) = identity
  }

  def optional[S, T, A, B, F[_, _]](
      getOrModify: S => Either[T, A],
      reverseGet: ((S, B)) => T
  ) =
    new Optic[S, T, A, B, Affine] {
      type X = (T, S)
      def to: S => Either[T, (S, A)] = s => getOrModify(s).map(s -> _)
      def from: Either[T, (S, B)] => T = identity[T] ||| reverseGet
    }

  def getter[S, A](get: S => A): Optic[S, Unit, A, A, Forgetful] =
    new Optic[S, Unit, A, A, Forgetful] {
      type X = Nothing
      def to: S => A = get
      def from: A => Unit = _ => ()
    }

  def fromTraverse[T[_]: Traverse, A, B]: Optic[T[A], T[B], A, B, Forgetful] =
    new Optic[T[A], T[B], A, B, Forgetful] {
      type X = Nothing
      def to: T[A] => A = ???
      def from: B => T[B] = ???
    }

  extension [S, T, A, B, F[_, _]: Accessor](o: Optic[S, T, A, B, F]) {
    def get[X](s: S): A = summon[Accessor[F]].get(o.to(s))
  }

  extension [S, T, A, B, F[_, _]: ReverseAccessor](o: Optic[S, T, A, B, F]) {
    def reverseGet[X](b: B): T =
      o.from(summon[ReverseAccessor[F]].reverseGet(b))
  }

  extension [S, T, A, B, F[_, _]: Accessor: ReverseAccessor](
      o: Optic[S, T, A, B, F]
  ) {
    def reverse = new Optic[B, A, T, S, F] {
      val A = summon[Accessor[F]]
      val RA = summon[ReverseAccessor[F]]
      type X = o.X
      def to: B => F[X, T] = (b: B) => RA.reverseGet(o.from(RA.reverseGet(b)))

      def from: F[X, S] => A = (fs: F[X, S]) => A.get(o.to(A.get(fs)))
    }
  }

  extension [S, T, A, B, F[_, _]: ForgetfulFunctor](o: Optic[S, T, A, B, F]) {
    def modify[X](f: A => B): S => T =
      s => o.from(summon[ForgetfulFunctor[F]].map(o.to(s))(f))

    def replace[X](b: B): S => T = modify[X](_ => b)

  }

  extension [S, T, A, B, F[_, _]](o: Optic[S, T, A, B, F])(using FT: ForgetfulTraverse[F, Functor]) {
    def modifyF[X, G[_]](f: A => G[B])(using G: Functor[G]): S => G[T] =
      o.to.andThen(FT.traverse(G)(_)(f)).andThen(_.map(o.from))
  }

  extension [S, T, A, B, F[_, _]](o: Optic[S, T, A, B, F])(using FT: ForgetfulTraverse[F, Applicative]) {
    def modifyA[X, G[_]](f: A => G[B])(using G: Applicative[G]): S => G[T] =
      o.to.andThen(FT.traverse(G)(_)(f)).andThen(_.map(o.from))
  }

}

@main def entry() = {
  case class UP(a: Int, b: Boolean)

  val aLens = Optic.lens[UP, UP, Int, Int](_.a, t => t._1.copy(a = t._2))
  val bLens = Optic.second[Int, Boolean]
  val aPrism = Optic.prism2[UP, Int, Int](u => Some(u.a).filter(_ % 2 == 0), UP(_, false))
  import eo.Optic.modifyF
  val mod = aLens.modifyF(i => if (i % 2 == 0) None else Some(i + 1))
  val bm = bLens.replace(false)
  val split = aPrism.modifyA(i => if (i %4 == 0) None else Some(i + 3))

  println(mod(UP(1, false)))
  println(split(UP(4, false)))
  println(bm((3, true)))

  val g: Composer[Forgetful, Affine] = summon[Composer[Forgetful, Affine]]

}