Implementing yield (yield return) using Scala continuations

Question

How might one implement C# yield return using Scala continuations? I'd like to be able to write Scala Iterators in the same style. A stab is in the comments on this Scala news post, but it doesn't work (tried using the Scala 2.8.0 beta). Answers in a related question suggest this is possible, but although I've been playing with delimited continuations for a while, I can't seem to exactly wrap my head around how to do this.

Paleface

What doesn't work about that example? It doesn't compile, or it doesn't produce the expected results? There is a mention that, for it to work, it might be necessary to have a CPS-aware foreach, but, at any rate, it would be helpful knowing what the problem is. – Binah

You may want to check Miles Sabin's answer to a similar question I had #2138119. Not sure it that gets you closer. – Grummet

There are at least two other answers to this question: #1656238 #2138119 There's also an answer to the issue of how to make it compatible with for statements: #8934726 – Horsemint

Answer 1

Before we introduce continuations we need to build some infrastructure. Below is a trampoline that operates on Iteration objects. An iteration is a computation that can either Yield a new value or it can be Done.

sealed trait Iteration[+R]
case class Yield[+R](result: R, next: () => Iteration[R]) extends Iteration[R]
case object Done extends Iteration[Nothing]

def trampoline[R](body: => Iteration[R]): Iterator[R] = {
  def loop(thunk: () => Iteration[R]): Stream[R] = {
    thunk.apply match {
      case Yield(result, next) => Stream.cons(result, loop(next))
      case Done => Stream.empty
    }
  }
  loop(() => body).iterator
}

The trampoline uses an internal loop that turns the sequence of Iteration objects into a Stream. We then get an Iterator by calling iterator on the resulting stream object. By using a Stream our evaluation is lazy; we don't evaluate our next iteration until it is needed.

The trampoline can be used to build an iterator directly.

val itr1 = trampoline {
  Yield(1, () => Yield(2, () => Yield(3, () => Done)))
}

for (i <- itr1) { println(i) }

That's pretty horrible to write, so let's use delimited continuations to create our Iteration objects automatically.

We use the shift and reset operators to break the computation up into Iterations, then use trampoline to turn the Iterations into an Iterator.

import scala.continuations._
import scala.continuations.ControlContext.{shift,reset}

def iterator[R](body: => Unit @cps[Iteration[R],Iteration[R]]): Iterator[R] =
  trampoline {
    reset[Iteration[R],Iteration[R]] { body ; Done }
  }

def yld[R](result: R): Unit @cps[Iteration[R],Iteration[R]] =
  shift((k: Unit => Iteration[R]) => Yield(result, () => k(())))

Now we can rewrite our example.

val itr2 = iterator[Int] {
  yld(1)
  yld(2)
  yld(3)
}

for (i <- itr2) { println(i) }

Much better!

Now here's an example from the C# reference page for yield that shows some more advanced usage. The types can be a bit tricky to get used to, but it all works.

def power(number: Int, exponent: Int): Iterator[Int] = iterator[Int] {
  def loop(result: Int, counter: Int): Unit @cps[Iteration[Int],Iteration[Int]] = {
    if (counter < exponent) {
      yld(result)
      loop(result * number, counter + 1)
    }
  }
  loop(number, 0)
}

for (i <- power(2, 8)) { println(i) }

Theall

I'd like to see the output of scalac -print for iterator, yld, and the assignment to itr2. Could someone with the plugin add this to the answer? – Whitney

I was just trying to apply this so I had the code running and handy. See gist.github.com/297230 for the output (scroll to the bottom). – Grummet

I'd rename iterator to yldIterator or something like that, to avoid confusion. :-) – Binah

I just saw that this was adapted into the grizzled scala library. – Binah

Grizzled removed generators in 1.1.6 as it relied "on the unsupported and unmaintained Scala continuations plugin." read here: github.com/bmc/grizzled-scala/blob/master/CHANGELOG.md – Induction

Answer 2

I managed to discover a way to do this, after a few more hours of playing around. I thought this was simpler to wrap my head around than all the other solutions I've seen thus far, though I did afterward very much appreciate Rich's and Miles' solutions.

def loopWhile(cond: =>Boolean)(body: =>(Unit @suspendable)): Unit @suspendable = {
  if (cond) {
    body
    loopWhile(cond)(body)
  }
}

  class Gen {
    var prodCont: Unit => Unit = { x: Unit => prod }
    var nextVal = 0
    def yld(i: Int) = shift { k: (Unit => Unit) => nextVal = i; prodCont = k }
    def next = { prodCont(); nextVal }
    def prod = {
      reset {
        // following is generator logic; can be refactored out generically
        var i = 0
        i += 1
        yld(i)
        i += 1
        yld(i)
        // scala continuations plugin can't handle while loops, so need own construct
        loopWhile (true) {
          i += 1
          yld(i)
        }
      }
    }
  }
  val it = new Gen
  println(it.next)
  println(it.next)
  println(it.next)

Paleface

Scala continuations can't handle while loops? Ouch! – Binah

Indeed. :( Hopefully that one's a work-in-progress, but I believe for-comprehensions are definitely not compatible with shift, as that would imply ripping apart the map/foreach/etc. – Paleface

Not any longer. Calling cps code from within the while loop is possible for some time already. For-comprehensions are still unsupported though (I don't really think they'll gain support ever) – Herrod

@PrzemekPokrywka: ...unless Scala gets her own VM, or Java 12 is released, or something. – Gilliland