Test for await in while (currently forbidden).

Lookup the ValDef symbol from Assign#rhs.symbol.

4 files changed, 40 insertions, 3 deletions

diff --git a/src/main/scala/scala/async/AsyncAnalysis.scala b/src/main/scala/scala/async/AsyncAnalysis.scala
index c160776..88a1bb0 100644
--- a/src/main/scala/scala/async/AsyncAnalysis.scala
+++ b/src/main/scala/scala/async/AsyncAnalysis.scala
@@ -102,10 +102,10 @@ private[async] final case class AsyncAnalysis[C <: Context](override val c: C) e
           if (isAwait(vd.rhs)) valDefsToLift += vd
         case as: Assign                                    =>
           if (isAwait(as.rhs)) {
-            assert(as.symbol != null, "internal error: null symbol for Assign tree:" + as)
+            assert(as.lhs.symbol != null, "internal error: null symbol for Assign tree:" + as +  " " + as.lhs.symbol)
 
             // TODO test the orElse case, try to remove the restriction.
-            val (vd, defBlockId) = valDefChunkId.getOrElse(as.symbol, c.abort(as.pos, "await may only be assigned to a var/val defined in the async block. " + as.symbol))
+            val (vd, defBlockId) = valDefChunkId.getOrElse(as.lhs.symbol, c.abort(as.pos, s"await may only be assigned to a var/val defined in the async block. ${as.lhs} ${as.lhs.symbol}"))
             valDefsToLift += vd
           }
           super.traverse(tree)
diff --git a/src/main/scala/scala/async/ExprBuilder.scala b/src/main/scala/scala/async/ExprBuilder.scala
index 4b81d1e..8ea7ecf 100644
--- a/src/main/scala/scala/async/ExprBuilder.scala
+++ b/src/main/scala/scala/async/ExprBuilder.scala
@@ -269,7 +269,7 @@ final case class ExprBuilder[C <: Context, FS <: FutureSystem](override val c: C
     def checkForUnsupportedAwait(tree: c.Tree) = if (tree exists {
       case Apply(fun, _) if isAwait(fun) => true
       case _                             => false
-    }) c.abort(tree.pos, "await unsupported in this position") //throw new FallbackToCpsException
+    }) c.abort(tree.pos, "await must not be used in this position") //throw new FallbackToCpsException
 
     def builderForBranch(tree: c.Tree, state: Int, nextState: Int): AsyncBlockBuilder = {
       val (branchStats, branchExpr) = tree match {
diff --git a/src/test/scala/scala/async/TreeInterrogation.scala b/src/test/scala/scala/async/TreeInterrogation.scala
index 16cbcf7..cf5948c 100644
--- a/src/test/scala/scala/async/TreeInterrogation.scala
+++ b/src/test/scala/scala/async/TreeInterrogation.scala
@@ -34,4 +34,25 @@ class TreeInterrogation {
     }
     varDefs.map(_.decoded).toSet mustBe(Set("state$async", "onCompleteHandler$async", "await$1", "await$2"))
   }
+
+
+  // @Test
+  def sandbox() {
+    val cm = reflect.runtime.currentMirror
+    val tb = mkToolbox("-cp target/scala-2.10/classes")
+    val tree = tb.parse(
+      """| import _root_.scala.async.AsyncId._
+        | async {
+        |    var x = 0
+        | var y = 0
+        | while (x <= 2) {
+        |   y = await(x)
+        |   x += 1
+        | }
+        | y
+        | }""".stripMargin)
+    val tree1 = tb.typeCheck(tree)
+
+    println(cm.universe.show(tree1))
+  }
 }
diff --git a/src/test/scala/scala/async/neg/NakedAwait.scala b/src/test/scala/scala/async/neg/NakedAwait.scala
index a0c4e4d..d400729 100644
--- a/src/test/scala/scala/async/neg/NakedAwait.scala
+++ b/src/test/scala/scala/async/neg/NakedAwait.scala
@@ -117,4 +117,20 @@ class NakedAwait {
       """.stripMargin
     }
   }
+
+  @Test
+  def whileBody() {
+    expectError("await must not be used in this position") {
+      """ import _root_.scala.async.AsyncId._
+        | async {
+        |    var x = 0
+        | var y = 0
+        | while (x <= 2) {
+        |   y = await(x)
+        |   x += 1
+        | }
+        | y
+        | }""".stripMargin
+    }
+  }
 }