June 20, 2011

After the requisite compiler gymnastics, Scala continuations tend not to play well with Scala collections. This can be frustrating when you want to use iterable methods such as `foreach`, `flatMap`, and so forth. These can be written manually, though getting the CPS types right can be a bit tricky.

With a huge thanks to Tiark Rompf for inspiring the cpsIterable implicit conversion during his done.scala presentation, I have begun the extension of this pattern to (eventually) cover many more of the collections methods we all know and love.

``````implicit def cpsIterable[A, Repr](xs: IterableLike[A, Repr]) = new {
def cps = new {
def foreach[B](f: A => Any@cpsParam[Unit, Unit]): Unit@cpsParam[Unit, Unit] = {
val it = xs.iterator
while(it.hasNext) f(it.next)
}
def map[B, That](f: A => B@cpsParam[Unit, Unit])(implicit cbf: CanBuildFrom[Repr, B, That]):
That@cpsParam[Unit, Unit] = {
val b = cbf(xs.repr)
foreach(b += f(_))
b.result
}
def flatMap[B, That](f: A => GenTraversableOnce[B]@cpsParam[Unit, Unit])
(implicit cbf: CanBuildFrom[Repr, B, That]): That@cpsParam[Unit, Unit] = {
val b = cbf(xs.repr)
for (x <- this) b ++= f(x)
b.result
}
}
}``````

To demonstrate this, I use the `shifty` method, which encapsulates a sample `shift`-based method to test within an iterable call.

``````def shifty(x: String): Int@cpsParam[Unit, Unit] = shift { k: (Int => Unit) =>
println(x)
k(x.toInt)
x.toInt
}``````

With a sample collection of values, the `cpsIterable` implicit conversion can be used to enable both `foreach` calls and for comprehensions.

``````val list = List("1", "2", "3")

def contForeach() = reset {
list.cps.foreach { x =>
shifty(x)
}
}

def contForComprehension() = reset {
println {
for(x <- list.cps) yield {
shifty(x)
}
}
}``````

Additionally, with `map` and `flatMap` implementations taken almost straight out of the Scala source code, `cpsIterable` allows these methods to straddle continuations.

``````def contMap() = reset {
list.cps.map(x => shifty(x))
()
}

def contFlatMap() = reset {
list.cps.flatMap(x => List(shifty(x)))
()
}``````

A drawback of this approach is the difficulty in transporting state outside the bounds of the `reset` delimitation. Addressing this would require more complex type annotations than the simple `@cpsParam[Unit, Unit]` used here, and would preclude the implicit conversion from being so nicely generic.