1 files changed, 47 insertions, 26 deletions
diff --git a/src/main/scala/scala/async/internal/ExprBuilder.scala b/src/main/scala/scala/async/internal/ExprBuilder.scala
index 16b9207..fe62cd6 100644
--- a/src/main/scala/scala/async/internal/ExprBuilder.scala
+++ b/src/main/scala/scala/async/internal/ExprBuilder.scala
@@ -3,7 +3,6 @@
  */
 package scala.async.internal
 
-import scala.reflect.macros.Context
 import scala.collection.mutable.ListBuffer
 import collection.mutable
 import language.existentials
@@ -34,18 +33,17 @@ trait ExprBuilder {
 
     var stats: List[Tree]
 
-    def statsAnd(trees: List[Tree]): List[Tree] = {
-      val body = stats match {
+    def treesThenStats(trees: List[Tree]): List[Tree] = {
+      (stats match {
         case init :+ last if tpeOf(last) =:= definitions.NothingTpe =>
-          adaptToUnit(init :+ Typed(last, TypeTree(definitions.AnyTpe)))
+          adaptToUnit((trees ::: init) :+ Typed(last, TypeTree(definitions.AnyTpe)))
         case _ =>
-          adaptToUnit(stats)
-      }
-      Try(body, Nil, adaptToUnit(trees)) :: Nil
+          adaptToUnit(trees ::: stats)
+      }) :: Nil
     }
 
     final def allStats: List[Tree] = this match {
-      case a: AsyncStateWithAwait => statsAnd(a.awaitable.resultValDef :: Nil)
+      case a: AsyncStateWithAwait => treesThenStats(a.awaitable.resultValDef :: Nil)
       case _ => stats
     }
 
@@ -63,7 +61,7 @@ trait ExprBuilder {
       List(nextState)
 
     def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef = {
-      mkHandlerCase(state, statsAnd(mkStateTree(nextState, symLookup) :: Nil))
+      mkHandlerCase(state, treesThenStats(mkStateTree(nextState, symLookup) :: Nil))
     }
 
     override val toString: String =
@@ -95,14 +93,17 @@ trait ExprBuilder {
       val fun = This(tpnme.EMPTY)
       val callOnComplete = futureSystemOps.onComplete[Any, Unit](c.Expr[futureSystem.Fut[Any]](awaitable.expr),
         c.Expr[futureSystem.Tryy[Any] => Unit](fun), c.Expr[futureSystem.ExecContext](Ident(name.execContext))).tree
-      val tryGetOrCallOnComplete =
-        if (futureSystemOps.continueCompletedFutureOnSameThread)
-          If(futureSystemOps.isCompleted(c.Expr[futureSystem.Fut[_]](awaitable.expr)).tree,
-            adaptToUnit(ifIsFailureTree[T](futureSystemOps.getCompleted[Any](c.Expr[futureSystem.Fut[Any]](awaitable.expr)).tree) :: Nil),
+      val tryGetOrCallOnComplete: List[Tree] =
+        if (futureSystemOps.continueCompletedFutureOnSameThread) {
+          val tempName = name.fresh(name.completed)
+          val initTemp = ValDef(NoMods, tempName, TypeTree(futureSystemOps.tryType[Any]), futureSystemOps.getCompleted[Any](c.Expr[futureSystem.Fut[Any]](awaitable.expr)).tree)
+          val ifTree = If(Apply(Select(Literal(Constant(null)), TermName("ne")), Ident(tempName) :: Nil),
+            adaptToUnit(ifIsFailureTree[T](Ident(tempName)) :: Nil),
             Block(toList(callOnComplete), Return(literalUnit)))
-        else
-          Block(toList(callOnComplete), Return(literalUnit))
-      mkHandlerCase(state, stats ++ List(mkStateTree(onCompleteState, symLookup), tryGetOrCallOnComplete))
+          initTemp :: ifTree :: Nil
+        } else
+          toList(callOnComplete) ::: Return(literalUnit) :: Nil
+      mkHandlerCase(state, stats ++ List(mkStateTree(onCompleteState, symLookup)) ++ tryGetOrCallOnComplete)
     }
 
     private def tryGetTree(tryReference: => Tree) =
@@ -237,10 +238,8 @@ trait ExprBuilder {
     var stateBuilder = new AsyncStateBuilder(startState, symLookup)
     var currState    = startState
 
-    def checkForUnsupportedAwait(tree: Tree) = if (tree exists {
-      case Apply(fun, _) if isAwait(fun) => true
-      case _                             => false
-    }) c.abort(tree.pos, "await must not be used in this position")
+    def checkForUnsupportedAwait(tree: Tree) = if (containsAwait(tree))
+      c.abort(tree.pos, "await must not be used in this position")
 
     def nestedBlockBuilder(nestedTree: Tree, startState: Int, endState: Int) = {
       val (nestedStats, nestedExpr) = statsAndExpr(nestedTree)
@@ -253,12 +252,17 @@ trait ExprBuilder {
         case LabelDef(name, _, _) => name.toString.startsWith("case")
         case _ => false
       }
-      val (before, _ :: after) = (stats :+ expr).span(_ ne t)
-      before.reverse.takeWhile(isPatternCaseLabelDef) ::: after.takeWhile(isPatternCaseLabelDef)
+      val span = (stats :+ expr).filterNot(isLiteralUnit).span(_ ne t)
+      span match {
+        case (before, _ :: after) =>
+          before.reverse.takeWhile(isPatternCaseLabelDef) ::: after.takeWhile(isPatternCaseLabelDef)
+        case _ =>
+          stats :+ expr
+      }
     }
 
     // populate asyncStates
-    for (stat <- (stats :+ expr)) stat match {
+    def add(stat: Tree): Unit = stat match {
       // the val name = await(..) pattern
       case vd @ ValDef(mods, name, tpt, Apply(fun, arg :: Nil)) if isAwait(fun) =>
         val onCompleteState = nextState()
@@ -317,10 +321,13 @@ trait ExprBuilder {
         asyncStates ++= builder.asyncStates
         currState = afterLabelState
         stateBuilder = new AsyncStateBuilder(currState, symLookup)
+      case b @ Block(stats, expr) =>
+        (stats :+ expr) foreach (add)
       case _ =>
         checkForUnsupportedAwait(stat)
         stateBuilder += stat
     }
+    for (stat <- (stats :+ expr)) add(stat)
     val lastState = stateBuilder.resultSimple(endState)
     asyncStates += lastState
   }
@@ -359,8 +366,8 @@ trait ExprBuilder {
         val caseForLastState: CaseDef = {
           val lastState = asyncStates.last
           val lastStateBody = c.Expr[T](lastState.body)
-          val rhs = futureSystemOps.completeProm(
-            c.Expr[futureSystem.Prom[T]](symLookup.memberRef(name.result)), futureSystemOps.tryySuccess[T](lastStateBody))
+          val rhs = futureSystemOps.completeWithSuccess(
+            c.Expr[futureSystem.Prom[T]](symLookup.memberRef(name.result)), lastStateBody)
           mkHandlerCase(lastState.state, Block(rhs.tree, Return(literalUnit)))
         }
         asyncStates.toList match {
@@ -394,7 +401,10 @@ trait ExprBuilder {
        *     }
        */
       private def resumeFunTree[T: WeakTypeTag]: Tree = {
-        val body = Match(symLookup.memberRef(name.state), mkCombinedHandlerCases[T] ++ initStates.flatMap(_.mkOnCompleteHandler[T]))
+        val stateMemberSymbol = symLookup.stateMachineMember(name.state)
+        val stateMemberRef = symLookup.memberRef(name.state)
+        val body = Match(stateMemberRef, mkCombinedHandlerCases[T] ++ initStates.flatMap(_.mkOnCompleteHandler[T]) ++ List(CaseDef(Ident(nme.WILDCARD), EmptyTree, Throw(Apply(Select(New(Ident(defn.IllegalStateExceptionClass)), termNames.CONSTRUCTOR), List())))))
+
         Try(
           body,
           List(
@@ -464,13 +474,24 @@ trait ExprBuilder {
   private def tpeOf(t: Tree): Type = t match {
     case _ if t.tpe != null => t.tpe
     case Try(body, Nil, _) => tpeOf(body)
+    case Block(_, expr) => tpeOf(expr)
+    case Literal(Constant(value)) if value == () => definitions.UnitTpe
+    case Return(_) => definitions.NothingTpe
     case _ => NoType
   }
 
   private def adaptToUnit(rhs: List[Tree]): c.universe.Block = {
     rhs match {
+      case (rhs: Block) :: Nil if tpeOf(rhs) <:< definitions.UnitTpe =>
+        rhs
       case init :+ last if tpeOf(last) <:< definitions.UnitTpe =>
         Block(init, last)
+      case init :+ (last @ Literal(Constant(()))) =>
+        Block(init, last)
+      case init :+ (last @ Block(_, Return(_) | Literal(Constant(())))) =>
+        Block(init, last)
+      case init :+ (Block(stats, expr)) =>
+        Block(init, Block(stats :+ expr, literalUnit))
       case _ =>
         Block(rhs, literalUnit)
     }