https://github.com/codeport/scala/wiki/Seriously-Good-Software-%282020.7.3%29/

Speed3.scala

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

case class Container(key: String, parentKey: String, amount: Double, size: Int)

case class World(containers: Map[String, Container])
object World {
  def apply(keys: String*): World = World(keys.map { k =>
    k -> Container(k, k, 0.0, 1)
  }.toMap)
}

object Speed3 {
  def main(args: Array[String]): Unit = {
    
    import syntax._
    
    val world = World("a", "b", "c", "d")

    val setup = for {
      _ <- "a" addWater 12
      _ <- "d" addWater 8
      _ <- "a" connectTo "b"
      _ <- "b" connectTo "c"
    } yield ()

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

    println((program run world).value)
  }

}

object ops {
  def findRootAndCompress(key: String): State[World, Container] = {
    State.inspect((_: World) containers key).flatMap{ c =>
      if (c.parentKey === key) State.pure[World, Container](c)
      else findRootAndCompress(c.parentKey)
    }
  }

  def getAmount(key: String): State[World, Double] = findRootAndCompress(key).map(_.amount)

  def addWater(key: String, amount: Double): State[World, Unit] = {
    findRootAndCompress(key).flatMap{ root =>
      val newAmount = root.amount + amount / root.size
      State.modify{(w: World) =>
        World(w.containers.updated(root.parentKey, root.copy(amount = newAmount)))
      }
    }
  }

  def connectTo(thisKey: String, otherKey: String): State[World, Unit] = for {
    root1 <- findRootAndCompress(thisKey)
    root2 <- findRootAndCompress(otherKey)
    (rootKey1, rootKey2) = (root1.key, root2.key)
      
    _ <- if (rootKey1 === rootKey2) State.empty[World, Unit] else {
      val size1 = root1.size
      val size2 = root2.size
      val newSize = size1 + size2
      val newAmount = (root1.amount * size1 + root2.amount * size2) / newSize

      State.modify[World] { world => World(
        if (size1 <= size2) world.containers
          .updated(rootKey1, root1.copy(parentKey = rootKey2))
          .updated(rootKey2, root2.copy(amount = newAmount, size = newSize))
        else world.containers
          .updated(rootKey2, root2.copy(parentKey = rootKey1))
          .updated(rootKey1, root1.copy(amount = newAmount, size = newSize))
      )}
    }
  } yield ()
}

object syntax {
  implicit class ContainerKeyOps(val key: String) extends AnyVal {
    def addWater(amount: Double): State[World, Unit] = ops.addWater(key, amount)
    def connectTo(that: String): State[World, Unit] = ops.connectTo(key, that)
    def getAmount: State[World, Double] = ops.getAmount(key)
  }
}