Fix crashers in do/while and while(await(..))

src/main/scala/scala/async/internal/AnfTransform.scala

@@ -162,7 +162,14 @@ private[async] trait AnfTransform {
       def _transformToList(tree: Tree): List[Tree] = trace(tree) {
         val containsAwait = tree exists isAwait
         if (!containsAwait) {
-          List(tree)
+          tree match {
+            case Block(stats, expr) =>
+              // avoids nested block in `while(await(false)) ...`.
+              // TODO I think `containsAwait` really should return true if the code contains a label jump to an enclosing
+              // while/doWhile and there is an await *anywhere* inside that construct.
+              stats :+ expr
+            case _ => List(tree)
+          }
         } else tree match {
           case Select(qual, sel) =>
             val stats :+ expr = linearize.transformToList(qual)

src/main/scala/scala/async/internal/ExprBuilder.scala

@@ -127,7 +127,11 @@ trait ExprBuilder {
     private var nextJumpState: Option[Int] = None
 
     def +=(stat: Tree): this.type = {
-      assert(nextJumpState.isEmpty, s"statement appeared after a label jump: $stat")
+      stat match {
+        case Literal(Constant(())) => // This case occurs in do/while
+        case _ =>
+          assert(nextJumpState.isEmpty, s"statement appeared after a label jump: $stat")
+      }
       def addStat() = stats += stat
       stat match {
         case Apply(fun, Nil) =>
@@ -228,7 +232,7 @@ trait ExprBuilder {
         currState = afterAwaitState
         stateBuilder = new AsyncStateBuilder(currState, symLookup)
 
-      case If(cond, thenp, elsep) if stat exists isAwait =>
+      case If(cond, thenp, elsep) if (stat exists isAwait) || containsForiegnLabelJump(stat) =>
         checkForUnsupportedAwait(cond)
 
         val thenStartState = nextState()

src/main/scala/scala/async/internal/TransformUtils.scala

@@ -96,6 +96,19 @@ private[async] trait TransformUtils {
     treeInfo.isExprSafeToInline(tree)
   }
 
+  // `while(await(x))` ... or `do { await(x); ... } while(...)` contain an `If` that loops;
+  // we must break that `If` into states so that it convert the label jump into a state machine
+  // transition
+  final def containsForiegnLabelJump(t: Tree): Boolean = {
+    val labelDefs = t.collect {
+      case ld: LabelDef => ld.symbol
+    }.toSet
+    t.exists {
+      case rt: RefTree => !(labelDefs contains rt.symbol)
+      case _ => false
+    }
+  }
+
   /** Map a list of arguments to:
     * - A list of argument Trees
     * - A list of auxillary results.

src/test/scala/scala/async/TreeInterrogation.scala

@@ -66,43 +66,18 @@ object TreeInterrogation extends App {
   withDebug {
     val cm = reflect.runtime.currentMirror
     val tb = mkToolbox("-cp ${toolboxClasspath} -Xprint:typer -uniqid")
-    import scala.async.internal.AsyncTestLV._
+    import scala.async.internal.AsyncId._
     val tree = tb.parse(
       """
-        | import scala.async.internal.AsyncTestLV._
-        | import scala.async.internal.AsyncTestLV
-        |
-        | case class MCell[T](var v: T)
-        | val f = async { MCell(1) }
-        |
-        | def m1(x: MCell[Int], y: Int): Int =
-        |   async { x.v + y }
-        | case class Cell[T](v: T)
-        |
+        | import scala.async.internal.AsyncId._
         | async {
-        |   // state #1
-        |   val a: MCell[Int] = await(f)     // await$13$1
-        |   // state #2
-        |   var y = MCell(0)
-        |
-        |   while (a.v < 10) {
-        |     // state #4
-        |     a.v = a.v + 1
-        |     y = MCell(await(a).v + 1)      // await$14$1
-        |     // state #7
+        |   var b = true
+        |   while(await(b)) {
+        |     b = false
         |   }
-        |
-        |   // state #3
-        |   assert(AsyncTestLV.log.exists(entry => entry._1 == "await$14$1"))
-        |
-        |   val b = await(m1(a, y.v))        // await$15$1
-        |   // state #8
-        |   assert(AsyncTestLV.log.exists(_ == ("a$1" -> MCell(10))))
-        |   assert(AsyncTestLV.log.exists(_ == ("y$1" -> MCell(11))))
-        |   b
+        |   await(b)
         | }
         |
-        |
         | """.stripMargin)
     println(tree)
     val tree1 = tb.typeCheck(tree.duplicate)

src/test/scala/scala/async/run/ifelse0/WhileSpec.scala

@@ -76,4 +76,44 @@ class WhileSpec {
     }
     result mustBe ()
   }
+
+  @Test def doWhile() {
+    import AsyncId._
+    val result = async {
+      var b = 0
+      var x = ""
+      await(do {
+        x += "1"
+        x += await("2")
+        x += "3"
+        b += await(1)
+      } while (b < 2))
+      await(x)
+    }
+    result mustBe "123123"
+  }
+
+  @Test def whileAwaitCondition() {
+    import AsyncId._
+    val result = async {
+      var b = true
+      while(await(b)) {
+        b = false
+      }
+      await(b)
+    }
+    result mustBe false
+  }
+
+  @Test def doWhileAwaitCondition() {
+    import AsyncId._
+    val result = async {
+      var b = true
+      do {
+        b = false
+      } while(await(b))
+      b
+    }
+    result mustBe false
+  }
 }