Small cleanups after code review

 - More internal docs
 - Be more frugal with the `NoAwait` attachment, for some AST node
   types this is implied.
 - Just use `x`, rather than what was effectively
   `x.reverseMap(identity).reverse`

3 files changed, 40 insertions, 12 deletions

diff --git a/src/main/scala/scala/async/internal/AnfTransform.scala b/src/main/scala/scala/async/internal/AnfTransform.scala
index cc77ec7..4545ca6 100644
--- a/src/main/scala/scala/async/internal/AnfTransform.scala
+++ b/src/main/scala/scala/async/internal/AnfTransform.scala
@@ -220,7 +220,7 @@ private[async] trait AnfTransform {
 
             case Block(stats, expr)                                    =>
               val trees = stats.flatMap(linearize.transformToList).filterNot(isLiteralUnit) ::: linearize.transformToList(expr)
-              eliminateLabelParameters(trees)
+              eliminateMatchEndLabelParameter(trees)
 
             case ValDef(mods, name, tpt, rhs) =>
               if (containsAwait(rhs)) {
@@ -273,7 +273,15 @@ private[async] trait AnfTransform {
       }
 
       // Replace the label parameters on `matchEnd` with use of a `matchRes` temporary variable
-      def eliminateLabelParameters(statsExpr: List[Tree]): List[Tree] = {
+      //
+      // CaseDefs are translated to labels without parmeters. A terminal label, `matchEnd`, accepts
+      // a parameter which is the result of the match (this is regular, so even Unit-typed matches have this).
+      //
+      // For our purposes, it is easier to:
+      //   - extract a `matchRes` variable
+      //   - rewrite the terminal label def to take no parameters, and instead read this temp variable
+      //   - change jumps to the terminal label to an assignment and a no-arg label application
+      def eliminateMatchEndLabelParameter(statsExpr: List[Tree]): List[Tree] = {
         import internal.{methodType, setInfo}
         val caseDefToMatchResult = collection.mutable.Map[Symbol, Symbol]()
 
@@ -304,7 +312,7 @@ private[async] trait AnfTransform {
             )
         }
         matchResults.toList match {
-          case Nil => statsExpr0.reverse
+          case Nil => statsExpr
           case r1 :: Nil => (r1 +: statsExpr0.reverse) :+ atPos(tree.pos)(gen.mkAttributedIdent(r1.symbol))
           case _ => c.error(macroPos, "Internal error: unexpected tree encountered during ANF transform " + statsExpr); statsExpr
         }
diff --git a/src/main/scala/scala/async/internal/AsyncTransform.scala b/src/main/scala/scala/async/internal/AsyncTransform.scala
index f491403..af290e4 100644
--- a/src/main/scala/scala/async/internal/AsyncTransform.scala
+++ b/src/main/scala/scala/async/internal/AsyncTransform.scala
@@ -98,10 +98,11 @@ trait AsyncTransform {
     }
 
     val isSimple = asyncBlock.asyncStates.size == 1
-    if (isSimple)
+    val result = if (isSimple)
       futureSystemOps.spawn(body, execContext) // generate lean code for the simple case of `async { 1 + 1 }`
     else
       startStateMachine
+    cleanupContainsAwaitAttachments(result)
   }
 
   def logDiagnostics(anfTree: Tree, states: Seq[String]) {
diff --git a/src/main/scala/scala/async/internal/TransformUtils.scala b/src/main/scala/scala/async/internal/TransformUtils.scala
index ed8b103..90419d3 100644
--- a/src/main/scala/scala/async/internal/TransformUtils.scala
+++ b/src/main/scala/scala/async/internal/TransformUtils.scala
@@ -396,13 +396,18 @@ private[async] trait TransformUtils {
    * in search of a sub tree that was decorated with the cached answer.
    */
   final def containsAwaitCached(t: Tree): Tree => Boolean = {
+    def treeCannotContainAwait(t: Tree) = t match {
+      case _: Ident | _: TypeTree | _: Literal => true
+      case _ => false
+    }
+    def shouldAttach(t: Tree) = !treeCannotContainAwait(t)
     val symtab = c.universe.asInstanceOf[scala.reflect.internal.SymbolTable]
-    def attachContainsAwait(t: Tree): Unit = {
+    def attachContainsAwait(t: Tree): Unit = if (shouldAttach(t)) {
       val t1 = t.asInstanceOf[symtab.Tree]
       t1.updateAttachment(ContainsAwait)
       t1.removeAttachment[NoAwait.type]
     }
-    def attachNoAwait(t: Tree): Unit = {
+    def attachNoAwait(t: Tree): Unit = if (shouldAttach(t)) {
       val t1 = t.asInstanceOf[symtab.Tree]
       t1.updateAttachment(NoAwait)
     }
@@ -423,15 +428,16 @@ private[async] trait TransformUtils {
     markContainsAwaitTraverser.traverse(t)
 
     (t: Tree) => {
-      val symtab = c.universe.asInstanceOf[scala.reflect.internal.SymbolTable]
       object traverser extends Traverser {
         var containsAwait = false
         override def traverse(tree: Tree): Unit = {
-          if (tree.asInstanceOf[symtab.Tree].hasAttachment[NoAwait.type])
-            ()
-          else if (tree.asInstanceOf[symtab.Tree].hasAttachment[ContainsAwait.type])
-            containsAwait = true
-          else super.traverse(tree)
+          def castTree = tree.asInstanceOf[symtab.Tree]
+          if (!castTree.hasAttachment[NoAwait.type]) {
+            if (castTree.hasAttachment[ContainsAwait.type])
+              containsAwait = true
+            else if (!treeCannotContainAwait(t))
+              super.traverse(tree)
+          }
         }
       }
       traverser.traverse(t)
@@ -439,6 +445,19 @@ private[async] trait TransformUtils {
     }
   }
 
+  final def cleanupContainsAwaitAttachments(t: Tree): t.type = {
+    val symtab = c.universe.asInstanceOf[scala.reflect.internal.SymbolTable]
+    t.foreach {t =>
+      t.asInstanceOf[symtab.Tree].removeAttachment[ContainsAwait.type]
+      t.asInstanceOf[symtab.Tree].removeAttachment[NoAwait.type]
+    }
+    t
+  }
+
+  // First modification to translated patterns:
+  //  - Set the type of label jumps to `Unit`
+  //  - Propagate this change to trees known to directly enclose them:
+  //    ``If` / `Block`) adjust types of enclosing
   final def adjustTypeOfTranslatedPatternMatches(t: Tree, owner: Symbol): Tree = {
     import definitions.UnitTpe
     typingTransform(t, owner) {