How to implement cost estimation for deferred computations with `flatMap`? - scala

Question

I have implemented a Calculation class that takes 2 parameters: a calculation input that is a call-by-name param and also a cost. When I try to flatMap calculations, the first part of it gets executed. Is it possible to defer everything in flatMap and still provide the total cost?
class Calculation[+R](input: => R, val cost: Int = 0) {
def value: R = input
def map[A](f: R => A): Calculation[A] =
new Calculation(f(input), cost)
def flatMap[A](f: R => Calculation[A]): Calculation[A] = {
val step = f(input)
new Calculation(step.value, cost + step.cost)
}
}
object Rextester extends App {
val f1 = new Calculation({
println("F1")
"F1"
})
val f2 = f1.flatMap(s => new Calculation({
println("F2")
s + " => F2"
}))
println(f2.cost)
}
Once f2 is declared (flatMap is called), we can see that "F1" will be printed. The printed cost is "15", which is correct, but I would like to have the actual calculation fully deferred, meaning that I shouldn't see the f1 being executed when I calculate the cost.

Answer 1

You just need a little more laziness, so that the cost isn't eagerly evaluated in the flatMap:
class Calculation[+R](input: => R, c: => Int = 0) {
def value: R = input
lazy val cost: Int = c
def map[A](f: R => A): Calculation[A] =
new Calculation(f(input), cost)
def flatMap[A](f: R => Calculation[A]): Calculation[A] = {
lazy val step = f(value)
new Calculation(step.value, cost + step.cost)
}
}
Note that this still might not have exactly the semantics you want (e.g. calling f2.value twice in a row will results in both F1 and F2 being printed the first time, and only F2 the second), but it does keep the side effect from occurring when f2 is defined.

Answer 2

If I understand your requirement
defer everything in flatMap and still provide the total cost
correctly, then you want to compute an estimate for the total costs before making any computations. I don't see how this is supposed to work with the signature flatMap[A](f: R => Calculation[A]): Calculation[A] - your cost is attached to Calculation[A], and your Calculation[A] depends on a concrete instance of R, so you cannot compute the cost before computing R.
Constant costs for computation steps
Here is a completely different proposal:
sealed trait Step[-A, +B] extends (A => B) { outer =>
def estimatedCosts: Int
def andThen[U >: B, C](next: Step[U, C]): Step[A, C] = new Step[A, C] {
def apply(a: A): C = next(outer(a))
def estimatedCosts = outer.estimatedCosts + next.estimatedCosts
}
def result(implicit u_is_a: Unit <:< A): B = this(u_is_a(()))
}
type Computation[+R] = Step[Unit, R]
The trait Step represents a computation step, for which the costs do not depend on the input. It's essentially just a Function[A, B] with an integer value attached to it. Your Computation[R] then becomes a special case, namely Step[Unit, R].
Here is how you can use it:
val x = new Step[Unit, Int] {
def apply(_u: Unit) = 42
def estimatedCosts = 0
}
val mul = new Step[Int, Int] {
def apply(i: Int) = {
println("<computing> adding is easy")
i + 58
}
def estimatedCosts = 10
}
val sqrt = new Step[Int, Double] {
def apply(i: Int) = {
println("<computing> finding square roots is difficult")
math.sqrt(i)
}
def estimatedCosts = 50
}
val c: Computation[Double] = x andThen mul andThen sqrt
println("Estimated costs: " + c.estimatedCosts)
println("(nothing computed so far)")
println(c.result)
If you run it, you obtain:
Estimated costs: 60
(nothing computed so far)
<computing> adding is easy
<computing> finding square roots is difficult
10.0
What it does is the following:
It starts with value 42, adds 58 to it, and then computes the square root of the sum
Addition is set to cost 10 units, square root costs 50.
It gives you the cost estimate of 60 units, without performing any computations.
Only when you invoke .result does it compute the actual result 10.0
Admittedly, it's not very useful for anything except very coarse order-of-magnitude estimates. It's so coarse that even using Ints barely makes any sense.
Non-constant costs per step
You can make your cost estimates more accurate by keeping track of a size estimate as follows:
trait Step[-A, +B] extends (A => B) {
def outputSizeEstimate(inputSizeEstimate: Int): Int
def costs(inputSizeEstimate: Int): Int
}
trait Computation[+R] { outer =>
def result: R
def resultSize: Int
def estimatedCosts: Int
def map[S](step: Step[R, S]): Computation[S] = new Computation[S] {
def result: S = step(outer.result)
def estimatedCosts: Int = outer.estimatedCosts + step.costs(outer.resultSize)
def resultSize: Int = step.outputSizeEstimate(outer.resultSize)
}
}
val x = new Computation[List[Int]] {
def result = (0 to 10).toList
def resultSize = 10
def estimatedCosts = 10
}
val incrementEach = new Step[List[Int], List[Int]] {
def outputSizeEstimate(inputSize: Int) = inputSize
def apply(xs: List[Int]) = {
println("incrementing...")
xs.map(1.+)
}
def costs(inputSize: Int) = 3 * inputSize
}
val timesSelf = new Step[List[Int], List[(Int, Int)]] {
def outputSizeEstimate(n: Int) = n * n
def apply(xs: List[Int]) = {
println("^2...")
for (x <- xs; y <- xs) yield (x, y)
}
def costs(n: Int) = 5 * n * n
}
val addPairs = new Step[List[(Int, Int)], List[Int]] {
def outputSizeEstimate(n: Int) = n
def apply(xs: List[(Int, Int)]) = {
println("adding...")
xs.map{ case (a, b) => a + b }
}
def costs(n: Int) = 7 * n
}
val y = x map incrementEach map timesSelf map addPairs
println("Estimated costs (manually): " + (10 + 30 + 500 + 700))
println("Estimated costs (automatically): " + y.estimatedCosts)
println("(nothing computed so far)")
println(y.result)
The output looks encouraging:
Estimated costs (manually): 1240
Estimated costs (automatically): 1240
(nothing computed so far)
incrementing...
^2...
adding...
List(2, 3, 4, 5, 6, 7, 8, 9, 10, 11, ...[omitted]..., 20, 21, 22)
Note that the approach is not restricted to lists and integers: the size estimates can be arbitrarily complicated. For example, they could be dimensions of matrices or tensors. Actually, they don't have to be sizes at all. Those estimates could just as well contain any other kind of "static conservative estimates", like types or logical predicates.
Non-constant costs, using Writer
Using the Writer monad from Cats, we can express the same idea more succinctly by replacing the two methods outputSizeEstimate and costs on a Step by a single method that takes an Int and returns a Writer[Int, Int]:
Writers .value corresponds to size estimate for the output
Writers .written corresponds to costs of the step (which might depend on the input size)
Full code:
import cats.data.Writer
import cats.syntax.writer._
import cats.instances.int._
object EstimatingCosts extends App {
type Costs = Int
type Size = Int
trait Step[-A, +B] extends (A => B) {
def sizeWithCosts(inputSizeEstimate: Size): Writer[Costs, Size]
}
object Step {
def apply[A, B]
(sizeCosts: Size => (Size, Costs))
(mapResult: A => B)
: Step[A, B] = new Step[A, B] {
def apply(a: A) = mapResult(a)
def sizeWithCosts(s: Size) = { val (s2, c) = sizeCosts(s); Writer(c, s2) }
}
}
trait Computation[+R] { outer =>
def result: R
def sizeWithCosts: Writer[Costs, Size]
def size: Size = sizeWithCosts.value
def costs: Costs = sizeWithCosts.written
def map[S](step: Step[R, S]): Computation[S] = new Computation[S] {
lazy val result: S = step(outer.result)
lazy val sizeWithCosts = outer.sizeWithCosts.flatMap(step.sizeWithCosts)
}
}
object Computation {
def apply[A](initialSize: Size, initialCosts: Costs)(a: => A) = {
new Computation[A] {
lazy val result = a
lazy val sizeWithCosts = Writer(initialCosts, initialSize)
}
}
}
val x = Computation(10, 10){ (0 to 10).toList }
val incrementEach = Step(n => (n, 3 * n)){ (xs: List[Int]) =>
println("incrementing...")
xs.map(1.+)
}
val timesSelf = Step(n => (n * n, 5 * n * n)) { (xs: List[Int]) =>
println("^2...")
for (x <- xs; y <- xs) yield (x, y)
}
val addPairs = Step(n => (n, 7 * n)) { (xs: List[(Int, Int)]) =>
println("adding...")
xs.map{ case (a, b) => a + b }
}
val y = x map incrementEach map timesSelf map addPairs
println("Estimated costs (manually): " + (10 + 30 + 500 + 700))
println("Estimated costs (automatically): " + y.costs)
println("(nothing computed so far)")
println(y.result)
}
The output stays exactly the same as in the previous section.
PS: I think I came up with a more concise way to summarize this entire answer:
Use the product category of the ordinary ambient Scala category (types and functions) with the monoid of endomorphisms on object Int in the Kleisli category of Writer[Int, ?].
In some hypothetical language, the answer might have been:
Use Sc * End{Kl(Writer[Int, ?])}[Int].

Answer 3

First of all there is no reason to re-invent your own Functor and FlatMap and I would strongly advice you to use existing implementation.
If you need deffered computation than cats.Writer[Int, ?] is your friend.
With its support you can write your cost as well as get functor and monad instances.
Let me give you an example. We start with some initial cost
val w = Writer.put("F1")(0)
w.flatMap(v => Writer.value(v + "F2"))