Handle while loops as expressions in ANF transform.

Append a `()`, as we do for `Unit` returning `if`-s and `try-s`

We don't currently support `await` in try/catch, otherwise I'd
write tests for that case, too.

2 files changed, 20 insertions, 2 deletions

diff --git a/src/main/scala/scala/async/internal/AnfTransform.scala b/src/main/scala/scala/async/internal/AnfTransform.scala
index d513499..f7f691d 100644
--- a/src/main/scala/scala/async/internal/AnfTransform.scala
+++ b/src/main/scala/scala/async/internal/AnfTransform.scala
@@ -68,6 +68,8 @@ private[async] trait AnfTransform {
 
       def _transformToList(tree: Tree): List[Tree] = trace(tree) {
         val stats :+ expr = anf.transformToList(tree)
+        def statsExprUnit =
+          stats :+ expr :+ localTyper.typedPos(expr.pos)(Literal(Constant(())))
         expr match {
           case Apply(fun, args) if isAwait(fun) =>
             val valDef = defineVal(name.await, expr, tree.pos)
@@ -77,7 +79,7 @@ private[async] trait AnfTransform {
             // if type of if-else is Unit don't introduce assignment,
             // but add Unit value to bring it into form expected by async transform
             if (expr.tpe =:= definitions.UnitTpe) {
-              stats :+ expr :+ localTyper.typedPos(expr.pos)(Literal(Constant(())))
+              statsExprUnit
             } else {
               val varDef = defineVar(name.ifRes, expr.tpe, tree.pos)
               def branchWithAssign(orig: Tree) = localTyper.typedPos(orig.pos) {
@@ -90,12 +92,14 @@ private[async] trait AnfTransform {
               val ifWithAssign = treeCopy.If(tree, cond, branchWithAssign(thenp), branchWithAssign(elsep)).setType(definitions.UnitTpe)
               stats :+ varDef :+ ifWithAssign :+ gen.mkAttributedStableRef(varDef.symbol).setType(tree.tpe).setPos(tree.pos)
             }
+          case LabelDef(name, params, rhs) =>
+            statsExprUnit
 
           case Match(scrut, cases) =>
             // if type of match is Unit don't introduce assignment,
             // but add Unit value to bring it into form expected by async transform
             if (expr.tpe =:= definitions.UnitTpe) {
-              stats :+ expr :+ localTyper.typedPos(expr.pos)(Literal(Constant(())))
+              statsExprUnit
             }
             else {
               val varDef = defineVar(name.matchRes, expr.tpe, tree.pos)
diff --git a/src/test/scala/scala/async/run/ifelse0/WhileSpec.scala b/src/test/scala/scala/async/run/ifelse0/WhileSpec.scala
index 666c373..4b3c2aa 100644
--- a/src/test/scala/scala/async/run/ifelse0/WhileSpec.scala
+++ b/src/test/scala/scala/async/run/ifelse0/WhileSpec.scala
@@ -62,4 +62,18 @@ class WhileSpec {
     }
     result mustBe (100)
   }
+
+  @Test
+  def whileExpr() {
+    import AsyncId._
+
+    val result = async {
+      var cond = true
+      while (cond) {
+        cond = false
+        await { 22 }
+      }
+    }
+    result mustBe ()
+  }
 }