Lens + State

Example.scala

object Example {

  def doStuff(partition: String, itemId: String): State[ExampleState, Seq[Item]] =
    for {
      _ <- getOrCreateItem(partition, itemId)
      _ <- setPending(partition, itemId)
      pending <- unconsPending(partition)
    } yield pending

  private def getOrCreateItem(partition: String, item: String): State[ExampleState, Item] =
    Lens.item(partition, item).getOrElseSet(Eval.now(Item(item, pending = false)))

  private def setPending(partition: String, item: String): State[ExampleState, Unit] =
    Lens.pending(partition, item).tryModifyS(_ => Eval.now(true))

  private def unconsPending(partition: String): State[ExampleState, Seq[Item]] =
    Lens.items(partition).applyS { itemsMap =>
      val (pending, notPending) = itemsMap.partition { case (_, item) => item.pending }

      Eval.now(notPending -> pending.values.toSeq)
    }

}

object Lens {

  def items(partition: String): Lens[ExampleState, Map[String, Item]] =
    (ExampleState.partitionsLens composeLens
      mapValue(partition, Partition.empty) composeLens
      Partition.itemsLens)

  def item(partition: String, item: String): Optional[ExampleState, Item] =
    (ExampleState.partitionsLens composeLens
      mapValue(partition, Partition.empty) composeLens
      Partition.itemsLens composeOptional
      mapValue(item))

  def itemId(partition: String, item: String): Optional[ExampleState, String] =
    (ExampleState.partitionsLens composeLens
      mapValue(partition, Partition.empty) composeLens
      Partition.itemsLens composeOptional
      mapValue(item) composeLens
      Item.idLens)

  def pending(partition: String, item: String): Optional[ExampleState, Boolean] =
    (ExampleState.partitionsLens composeLens
      mapValue(partition, Partition.empty) composeLens
      Partition.itemsLens composeOptional
      mapValue(item) composeLens
      Item.pendingLens)

}

Models.scala

case class ExampleState(partitions: Map[String, Partition], finished: Boolean)

case class Partition(items: Map[String, Item])

case class Item(id: String, pending: Boolean)

object ExampleState {
  val partitionsLens: Lens[ExampleState, Map[String, Partition]] = GenLens[ExampleState](_.partitions)
  val finishedLens: Lens[ExampleState, Boolean] = GenLens[ExampleState](_.finished)
}

object Partition {

  val empty: Partition = Partition(Map.empty)

  val itemsLens: Lens[Partition, Map[String, Item]] = GenLens[Partition](_.items)
}

object Item {
  val idLens: Lens[Item, String] = GenLens[Item](_.id)
  val pendingLens: Lens[Item, Boolean] = GenLens[Item](_.pending)
}

package.scala

package object mono {

  def mapValue[K, V](key: K): Optional[Map[K, V], V] =
    Optional[Map[K, V], V](_.get(key))(v => _ + (key -> v))

  def mapValue[K, V](key: K, default: V): Lens[Map[K, V], V] =
    Lens[Map[K, V], V](_.getOrElse(key, default))(v => _ + (key -> v))

  implicit class LensOps[S, A](lens: Lens[S, A]) {
    def inspectS[F[_]: Applicative, B](f: A => F[B]): StateT[F, S, B] = StateT.inspectF { s: S =>
      f(s.applyLens(lens).get)
    }

    def modifyS[F[_]: Applicative](f: A => F[A]): StateT[F, S, Unit] = StateT.modifyF { s: S =>
      s.applyLens(lens).modifyF(f)
    }

    def applyS[F[_]: Applicative, B](f: A => F[(A, B)]): StateT[F, S, B] = StateT { s: S =>
      val S = s.applyLens(lens)
      val a = S.get
      f(a).map(_.leftMap(S.set))
    }

    def setS[F[_]: Applicative](fa: F[A]): StateT[F, S, Unit] = modifyS(_ => fa)
    def getS[F[_]: Applicative]: StateT[F, S, A] = inspectS(_.pure[F])

  }

  implicit class OptionalOps[S, A](opt: Optional[S, A]) {
    def inspectS[F[_]: Applicative, B](f: A => F[B]): StateT[F, S, Option[B]] = StateT.inspectF { s: S =>
      s.applyOptional(opt)
        .getOption
        .traverse(f)
    }

    def tryModifyS[F[_]: Applicative](f: A => F[A]): StateT[F, S, Unit] = StateT.modifyF { s: S =>
      s.applyOptional(opt).modifyF(f)
    }

    def modifyOptS[F[_]: Applicative](f: Option[A] => F[A]): StateT[F, S, Unit] = StateT.modifyF { s: S =>
      val S = s.applyOptional(opt)
      f(S.getOption).map(S.set)
    }


    def applyS[F[_]: Applicative, B](f: A => F[(A, B)]): StateT[F, S, Option[B]] = StateT { s: S =>
      val S = s.applyOptional(opt)
      val empty: F[(S, Option[B])] = (s -> none[B]).pure[F]

      def set(a: A): F[(S, Option[B])] = f(a).map(_.leftMap(S.set))

      S.getOption.fold(empty)(set)
    }

    def applyOptionS[F[_]: Applicative, B](f: Option[A] => F[(A, B)]): StateT[F, S, B] = StateT { s: S =>
      val S = s.applyOptional(opt)
      f(S.getOption).map(_.leftMap(S.set))
    }

    def getOrElseSet[F[_]: Applicative](default: F[A]): StateT[F, S, A] = applyOptionS {
      case None => default.map(a => a -> a)
      case Some(a) => (a -> a).pure[F]
    }
  }

}