In scala continuation, how to write a loop in CPS form?

Question

I'm trying to implement an example at:

https://portal.klewel.com/watch/webcast/scala-days-2019/talk/37/

using scala continuation:

object ReverseGrad_CPSImproved {

  import scala.util.continuations._

  case class Num(
      x: Double,
      var d: Double = 0.0
  ) {

    def +(that: Num) = shift { (cont: Num => Unit) =>
      val y = Num(x + that.x)

      cont(y)

      this.d += y.d
      that.d += y.d
    }

    def *(that: Num) = shift { (cont: Num => Unit) =>
      val y = Num(x * that.x)

      cont(y)

      this.d += that.x * y.d
      that.d += this.x * y.d
    }
  }

  object Num {

    implicit def fromX(x: Double): Num = Num(x)
  }

  def grad(f: Num => Num @cps[Unit])(x: Double): Double = {

    val _x = Num(x)
    reset { f(_x).d = 1.0 }

    _x.d
  }
}

This works as long as I'm using simple expression:

  it("simple") {

    val fn = { x: Num =>
      val result = (x + 3) * (x + 4)

      result
    }

    val gg = grad(fn)(3)

    println(gg)
  }

But once I started using loop it all fall apart:

  it("benchmark") {

    import scala.util.continuations._

    for (i <- 1 to 20) {

      val n = Math.pow(2, i).toInt

      val fn = { x: Num =>
        var result = x + 1

        for (j <- 2 to n) {
          result = result * (x + j)
        }

        result
      }

      val nanoFrom = System.nanoTime()
      val gg = grad(fn)(3)
      val nanoTo = System.nanoTime()

      println(s"diff = $gg,\t time = ${nanoTo - nanoFrom}")
    }
  }


[Error] /home/peng/git-spike/scalaspike/meta/src/test/scala/com/tribbloids/spike/meta/multistage/lms/ReverseGrad_CPSImproved.scala:78: found cps expression in non-cps position
one error found

I'm under the impression that the continuation library should have its own loop implementation that can be rewritten into a recursion, but I cannot find it anywhere in the latest version (scala 2.12). What's the easiest way to use loop in this case?

Answer 1

In CPS you have to rewrite your code so that you will NOT perform a nested/iterative/recursive call in the same context and instead perform just one step of the computation and pass the partial result forward.

E.g. if you wanted to calculate a product of numbers A to B you could implement it this way:

import scala.util.continuations._

case class Num(toDouble: Double) {

  def get = shift { cont: (Num => Num) =>
    cont(this)
  } 

  def +(num: Num) = reset {
    val a  = num.get
    Num(toDouble + a.toDouble)
  }

  def *(num: Num) = reset {
    val a  = num.get
    Num(toDouble * a.toDouble)
  }
}

// type annotation required because of recursive call
def product(from: Int, to: Int): Num @cps[Num] = reset { 
  if (from > to) Num(1.toDouble)
  else Num(from.toDouble) * product(from + 1, to)
}

def run: Num = reset {
  product(2, 10)
}

println(run)

(see this scastie).

The most interesting is this fragment:

reset {
  if (from > to) Num(1.toDouble)
  else Num(from.toDouble) * product(from + 1, to)
}

Here, the compiler (plugin) rewrites this to be something similar to:

input: (Num => Num) => {
  if (from > to) Num(1.toDouble)
  else {
    Num(from.toDouble) * product(from + 1, to) // this is virtually (Num => Num) => Num function!
  } (input)
}

Compiler can do this because:

• it observes the content of shift and reset calls
  • both create something that takes some parameter A and returns intermediate result B (usable in e.g. inside this or another reset) and final result C (what you get when you run the final result of the composition) (denoted as A @ cpsParam[B, C] - if B =:= C you can use a type alias A @ cps[A])
  • reset makes it easier to not go insane with passing parameters around as it handles taking parameter (A in A @ cpsParam[B, C]) and passing it to all nested CPS calls and obtaining the intermediate result (so B in A @ cpsParam[B, C]) and making whole block returning the final result - C A @ cpsParam[B, C]
  • shift lifts function (A => B) => C into A @ cpsParam[B, C]
• when it sees that the return type is Input @cpsParam[Output1, Output2] it knows that is should rewrite the code to introduce a parameter and pass it there

In practice, it s slighly more complex underneath, but that's basically it.

Meanwhile you do your

        for (j <- 2 to n) {
          result = result * (x + j)
        }

outside of this context, where compiler doesn't perform any transformations. You have to at least compose all that CPS operations within reset. (Additionally you run things in a loop and mutation which can also be delegated to CPS).

That said CPS (as in: this particular implementation) is dead. It was dropped in Scala 2.13, nobody supports it and using some trampoline-based monad (e.g. Cont from Cats) is much easier to understand, so the only places I still see it are outdated courses or articles about historical trivia.