Speed2_State.scala

package speed2

import cats.data.State

case class World(containers: Map[Symbol, Container])
case class Container(amount: Double, next: Symbol)

object Container {
  def apply(name: Symbol): Container = Container(0.0, name)
}

object Speed2 extends App {

  def getAmount(key: Symbol): State[World, Double] = updateGroup(key).inspect(_.containers(key).amount)

  def updateGroup(key: Symbol): State[World, Unit] = State.modify{ world =>
    val containers = world.containers
    @annotation.tailrec
    def findLoop(current: Symbol, target: Symbol, acc: List[Symbol]): List[Symbol] = {
      val c = containers(current)
      if (c.next == target) current :: acc
      else findLoop(c.next, target, current :: acc)
    }

    val members = findLoop(key, key, List.empty)
    val newAmount = members.map(containers(_).amount).sum / members.size
    World(members.foldLeft(containers){ (containers, member) =>
      containers.updated(member, containers(member).copy(amount = newAmount))
    })
  }

  def addWater(key: Symbol, amount: Double): State[World, Unit] = State.modify{ world =>
    val c = world.containers(key)
    World(world.containers.updated(key, c.copy(amount = c.amount + amount)))
  }

  def connetTo(key1: Symbol, key2: Symbol): State[World, Unit] = State.modify{ world =>
    val containers = world.containers
    val next1 = containers(key1).next
    val next2 = containers(key2).next
    val newNext1 = next2
    val newNext2 = next1

    World(containers
      .updated(key1, containers(key1).copy(next = newNext1))
      .updated(key2, containers(key2).copy(next = newNext2))
    )
  }



  val world = World(Map(
    'a -> Container('a),
    'b -> Container('b),
    'c -> Container('c),
    'd -> Container('d),
  ))

  val setup = for {
    _ <- addWater('a, 12)
    _ <- addWater('d, 8)
    _ <- connetTo('a, 'b)
    _ <- connetTo('b, 'c)
  } yield ()

  val program = for {
    _ <- setup
    a <- getAmount('a)
    d <- getAmount('d)
  } yield (a, d)


  println((program runA world).value)
}