Cleanup code generation by avoiding redundant blocks

Don't bother adding `{ ...; () }` if we can use the original tree(s) instead,
e.g. if the last tree in `...` conforms to `Unit`.

This makes the debug output of the macro a little easier to read.

2 files changed, 44 insertions, 19 deletions

diff --git a/src/main/scala/scala/async/internal/AsyncTransform.scala b/src/main/scala/scala/async/internal/AsyncTransform.scala
index ec3a2a1..baa3fc2 100644
--- a/src/main/scala/scala/async/internal/AsyncTransform.scala
+++ b/src/main/scala/scala/async/internal/AsyncTransform.scala
@@ -70,12 +70,11 @@ trait AsyncTransform {
     for ((state, flds) <- assignsOf) {
       val assigns = flds.map { fld =>
         val fieldSym = fld.symbol
-        Block(
-          List(
-            asyncBase.nullOut(c.universe)(c.Expr[String](Literal(Constant(fieldSym.name.toString))), c.Expr[Any](Ident(fieldSym))).tree
-          ),
-          Assign(gen.mkAttributedStableRef(thisType(fieldSym.owner), fieldSym), mkZero(fieldSym.info))
-        )
+        val assign = Assign(gen.mkAttributedStableRef(thisType(fieldSym.owner), fieldSym), mkZero(fieldSym.info))
+        asyncBase.nullOut(c.universe)(c.Expr[String](Literal(Constant(fieldSym.name.toString))), c.Expr[Any](Ident(fieldSym))).tree match {
+          case Literal(Constant(value: Unit)) => assign
+          case x => Block(x :: Nil, assign)
+        }
       }
       val asyncState = asyncBlock.asyncStates.find(_.state == state).get
       asyncState.stats = assigns ++ asyncState.stats
diff --git a/src/main/scala/scala/async/internal/ExprBuilder.scala b/src/main/scala/scala/async/internal/ExprBuilder.scala
index 4e521c9..f24e37e 100644
--- a/src/main/scala/scala/async/internal/ExprBuilder.scala
+++ b/src/main/scala/scala/async/internal/ExprBuilder.scala
@@ -34,8 +34,13 @@ trait ExprBuilder {
 
     var stats: List[Tree]
 
+    def statsAnd(trees: List[Tree]): List[Tree] = {
+      val body = adaptToUnit(stats)
+      Try(body, Nil, adaptToUnit(trees)) :: Nil
+    }
+
     final def allStats: List[Tree] = this match {
-      case a: AsyncStateWithAwait => stats :+ a.awaitable.resultValDef
+      case a: AsyncStateWithAwait => statsAnd(a.awaitable.resultValDef :: Nil)
       case _ => stats
     }
 
@@ -52,8 +57,9 @@ trait ExprBuilder {
     def nextStates: List[Int] =
       List(nextState)
 
-    def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef =
-      mkHandlerCase(state, stats :+ mkStateTree(nextState, symLookup))
+    def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef = {
+      mkHandlerCase(state, statsAnd(mkStateTree(nextState, symLookup) :: Nil))
+    }
 
     override val toString: String =
       s"AsyncState #$state, next = $nextState"
@@ -87,10 +93,10 @@ trait ExprBuilder {
       val tryGetOrCallOnComplete =
         if (futureSystemOps.continueCompletedFutureOnSameThread)
           If(futureSystemOps.isCompleted(c.Expr[futureSystem.Fut[_]](awaitable.expr)).tree,
-            Block(ifIsFailureTree[T](futureSystemOps.getCompleted[Any](c.Expr[futureSystem.Fut[Any]](awaitable.expr)).tree) :: Nil, literalUnit),
-            Block(callOnComplete :: Nil, Return(literalUnit)))
+            adaptToUnit(ifIsFailureTree[T](futureSystemOps.getCompleted[Any](c.Expr[futureSystem.Fut[Any]](awaitable.expr)).tree) :: Nil),
+            Block(toList(callOnComplete), Return(literalUnit)))
         else
-          Block(callOnComplete :: Nil, Return(literalUnit))
+          Block(toList(callOnComplete), Return(literalUnit))
       mkHandlerCase(state, stats ++ List(mkStateTree(onCompleteState, symLookup), tryGetOrCallOnComplete))
     }
 
@@ -110,11 +116,11 @@ trait ExprBuilder {
      */
     def ifIsFailureTree[T: WeakTypeTag](tryReference: => Tree) =
       If(futureSystemOps.tryyIsFailure(c.Expr[futureSystem.Tryy[T]](tryReference)).tree,
-        Block(futureSystemOps.completeProm[T](
+        Block(toList(futureSystemOps.completeProm[T](
           c.Expr[futureSystem.Prom[T]](symLookup.memberRef(name.result)),
           c.Expr[futureSystem.Tryy[T]](
             TypeApply(Select(tryReference, newTermName("asInstanceOf")),
-              List(TypeTree(futureSystemOps.tryType[T]))))).tree :: Nil,
+              List(TypeTree(futureSystemOps.tryType[T]))))).tree),
           Return(literalUnit)),
         Block(List(tryGetTree(tryReference)), mkStateTree(nextState, symLookup))
       )
@@ -382,12 +388,12 @@ trait ExprBuilder {
                   val t = c.Expr[Throwable](Ident(name.t))
                   val complete = futureSystemOps.completeProm[T](
                     c.Expr[futureSystem.Prom[T]](symLookup.memberRef(name.result)), futureSystemOps.tryyFailure[T](t)).tree
-                  Block(complete :: Nil, Return(literalUnit))
+                  Block(toList(complete), Return(literalUnit))
                 })), EmptyTree)
 
       def forever(t: Tree): Tree = {
         val labelName = name.fresh("while$")
-        LabelDef(labelName, Nil, Block(t :: Nil, Apply(Ident(labelName), Nil)))
+        LabelDef(labelName, Nil, Block(toList(t), Apply(Ident(labelName), Nil)))
       }
 
       /**
@@ -405,7 +411,7 @@ trait ExprBuilder {
       def onCompleteHandler[T: WeakTypeTag]: Tree = {
         val onCompletes = initStates.flatMap(_.mkOnCompleteHandler[T]).toList
         forever {
-          Block(resumeFunTree :: Nil, literalUnit)
+          adaptToUnit(toList(resumeFunTree))
         }
       }
     }
@@ -422,12 +428,32 @@ trait ExprBuilder {
     Assign(symLookup.memberRef(name.state), Literal(Constant(nextState)))
 
   private def mkHandlerCase(num: Int, rhs: List[Tree]): CaseDef =
-    mkHandlerCase(num, Block(rhs, literalUnit))
+    mkHandlerCase(num, adaptToUnit(rhs))
+
+  private def tpeOf(t: Tree): Type = t match {
+    case _ if t.tpe != null => t.tpe
+    case Try(body, Nil, _) => tpeOf(body)
+    case _ => NoType
+  }
+
+  private def adaptToUnit(rhs: List[Tree]): c.universe.Block = {
+    rhs match {
+      case init :+ last if tpeOf(last) <:< definitions.UnitTpe =>
+        Block(init, last)
+      case _ =>
+        Block(rhs, literalUnit)
+    }
+  }
 
   private def mkHandlerCase(num: Int, rhs: Tree): CaseDef =
     CaseDef(Literal(Constant(num)), EmptyTree, rhs)
 
-  def literalUnit = Literal(Constant(()))
+  def literalUnit = Literal(Constant(())) // a def to avoid sharing trees
+
+  def toList(tree: Tree): List[Tree] = tree match {
+    case Block(stats, Literal(Constant(value))) if value == () => stats
+    case _ => tree :: Nil
+  }
 
   def literalNull = Literal(Constant(null))
 }