Control flows in Heaven with Scala 2.8 continuations - Part 2

My previous post in the series covered the continuation concept, the Scala 2.8 support for it and how the control flows in a program using them.

Let's have a look now at how the continuation typechecking and compilation works as this will allow us to use Scala 2.8 continuations with non-Unit types. I'm not grounding my reasoning on the formal typing rules in the paper but on my experiments and understanding so there could well be errors in the following explanation (don't hesitate to let me know!). My goal, though, is to gain and give back just an intuitive and practical understanding of how to use continuations with non-Unit types and how to help the type inferrer when it can't guess the correct types.

Let's first assume that the type of the continuation function built by the shift call and passed back to the redirect function is ContinuationFunctionType = A => B for some types A and B. Then the shift and redirect call types are the following:

• shift( redirectFunction : ContinuationFunctionType => C ) : A @cps[B, C]; please note that the redirect function can return a value of type C, i.e. belonging to a different type than the value returned by the continuation (B); also note that the return type of the shift call is A because it returns when the continuation function is invoked, and it returns precisely the continuation function's input.
  
• reset( toplevelContinuationFunction : () => B @cps[B, C] ) : C; please note that toplevelContinuationFunction returns B @cps[B, C] and not A @cps[B, C]; the B return type is because this function is the end of the continuation body (which returns a value of type B) and the C type in second position inside the annotation is because, as we've seen from the shift call type, it represents the reset return type.

The typechecker can infer the type of the shift call and will try to trace it up (towards the root) in the call tree until it finds a corresponding reset call with compatible types. If it doesn't succeed, it will complain. Sometimes manual annotations can do the trick; if they don't work it means that there's something wrong in the code about the continuation types.

As usual, if it compiles then the continuation structure of your program will run ok (and, as usual, only if you have coded your intended logic correctly you'll get the result you wanted, as no typechecker can read your mind yet).

Here's a modified (and simplified as, for now, there's only one shift call involved) version of Example1.scala in my previous post using different A, B and C types than Unit. Specifically, A = Int, B = Unit and C = String. You can see that the above typing rules are satisfied by explicit types and annotations and, indeed, the program compiles and runs fine:

import scala.continuations._
import ControlContext._

object Example2 {

private def l( msg : String ) {
  println( "[Thread '" + Thread.currentThread().getName() + "'] " + msg )
}

private def redirectFunction( cont : Int => Unit ) : String = {
  l( "Entered redirectFunction, calling continuation with 1" )

  /* 3 */

  val r : Unit = /* <-- 10 */ cont( 1 )

  /* 11 */

  l( "Exiting redirectFunction, continuation called, got result " + r + ", returning ok" )

  /* 12 */

  "ok"
}

private def packContAndUseItIfNeeded() : Int @cps[Unit, String] = {
  l( "Entered packContAndUseItIfNeeded, doing shift" )

  /* 2 */

  val r : Int @cps[Unit, String] = /* <-- 4 */ shift( redirectFunction )

  /* 5 */

  l(
    "Exiting packContAndUseItIfNeeded, shift done, " +
    "ACTUAL CONTINUATION RUNNING! Returning passed in value: " + r
  )

  /* 6 */

  r
}

private def toplevelContinuationFunction() : Unit @cps[Unit, String]= {
  l( "Entered reset and toplevelContinuationFunction, before shift" )

  /* 1 */

  val r1 : Int @cps[Unit, String] = /* <-- 7 */ packContAndUseItIfNeeded

  /* 8 */

  l(
    "In toplevelContinuationFunction, after second packContAndUseItIfNeeded, " +
    "ACTUAL CONTINUATION RUNNING!"
  )

  l( "Exiting toplevelContinuationFunction and reset, returning ()" )

  /* 9 */

  ()
}

def main( args: Array[String] ) : Unit = {
  l( "Before reset" )

  val x : String = /* <-- 13 */ reset( toplevelContinuationFunction )

  /* 14 */

  l( "After reset, terminating" )
}
}

As usual, read it and follow the flow through numbered comments, then just compile it and run it to see if logs match your understanding.

Types are not just used to check the program structure though, but to drive the cps-transform as well. In order to better understand the cps-transform, just use the plugin's compilation logging facility as in Rich Dougherty's post. Here's the result of cps-transforming the above program (with some bureocracy stripped out for easier reading):

object Example2 {
private def l(msg: String): Unit =
  println("[Thread '" + Thread.currentThread().getName() + "'] " + msg);

private def redirectFunction(cont: (Int) => Unit): String = {
  l( "Entered redirectFunction, calling continuation with 1" );
  val r: Unit = cont.apply(1);
  l( "Exiting redirectFunction, continuation called, got result " + r + ", returning ok" );
  "ok"
};

private def packContAndUseItIfNeeded(): ControlContext[Int,Unit,String] = {
  l( "Entered packContAndUseItIfNeeded, doing shift" );
  ControlContext.shiftR[Int, Unit, String]( {
    ((cont: (Int) => Unit) => redirectFunction( cont ))
  }).map[Int](((r: Int) => {
    l(
      "Exiting packContAndUseItIfNeeded, shift done, ACTUAL CONTINUATION RUNNING! " +
      "Returning passed in value: " + r
    );
    r
  }))
};

private def toplevelContinuationFunction(): ControlContext[Unit,Unit,String] = {
  l("Entered reset and toplevelContinuationFunction, before shift");
  packContAndUseItIfNeeded().map[Unit]( ((r1: Int) => {
    l(
      "In toplevelContinuationFunction, after second packContAndUseItIfNeeded, " +
      "ACTUAL CONTINUATION RUNNING!"
    );
    l( "Exiting toplevelContinuationFunction and reset, returning ()" );
    ()
  }))
};

def main(args: Array[String]): Unit = {
  l( "Before reset" );
  val x: String = ControlContext.reset[Unit, String]( toplevelContinuationFunction() );
  l( "After reset, terminating" )
}
}

1. Any expression having been assigned a type of the form A @cps[B, C] by the programmer or by the type inferrer (so not only shift calls but sub-calls containing shift calls as well) will be compiled into a new ControlContext[A, B, C] object. This object stores 2 things:
   1. The closure itself, i.e. the rest of the execution (if the shift call happens inside a sub-method, the rest of the method's body; if it happens inside a toplevel continuation function passed to a reset, up to the end of it's body);
   2. A reference to the redirect function passed to the shift call.
   The essential feature of these objects is that they exibit monadic operations, i.e. they define methods that enable them to be chained. This is essentially what happens in the transformed toplevelContinuationFunction, when the return value of packContAndUseItIfNeeded (i.e. the continuation portion contained in that method after the shift call) is chained, a map calls, with another continuation one built out of the lines in toplevelContinuationFunction after the packContAndUseItIfNeeded call. You can see that, essentially, the cps-transform is about "packaging" code in chainable continuations instead of executing it directly (a task that, incidentally, is made easier because Scala has closures). You can also see that the number of continuations built is always the least possible (and that's why the cps transform is a selective one) as, after the last call to packContAndUseItIfNeeded, there's is only one continuation built out of 2 calls to l().
   
2. reset is nothing more than a library function which call takes as an argument the chain of continuations just built, then simply starts executing the redirect function in the first continuation in the chain. The redirect function will then invoke the first continuation, which is able to hand control over to the next ones in the chain and execute up to the end of the toplevelContinuationFunction.
   

If you'd like to explore more about monads and understand how they relate to continuations, here's an amazing Monads are Elephants posts series by James Iry. You should also have a look at the definition of the cps-transform and of the Shift object in the original paper (also having a look at the type rules won't do any harm if you have the time to study them).

In the next (and last) post we'll explore a few variants of this sample that will enable you to

1. Invoke more than once shift in the toplevel continuation function;
   
2. Call the continuation on the result of another continuation call.