Seriously Good Software Ch 3 Exercise 3 (with cats and shapeless)

Exercise3.scala

import cats.data.StateT
import cats.implicits._
import typedef._

object Main extends App {
  
  import syntax._
  
  val tv = ApplianceId("tv")
  val radio = ApplianceId("radio")
  val grid = GridId("grid")
  
  val initialWorld = World(Map.empty, Map.empty)
  
  val setup: WorldState[Unit] = for {
    _ <- ops.newAppliance(tv, 150)
    _ <- ops.newAppliance(radio, 30)
    _ <- ops.newGrid(grid, 3000)
  } yield ()

  val program: WorldState[Unit] = for {
    _ <- setup
    _ <- tv plugInto grid
    _ <- radio plugInto grid
    g1 <- grid.get
    _ <- tv.on
    g2 <- grid.get
    _ <- radio.on
    g3 <- grid.get
  } yield {
    Seq(g1, g2, g3).map(_.residualPower) foreach println
  }
  
  println(program runS initialWorld)
}

object typedef {
  
  import shapeless.tag
  import shapeless.tag.@@
  
  trait ApplianceIdTag
  trait GridIdTag
  type ApplianceId = String @@ ApplianceIdTag
  type GridId = String @@ GridIdTag
  
  object ApplianceId {
    def apply(id: String): ApplianceId = tag[ApplianceIdTag][String](id)
  }
  
  object GridId {
    def apply(id: String): GridId = tag[GridIdTag][String](id)
  }
  
  type EitherS[A] = Either[String, A]
  type WorldState[A] = StateT[EitherS, World, A]
  
  case class World(grids: Map[GridId, Grid], appliances: Map[ApplianceId, Appliance])

  case class Appliance(powerAbsorbed: Int, gridId: Option[GridId], isOn: Boolean)

  case class Grid(maxPower: Int, residualPower: Int)

}

object syntax {
  
  implicit class ApplianceSyntax(val applianceId: ApplianceId) extends AnyVal {
    def plugInto(gridId: GridId): WorldState[Unit] = ops.plugInto(applianceId, gridId)
    def on: WorldState[Unit] = ops.on(applianceId)
    def off: WorldState[Unit] = ops.off(applianceId)
    def get: WorldState[Appliance] = ops.getAppliance(applianceId)
  }
  
  implicit class GridSynttax(val gridId: GridId) extends AnyVal {
    def get: WorldState[Grid] = ops.getGrid(gridId)
    def addPower(power: Int): WorldState[Unit] = ops.addPower(gridId, power)
  }
}

object ops {
  
  def updateAppliance(id: ApplianceId, appliance: Appliance): WorldState[Unit] = StateT.modify{ world =>
    world.copy(appliances = world.appliances.updated(id, appliance))
  }
  
  def updateGrid(id: GridId, grid: Grid): WorldState[Unit] = StateT.modify{ world =>
    world.copy(grids = world.grids.updated(id, grid))
  }
  
  def newAppliance(id: ApplianceId, powerAbsorbed: Int): WorldState[Unit] =
    updateAppliance(id, Appliance(powerAbsorbed, None, isOn = false))

  def newGrid(id: GridId, maxPower: Int): WorldState[Unit] =
    updateGrid(id, Grid(maxPower, maxPower))
  
  def getAppliance(id: ApplianceId): WorldState[Appliance] = StateT.inspectF{ world =>
    world.appliances.get(id).toRight(s"Appliance is not exist with id: $id")
  }
  
  def getGrid(id: GridId): WorldState[Grid] = StateT.inspectF{ world =>
    world.grids.get(id).toRight(s"Grid is not exist with id: $id")
  }
  
  def addPower(gridId: GridId, power: Int): WorldState[Unit] = getGrid(gridId).flatMap{ grid =>
    val newPower = grid.residualPower + power
    
    if (newPower < 0) StateT.setF[EitherS, World](Left("Not enough power."))
    else if (newPower > grid.maxPower) StateT.setF[EitherS, World](Left("Maximum power exceeded."))
    else updateGrid(gridId, grid.copy(residualPower = newPower))
  }
  
  def plugInto(applianceId: ApplianceId, gridId: GridId): WorldState[Unit] = for {
    appliance <- getAppliance(applianceId)
    _ <- appliance.gridId match {
      case Some(oldGridId) if appliance.isOn => addPower(oldGridId, appliance.powerAbsorbed)
      case _ => StateT.empty[EitherS, World, Unit]
    }
    _ <- if (appliance.isOn) addPower(gridId, -appliance.powerAbsorbed) else StateT.empty[EitherS, World, Unit]
    _ <- updateAppliance(applianceId, appliance.copy(gridId = Some(gridId)))
  } yield ()

  def on(applianceId: ApplianceId): WorldState[Unit] = getAppliance(applianceId).flatMap{ appliance =>
    appliance.gridId match {
      case None =>
        StateT.liftF[EitherS, World, Unit](Left("Cannot turn on when unconnected to any grid."))
      case Some(oldGridId) if !appliance.isOn => for {
        _ <- addPower(oldGridId, -appliance.powerAbsorbed)
        _ <- updateAppliance(applianceId, appliance.copy(isOn = true))
      } yield ()
      case Some(_) => StateT.empty[EitherS, World, Unit]
    }
  }

  def off(applianceId: ApplianceId): WorldState[Unit] = getAppliance(applianceId).flatMap{ appliance =>
    appliance.gridId match {
      case Some(oldGridId) if appliance.isOn => for {
        _ <- addPower(oldGridId, appliance.powerAbsorbed)
        _ <- updateAppliance(applianceId, appliance.copy(isOn = false))
      } yield ()
      case _ => StateT.empty[EitherS, World, Unit]
    }
  }
}