Generate cleaner code for Unit-typed if-else expressions

Instead of generating an assignment to a Unit-typed variable, just add
the Unit value after the if-else expression:

    [async] scala.async.Async.async[Unit]({
      val f: scala.concurrent.Future[Int] = AnfTestClass.this.base(y);
      if (y.>(0))
        State.result_=(scala.async.Async.await[Int](f).+(2))
      else
        State.result_=(scala.async.Async.await[Int](f).-(2))
    })
    [async] ANF transform expands to:
    {
      val f: scala.concurrent.Future[Int] = AnfTestClass.this.base(y);
      if (y.>(0))
        {
          val await$8: Int = scala.async.Async.await[Int](f);
          State.result_=(await$8.+(2))
        }
      else
        {
          val await$9: Int = scala.async.Async.await[Int](f);
          State.result_=(await$9.-(2))
        };
      ()
    }

1 files changed, 19 insertions, 13 deletions

diff --git a/src/main/scala/scala/async/AnfTransform.scala b/src/main/scala/scala/async/AnfTransform.scala
index 2dd6f8c..f29d6a1 100644
--- a/src/main/scala/scala/async/AnfTransform.scala
+++ b/src/main/scala/scala/async/AnfTransform.scala
@@ -16,21 +16,27 @@ class AnfTransform[C <: Context](override val c: C) extends TransformUtils(c) {
           stats :+ ValDef(NoMods, liftedName, TypeTree(), expr) :+ Ident(liftedName)
 
         case If(cond, thenp, elsep) =>
-          val liftedName = c.fresh("ifres$")
-          val varDef =
-            ValDef(Modifiers(Flag.MUTABLE), liftedName, TypeTree(expr.tpe), defaultValue(expr.tpe))
-          val thenWithAssign = thenp match {
-            case Block(thenStats, thenExpr) => Block(thenStats, Assign(Ident(liftedName), thenExpr))
-            case _ => Assign(Ident(liftedName), thenp)
+          // 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 :+ Literal(Constant(()))
           }
-          val elseWithAssign = elsep match {
-            case Block(elseStats, elseExpr) => Block(elseStats, Assign(Ident(liftedName), elseExpr))
-            case _ => Assign(Ident(liftedName), elsep)
+          else {
+            val liftedName = c.fresh("ifres$")
+            val varDef =
+              ValDef(Modifiers(Flag.MUTABLE), liftedName, TypeTree(expr.tpe), defaultValue(expr.tpe))
+            val thenWithAssign = thenp match {
+              case Block(thenStats, thenExpr) => Block(thenStats, Assign(Ident(liftedName), thenExpr))
+              case _ => Assign(Ident(liftedName), thenp)
+            }
+            val elseWithAssign = elsep match {
+              case Block(elseStats, elseExpr) => Block(elseStats, Assign(Ident(liftedName), elseExpr))
+              case _ => Assign(Ident(liftedName), elsep)
+            }
+            val ifWithAssign =
+              If(cond, thenWithAssign, elseWithAssign)
+            stats :+ varDef :+ ifWithAssign :+ Ident(liftedName)
           }
-          val ifWithAssign =
-            If(cond, thenWithAssign, elseWithAssign)
-          stats :+ varDef :+ ifWithAssign :+ Ident(liftedName)
-
         case _ =>
           stats :+ expr
       }