Fix detection of ill-nested awaits

 - Adapt detection logic to the post-erasure tree shapes
 - Lift restriction about || and && in favour of a
   rewrite to `If` in the ANF transform
 - import relevant parts of the test from scala-async.

Author: retronym

src/compiler/scala/tools/nsc/transform/UnCurry.scala

@@ -15,13 +15,12 @@ package tools.nsc
 package transform
 
 import scala.annotation.tailrec
-
 import symtab.Flags._
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
 import scala.reflect.internal.util.ListOfNil
-
 import PartialFunction.cond
+import scala.reflect.NameTransformer
 
 /*<export> */
 /** - uncurry all symbol and tree types (@see UnCurryPhase) -- this includes normalizing all proper types.
@@ -149,7 +148,7 @@ abstract class UnCurry extends InfoTransform
     /** Return non-local return key for given method */
     private def nonLocalReturnKey(meth: Symbol) =
       nonLocalReturnKeys.getOrElseUpdate(meth,
-        meth.newValue(unit.freshTermName("nonLocalReturnKey"), meth.pos, SYNTHETIC) setInfo ObjectTpe
+        meth.newValue(unit.freshTermName(nme.NON_LOCAL_RETURN_KEY_STRING), meth.pos, SYNTHETIC) setInfo ObjectTpe
       )
 
     /** Generate a non-local return throw with given return expression from given method.

src/compiler/scala/tools/nsc/transform/async/AnfTransform.scala

@@ -52,11 +52,13 @@ private[async] trait AnfTransform extends TransformUtils {
       tree match {
         case _: ClassDef | _: ModuleDef | _: Function | _: DefDef =>
           tree
-        case _: RefTree if tree.symbol.hasPackageFlag =>
-          tree
         case _ if !treeContainsAwait =>
           tree
-        case Apply(fun, args) if !isBooleanShortCircuit(fun.symbol) =>
+        case Apply(sel @ Select(fun, _), arg :: Nil) if isBooleanAnd(sel.symbol) && containsAwait(arg) =>
+          transform(treeCopy.If(tree, fun, arg, literalBool(false)))
+        case Apply(sel @ Select(fun, _), arg :: Nil) if isBooleanOr(sel.symbol) && containsAwait(arg) =>
+          transform(treeCopy.If(tree, fun, literalBool(true), arg))
+        case Apply(fun, args) =>
           val lastAwaitArgIndex: Int = args.lastIndexWhere(containsAwait)
           val simpleFun = transform(fun)
           var i = 0

src/compiler/scala/tools/nsc/transform/async/AsyncAnalysis.scala

@@ -13,7 +13,7 @@
 package scala.tools.nsc.transform.async
 
 import scala.collection.mutable.ListBuffer
-import scala.reflect.internal.Flags
+import scala.reflect.NameTransformer
 
 trait AsyncAnalysis extends TransformUtils  {
   import global._
@@ -35,8 +35,8 @@ trait AsyncAnalysis extends TransformUtils  {
       reportUnsupportedAwait(classDef, s"nested $kind")
     }
 
-    override def nestedModule(module: ModuleDef): Unit = {
-      reportUnsupportedAwait(module, "nested object")
+    override def nestedModuleClass(moduleClass: ClassDef): Unit = {
+      reportUnsupportedAwait(moduleClass, "nested object")
     }
 
     override def nestedMethod(defDef: DefDef): Unit = {
@@ -50,24 +50,19 @@ trait AsyncAnalysis extends TransformUtils  {
     override def function(function: Function): Unit = {
       reportUnsupportedAwait(function, "nested function")
     }
-
-    override def patMatFunction(tree: Match): Unit = {
-      reportUnsupportedAwait(tree, "nested function")
+    override def function(expandedFunction: ClassDef): Unit = {
+      reportUnsupportedAwait(expandedFunction, "nested function")
     }
 
     override def traverse(tree: Tree): Unit = {
       tree match {
         case Try(_, _, _) if containsAwait(tree)              =>
           reportUnsupportedAwait(tree, "try/catch")
           super.traverse(tree)
-        case Return(_)                                        =>
+        case Throw(Apply(fun, Ident(name) :: _)) if fun.symbol.isConstructor && fun.symbol.owner == definitions.NonLocalReturnControlClass && name.startsWith(nme.NON_LOCAL_RETURN_KEY_STRING) =>
           global.reporter.error(tree.pos, "return is illegal within a async block")
-        case DefDef(mods, _, _, _, _, _) if mods.hasFlag(Flags.LAZY) && containsAwait(tree) =>
-          reportUnsupportedAwait(tree, "lazy val initializer")
-        case ValDef(mods, _, _, _) if mods.hasFlag(Flags.LAZY) && containsAwait(tree) =>
+        case DefDef(mods, _, _, _, _, _) if tree.symbol.name.endsWith(nme.LAZY_SLOW_SUFFIX) && containsAwait(tree) =>
           reportUnsupportedAwait(tree, "lazy val initializer")
-        case CaseDef(_, guard, _) if guard exists isAwait     =>
-          reportUnsupportedAwait(tree, "pattern guard")
         case _                                                =>
           super.traverse(tree)
       }
@@ -91,7 +86,7 @@ trait AsyncAnalysis extends TransformUtils  {
       traverser(tree)
       badAwaits foreach {
         tree =>
-          reportError(tree.pos, s"await must not be used under a $whyUnsupported.")
+          reportError(tree.pos, s"${currentTransformState.Async_await.decodedName} must not be used under a $whyUnsupported.")
       }
       badAwaits.nonEmpty
     }

src/compiler/scala/tools/nsc/transform/async/LiveVariables.scala

@@ -82,15 +82,14 @@ trait LiveVariables extends ExprBuilder {
 
         override def nestedClass(classDef: ClassDef): Unit = capturingCheck(classDef)
 
-        override def nestedModule(module: ModuleDef): Unit = capturingCheck(module)
+        override def nestedModuleClass(moduleClass: ClassDef): Unit = capturingCheck(moduleClass)
 
         override def nestedMethod(defdef: DefDef): Unit = capturingCheck(defdef)
 
         override def byNameArgument(arg: Tree): Unit = capturingCheck(arg)
 
         override def function(function: Function): Unit = capturingCheck(function)
-
-        override def patMatFunction(tree: Match): Unit = capturingCheck(tree)
+        override def function(expandedFunction: ClassDef): Unit = capturingCheck(expandedFunction)
       }
 
       val findUses = new FindUseTraverser

src/compiler/scala/tools/nsc/transform/async/TransformUtils.scala

@@ -14,7 +14,6 @@ package scala.tools.nsc.transform.async
 
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
-import scala.reflect.internal.Mode
 
 /**
  * Utilities used in both `ExprBuilder` and `AnfTransform`.
@@ -36,8 +35,10 @@ private[async] trait TransformUtils extends AsyncTransformStates {
 
   def isAwait(fun: Tree): Boolean = fun.symbol == currentTransformState.Async_await
 
-  def isBooleanShortCircuit(sym: Symbol): Boolean =
-    sym.owner == definitions.BooleanClass && (sym == definitions.Boolean_and || sym == definitions.Boolean_or)
+  def isBooleanAnd(sym: Symbol): Boolean =
+    sym.owner == definitions.BooleanClass && sym == definitions.Boolean_and
+  def isBooleanOr(sym: Symbol): Boolean =
+    sym.owner == definitions.BooleanClass && sym == definitions.Boolean_or
 
   def isLabel(sym: Symbol): Boolean = sym != null && sym.isLabel
   def isCaseLabel(sym: Symbol): Boolean = sym != null && sym.isLabel && sym.name.startsWith("case")
@@ -67,6 +68,7 @@ private[async] trait TransformUtils extends AsyncTransformStates {
 
   def literalUnit: Tree = Literal(Constant(())).setType(definitions.UnitTpe) // a def to avoid sharing trees
   def literalBoxedUnit: Tree = gen.mkAttributedRef(definitions.BoxedUnit_UNIT)
+  def literalBool(b: Boolean): Tree = Literal(Constant(b)).setType(definitions.BooleanTpe)
 
   def isLiteralUnit(t: Tree): Boolean = t match {
     case Literal(Constant(())) => true
@@ -158,7 +160,7 @@ private[async] trait TransformUtils extends AsyncTransformStates {
     def nestedClass(classDef: ClassDef): Unit = {
     }
 
-    def nestedModule(module: ModuleDef): Unit = {
+    def nestedModuleClass(moduleClass: ClassDef): Unit = {
     }
 
     def nestedMethod(defdef: DefDef): Unit = {
@@ -170,20 +172,17 @@ private[async] trait TransformUtils extends AsyncTransformStates {
     def function(function: Function): Unit = {
     }
 
-    def patMatFunction(tree: Match): Unit = {
+    def function(expandedFunction: ClassDef): Unit = {
     }
 
     override def traverse(tree: Tree): Unit = {
       tree match {
-        case cd: ClassDef          => nestedClass(cd)
-        case md: ModuleDef         => nestedModule(md)
+        case cd: ClassDef          =>
+          if (cd.symbol.isAnonymousClass) function(cd)
+          else if (cd.symbol.isModuleClass) nestedModuleClass(cd)
+          else nestedClass(cd)
         case dd: DefDef            => nestedMethod(dd)
         case fun: Function         => function(fun)
-        case m@Match(EmptyTree, _) => patMatFunction(m) // Pattern matching anonymous function under -Xoldpatmat of after `restorePatternMatchingFunctions`
-        case Apply(fun, arg1 :: arg2 :: Nil) if isBooleanShortCircuit(fun.symbol) =>
-          traverse(fun)
-          traverse(arg1)
-          byNameArgument(arg2)
         case _                     => super.traverse(tree)
       }
     }

src/reflect/scala/reflect/internal/StdNames.scala

@@ -129,6 +129,8 @@ trait StdNames {
     val NESTED_IN_ANON_FUN: String             = NESTED_IN + ANON_FUN_NAME.toString.replace("$", "")
     val NESTED_IN_LAMBDA: String               = NESTED_IN + DELAMBDAFY_LAMBDA_CLASS_NAME.toString.replace("$", "")
 
+    val NON_LOCAL_RETURN_KEY_STRING: String    = "nonLocalReturnKey"
+
     /**
      * Ensures that name mangling does not accidentally make a class respond `true` to any of
      * isAnonymousClass, isAnonymousFunction, isDelambdafyFunction, e.g. by introducing "$anon".

test/async/neg/ill-nested-await.check

@@ -0,0 +1,37 @@
+ill-nested-await.scala:15: error: await must not be used under a nested method.
+    async { foo(0)(await(f(0))) }
+                   ^
+ill-nested-await.scala:20: error: await must not be used under a nested object.
+    async { object Nested { await(f(false)) } }
+                            ^
+ill-nested-await.scala:25: error: await must not be used under a nested trait.
+    async { trait Nested { await(f(false)) } }
+                           ^
+ill-nested-await.scala:30: error: await must not be used under a nested class.
+    async { class Nested { await(f(false)) } }
+                           ^
+ill-nested-await.scala:35: error: await must not be used under a nested method.
+    async { () => { await(f(false)) } }
+                    ^
+ill-nested-await.scala:40: error: await must not be used under a nested function.
+    async { { case 0 => { await(f(false)) } } : PartialFunction[Int, Boolean] }
+                          ^
+ill-nested-await.scala:45: error: await must not be used under a try/catch.
+    async { try { await(f(false)) } catch { case _: Throwable => } }
+                  ^
+ill-nested-await.scala:50: error: await must not be used under a try/catch.
+    async { try { () } catch { case _: Throwable => await(f(false)) } }
+                                                    ^
+ill-nested-await.scala:55: error: await must not be used under a try/catch.
+    async { try { () } finally { await(f(false)) } }
+                                 ^
+ill-nested-await.scala:60: error: await must not be used under a nested method.
+    async { def foo = await(f(false)) }
+                      ^
+ill-nested-await.scala:69: error: await must not be used under a lazy val initializer.
+    def foo(): Any = async { val x = { lazy val y = await(f(0)); y } }
+                                                    ^
+ill-nested-await.scala:9: error: `await` must be enclosed in an `async` block
+     await[Any](f(null))
+     ^
+12 errors found

test/async/neg/ill-nested-await.flags

@@ -0,0 +1 @@
+-Xasync

test/async/neg/ill-nested-await.scala

@@ -0,0 +1,71 @@
+import scala.concurrent._
+import ExecutionContext.Implicits.global
+import scala.tools.partest.async.Async._
+import Future.{successful => f}
+
+
+class NakedAwait {
+  def `await only allowed in async neg`(): Unit = {
+     await[Any](f(null))
+  }
+
+  def `await not allowed in by-name argument`(): Unit = {
+    // expectError("await must not be used under a by-name argument.") {
+    def foo(a: Int)(b: => Int) = 0
+    async { foo(0)(await(f(0))) }
+  }
+
+  def nestedObject(): Unit = {
+    // expectError("await must not be used under a nested object.") {
+    async { object Nested { await(f(false)) } }
+  }
+
+  def nestedTrait(): Unit = {
+    // expectError("await must not be used under a nested trait.") {
+    async { trait Nested { await(f(false)) } }
+  }
+
+  def nestedClass(): Unit = {
+    // expectError("await must not be used under a nested class.") {
+    async { class Nested { await(f(false)) } }
+  }
+
+  def nestedFunction(): Unit = {
+    // expectError("await must not be used under a nested function.") {
+    async { () => { await(f(false)) } }
+  }
+
+  def nestedPatMatFunction(): Unit = {
+    // expectError("await must not be used under a nested class.") { // TODO more specific error message
+    async { { case 0 => { await(f(false)) } } : PartialFunction[Int, Boolean] }
+  }
+
+  def tryBody(): Unit = {
+    // expectError("await must not be used under a try/catch.") {
+    async { try { await(f(false)) } catch { case _: Throwable => } }
+  }
+
+  def catchBody(): Unit = {
+    // expectError("await must not be used under a try/catch.") {
+    async { try { () } catch { case _: Throwable => await(f(false)) } }
+  }
+
+  def finallyBody(): Unit = {
+    // expectError("await must not be used under a try/catch.") {
+    async { try { () } finally { await(f(false)) } }
+  }
+
+  def nestedMethod(): Unit = {
+    // expectError("await must not be used under a nested method.") {
+    async { def foo = await(f(false)) }
+  }
+
+//  def returnIllegal(): Unit = {
+//    def foo(): Any = async { return false } //!!!
+//  }
+
+  def lazyValIllegal(): Unit = {
+    //expectError("await must not be used under a lazy val initializer")
+    def foo(): Any = async { val x = { lazy val y = await(f(0)); y } }
+  }
+}

test/junit/scala/tools/nsc/async/AnnotationDrivenAsync.scala

@@ -34,6 +34,32 @@ class AnnotationDrivenAsync {
     assertEquals(3, run(code))
   }
 
+  @Test
+  def testBooleanAndOr(): Unit = {
+    val code =
+      """import scala.concurrent._, duration.Duration, ExecutionContext.Implicits.global
+        |import scala.tools.partest.async.Async.{async, await}
+        |import Future.{successful => f}
+        |
+        |object Test {
+        |  var counter = 0
+        |  def ordered(i: Int, b: Boolean): Boolean = { assert(counter == i, (counter, i)); counter += 1; b }
+        |  def test: Future[Any] = async {
+        |    counter = 0; assert(!(ordered(0, false) && await(f(ordered(-1,  true)))))
+        |    counter = 0; assert(!(ordered(0, false) && await(f(ordered(-1, false)))))
+        |    counter = 0; assert( (ordered(0,  true) && await(f(ordered( 1,  true)))))
+        |    counter = 0; assert(!(ordered(0,  true) && await(f(ordered( 1, false)))))
+        |    counter = 0; assert( (ordered(0, false) || await(f(ordered( 1,  true)))))
+        |    counter = 0; assert(!(ordered(0, false) || await(f(ordered( 1, false)))))
+        |    counter = 0; assert( (ordered(0,  true) || await(f(ordered(-1, false)))))
+        |    counter = 0; assert( (ordered(0,  true) || await(f(ordered(-1,  true)))))
+        |    ()
+        |  }
+        |}
+        |""".stripMargin
+    assertEquals((), run(code))
+  }
+
   @Test
   def testBasicScalaConcurrentViaMacroFrontEnd(): Unit = {
     val code =