lens.scala

trait Semigroup[A]:
  def combine(x: A, y: A): A

trait Monoid[A] extends Semigroup[A]:
  def empty: A

trait Functor[F[_]]:
  def map[A, B](fa: F[A])(f: A => B): F[B]

trait Applicative[F[_]] extends Functor[F]:
  def pure[A](a: A): F[A]

  def map2[A, B, C](fa: F[A], fb: F[B])(f: (A, B) => C): F[C]

  override def map[A, B](fa: F[A])(f: A => B): F[B] =
    map2(fa, pure(()))((a, _) => f(a))

trait Monad[F[_]] extends Applicative[F]:
  def flatMap[A, B](fa: F[A])(f: A => F[B]): F[B]

  override def map2[A, B, C](fa: F[A], fb: F[B])(f: (A, B) => C): F[C] =
    flatMap(fa)(a => flatMap(fb)(b => pure(f(a, b))))

given [A]: Monoid[List[A]] with
  def empty: List[A] = List.empty
  def combine(x: List[A], y: List[A]): List[A] = x ++ y

opaque type Id[A] = A

object Id:
  inline def apply[A](a: A): Id[A] = a

  inline def value[A](fa: Id[A]): A = fa

given Monad[Id] with
  def pure[A](a: A): Id[A] = Id(a)
  def flatMap[A, B](fa: Id[A])(f: A => Id[B]): Id[B] =
    Id(f(Id.value(fa)))

opaque type Const[M, A] = M

object Const:
  inline def apply[M, A](m: M): Const[M, A] = m

  inline def value[M, A](fa: Const[M, A]): M = fa

given [M]: Functor[[A] =>> Const[M, A]] with
  def map[A, B](fa: Const[M, A])(f: A => B): Const[M, B] =
    Const(Const.value(fa))

given [M: Monoid]: Applicative[[A] =>> Const[M, A]] with
  def pure[A](a: A): Const[M, A] =
    Const(summon[Monoid[M]].empty)
  def map2[A, B, C](fa: Const[M, A], fb: Const[M, B])(f: (A, B) => C) =
    Const(summon[Monoid[M]].combine(Const.value(fa), Const.value(fb)))

trait Fold[S, A]:
  self =>

  def foldMap[M: Monoid](s: S)(f: A => M): M

  def toList(s: S): List[A] =
    foldMap(s)(a => List(a))

  def andThen[B](other: Fold[A, B]): Fold[S, B] =
    new Fold[S, B]:
      def foldMap[M: Monoid](s: S)(f: B => M): M =
        self.foldMap(s)(other.foldMap(_)(f))

trait AffineFold[S, A] extends Fold[S, A]:
  self =>

  def toOption(s: S): Option[A]

  override def foldMap[M: Monoid](s: S)(f: A => M): M =
    toOption(s).fold(summon[Monoid[M]].empty)(f)

  def andThen[B](other: AffineFold[A, B]): AffineFold[S, B] =
    new AffineFold[S, B]:
      def toOption(s: S): Option[B] =
        self.toOption(s).flatMap(other.toOption(_))

trait Getter[S, A] extends AffineFold[S, A]:
  self =>

  def get(s: S): A

  override def toOption(s: S): Option[A] =
    Some(get(s))

  def andThen[B](other: Getter[A, B]): Getter[S, B] =
    new Getter[S, B]:
      def get(s: S): B =
        other.get(self.get(s))

trait Setter[S, T, A, B]:
  self =>

  def modify(s: S)(f: A => B): T

  def replace(s: S)(b: B): T =
    modify(s)(_ => b)

  def andThen[C, D](other: Setter[A, B, C, D]): Setter[S, T, C, D] =
    new Setter[S, T, C, D]:
      def modify(s: S)(f: C => D): T =
        self.modify(s)(other.modify(_)(f))

trait Traverse[S, T, A, B] extends Setter[S, T, A, B] with Fold[S, A]:
  self =>

  def traverse[F[_]: Applicative](s: S)(f: A => F[B]): F[T]

  override def foldMap[M: Monoid](s: S)(f: A => M): M =
    Const.value(traverse[[A] =>> Const[M, A]](s)(a => Const(f(a))))

  override def modify(s: S)(f: A => B): T =
    Id.value(traverse[Id](s)(a => Id(f(a))))

  def andThen[C, D](other: Traverse[A, B, C, D]): Traverse[S, T, C, D] =
    new Traverse[S, T, C, D]:
      def traverse[F[_]: Applicative](s: S)(f: C => F[D]): F[T] =
        self.traverse(s)(other.traverse(_)(f))

trait AffineTraverse[S, T, A, B]
    extends Traverse[S, T, A, B]
    with AffineFold[S, A]:
  self =>

  def extract(s: S): Either[T, A]

  def replace(s: S)(b: B): T

  override def modify(s: S)(f: A => B): T =
    extract(s).fold(identity, a => replace(s)(f(a)))

  override def traverse[F[_]: Applicative](s: S)(f: A => F[B]): F[T] =
    extract(s).fold(
      summon[Applicative[F]].pure(_),
      a => summon[Applicative[F]].map(f(a))(replace(s)(_))
    )

  override def toOption(s: S): Option[A] =
    extract(s).toOption

  def andThen[C, D](other: AffineTraverse[A, B, C, D]): AffineTraverse[S, T, C, D] =
    new AffineTraverse[S, T, C, D]:
      def extract(s: S): Either[T, C] =
        self.extract(s).flatMap: a =>
          other.extract(a).left.map(self.replace(s)(_))

      override def replace(s: S)(d: D): T =
        self.modify(s)(other.replace(_)(d))

trait Lens[S, T, A, B] extends AffineTraverse[S, T, A, B] with Getter[S, A]:
  self =>

  def get(s: S): A

  override def extract(s: S): Either[T, A] =
    Right(get(s))

  def andThen[C, D](other: Lens[A, B, C, D]): Lens[S, T, C, D] =
    new Lens[S, T, C, D]:
      def get(s: S): C =
        other.get(self.get(s))

      override def replace(s: S)(d: D): T =
        self.modify(s)(other.replace(_)(d))

object Lens:
  def apply[S, A](get: S => A)(replace: (S, A) => S): Lens[S, S, A, A] =
    build(get)(replace)

  def build[S, T, A, B](get: S => A)(replace: (S, B) => T): Lens[S, T, A, B] =
    val funGet = get
    val funReplace = replace

    new Lens[S, T, A, B]:
      def get(s: S): A = funGet(s)
      override def replace(s: S)(b: B): T = funReplace(s, b)

trait Review[T, B]:
  self =>

  def put(b: B): T

  def andThen[D](other: Review[B, D]): Review[T, D] =
    new Review[T, D]:
      def put(d: D): T =
        self.put(other.put(d))

trait Prism[S, T, A, B] extends AffineTraverse[S, T, A, B] with Review[T, B]:
  self =>

  override def traverse[F[_]: Applicative](s: S)(f: A => F[B]): F[T] =
    extract(s).fold(
      summon[Applicative[F]].pure(_),
      a => summon[Applicative[F]].map(f(a))(put(_))
    )

  override def modify(s: S)(f: A => B): T =
    extract(s).fold(identity, a => put(f(a)))

  override def replace(s: S)(b: B): T =
    modify(s)(_ => b)

  def andThen[C, D](other: Prism[A, B, C, D]): Prism[S, T, C, D] =
    new Prism[S, T, C, D]:
      def extract(s: S): Either[T, C] =
        self.extract(s).flatMap: a =>
          other.extract(a).left.map(self.put(_))

      def put(d: D): T =
        self.put(other.put(d))

trait Grate[S, T, A, B] extends Setter[S, T, A, B] with Review[T, B]:
  self =>

  def grate(f: (S => A) => B): T

  override def modify(s: S)(f: A => B): T =
    grate(get => f(get(s)))

  override def put(b: B): T =
    grate(_ => b)

  def zipWith(s1: S, s2: S)(f: (A, A) => B): T =
    grate(get => f(get(s1), get(s2)))

  def andThen[C, D](other: Grate[A, B, C, D]): Grate[S, T, C, D] =
    new Grate[S, T, C, D]:
      def grate(f: (S => C) => D): T =
        self.grate(getSA => other.grate(getAC => f(s => getAC(getSA(s)))))

trait Iso[S, T, A, B] extends Lens[S, T, A, B] with Prism[S, T, A, B] with Grate[S, T, A, B]:
  self =>

  override def grate(f: (S => A) => B): T =
    put(f(get(_)))

  def andThen[C, D](other: Iso[A, B, C, D]): Iso[S, T, C, D] =
    new Iso[S, T, C, D]:
      def get(s: S): C =
        other.get(self.get(s))

      override def put(d: D): T =
        self.put(other.put(d))

final case class Foo(a: Int, b: String)
final case class Bar(c: Foo)

final case class Cont[R, A](f: (A => R) => R)

val fooA = Lens[Foo, Int](_.a)((s, a) => s.copy(a = a))
val fooB = Lens[Foo, String](_.b)((s, b) => s.copy(b = b))
val barC = Lens[Bar, Foo](_.c)((s, c) => s.copy(c = c))

barC.andThen(fooB).replace(Bar(Foo(1, "foo")))("bar")