Playing around with Cats' State monad

Main2.scala

package testing

import cats._
import cats.data._
import cats.implicits._

object Main2 extends App {

  case class Stats(warning: Int, error: Int)

  sealed trait Level
  case object Okay extends Level
  case object Notifiable extends Level
  case object Beware extends Level
  case object Deadly extends Level

  def classify(value: Int): State[Stats, Level] = State[Stats, Level]((stats: Stats) => {
    val level = if (value <= 5) Okay
    else if (value <= 10) Notifiable
    else if (value <= 20) Beware
    else Deadly

    val ns = if (level == Beware) stats.copy(warning = stats.warning + 1)
    else if (level == Deadly) stats.copy(error = stats.error + 1)
    else stats

    (ns, level)
  })

  case class Counter(times: Int)
  def count(value: Int): State[Counter, Unit] = State[Counter, Unit]((counter: Counter) => (Counter(counter.times + 1), ()))

  val list = (1 to 25).toList

  val res1: State[Stats, List[Level]] = list.traverse[State[Stats, ?], Level](classify)
  val r1 = res1.run(Stats(0, 0)).value
  println(r1)

  val res2 = list.traverse[State[Counter, ?], Unit](count)
  val r2 = res2.run(Counter(0)).value
  println(r2)

  case class BigState(stats: Stats, counter: Counter)

  def extStats[A](st: State[Stats, A]): State[BigState, A] = st.transformS[BigState](_.stats, (bs, st) => bs.copy(stats = st))
  def extCounter[A](st: State[Counter, A]): State[BigState, A] = st.transformS[BigState](_.counter, (bs, co) => bs.copy(counter = co))

  def classifyAndCount(value: Int): State[BigState, Level] = for {
    level <- extStats(classify(value))
    _ <- extCounter(count(value))
  } yield level

  val res3 = list.traverse[State[BigState, ?], Level](classifyAndCount)
  val r3 = res3.run(BigState(Stats(0, 0), Counter(0))).value
  println(r3)
  
}