Merge pull request #98 from retronym/ticket/73

Make `f(await(completedFuture))` execute `f` synchronously

Author: retronym

src/main/scala/scala/async/internal/AsyncTransform.scala

@@ -24,31 +24,24 @@ trait AsyncTransform {
 
     val anfTree = futureSystemOps.postAnfTransform(anfTree0)
 
-    val resumeFunTreeDummyBody = DefDef(Modifiers(), name.resume, Nil, List(Nil), Ident(definitions.UnitClass), Literal(Constant(())))
-
     val applyDefDefDummyBody: DefDef = {
       val applyVParamss = List(List(ValDef(Modifiers(Flag.PARAM), name.tr, TypeTree(futureSystemOps.tryType[Any]), EmptyTree)))
-      DefDef(NoMods, name.apply, Nil, applyVParamss, TypeTree(definitions.UnitTpe), Literal(Constant(())))
+      DefDef(NoMods, name.apply, Nil, applyVParamss, TypeTree(definitions.UnitTpe), literalUnit)
     }
 
     // Create `ClassDef` of state machine with empty method bodies for `resume` and `apply`.
     val stateMachine: ClassDef = {
       val body: List[Tree] = {
-        val stateVar = ValDef(Modifiers(Flag.MUTABLE | Flag.PRIVATE | Flag.LOCAL), name.state, TypeTree(definitions.IntTpe), Literal(Constant(0)))
+        val stateVar = ValDef(Modifiers(Flag.MUTABLE | Flag.PRIVATE | Flag.LOCAL), name.state, TypeTree(definitions.IntTpe), Literal(Constant(StateAssigner.Initial)))
         val result = ValDef(NoMods, name.result, TypeTree(futureSystemOps.promType[T](uncheckedBoundsResultTag)), futureSystemOps.createProm[T](uncheckedBoundsResultTag).tree)
         val execContextValDef = ValDef(NoMods, name.execContext, TypeTree(), execContext)
 
         val apply0DefDef: DefDef = {
           // We extend () => Unit so we can pass this class as the by-name argument to `Future.apply`.
-          // See SI-1247 for the the optimization that avoids creatio
-          DefDef(NoMods, name.apply, Nil, Nil, TypeTree(definitions.UnitTpe), Apply(Ident(name.resume), Nil))
-        }
-        val extraValDef: ValDef = {
-          // We extend () => Unit so we can pass this class as the by-name argument to `Future.apply`.
-          // See SI-1247 for the the optimization that avoids creatio
-          ValDef(NoMods, newTermName("extra"), TypeTree(definitions.UnitTpe), Literal(Constant(())))
+          // See SI-1247 for the the optimization that avoids creation.
+          DefDef(NoMods, name.apply, Nil, Nil, TypeTree(definitions.UnitTpe), Apply(Ident(name.apply), literalNull :: Nil))
         }
-        List(emptyConstructor, stateVar, result, execContextValDef) ++ List(resumeFunTreeDummyBody, applyDefDefDummyBody, apply0DefDef, extraValDef)
+        List(emptyConstructor, stateVar, result, execContextValDef) ++ List(applyDefDefDummyBody, apply0DefDef)
       }
 
       val tryToUnit = appliedType(definitions.FunctionClass(1), futureSystemOps.tryType[Any], typeOf[Unit])
@@ -90,8 +83,7 @@ trait AsyncTransform {
       val stateMachineSpliced: Tree = spliceMethodBodies(
         liftedFields,
         stateMachine,
-        atMacroPos(asyncBlock.onCompleteHandler[T]),
-        atMacroPos(asyncBlock.resumeFunTree[T].rhs)
+        atMacroPos(asyncBlock.onCompleteHandler[T])
       )
 
       def selectStateMachine(selection: TermName) = Select(Ident(name.stateMachine), selection)
@@ -131,10 +123,9 @@ trait AsyncTransform {
    *  @param  liftables  trees of definitions that are lifted to fields of the state machine class
    *  @param  tree       `ClassDef` tree of the state machine class
    *  @param  applyBody  tree of onComplete handler (`apply` method)
-   *  @param  resumeBody RHS of definition tree of `resume` method
    *  @return            transformed `ClassDef` tree of the state machine class
    */
-  def spliceMethodBodies(liftables: List[Tree], tree: ClassDef, applyBody: Tree, resumeBody: Tree): Tree = {
+  def spliceMethodBodies(liftables: List[Tree], tree: ClassDef, applyBody: Tree): Tree = {
     val liftedSyms = liftables.map(_.symbol).toSet
     val stateMachineClass = tree.symbol
     liftedSyms.foreach {
@@ -211,12 +202,6 @@ trait AsyncTransform {
         (ctx: analyzer.Context) =>
           val typedTree = fixup(dd, changeOwner(applyBody, callSiteTyper.context.owner, dd.symbol), ctx)
           typedTree
-
-      case dd@DefDef(_, name.resume, _, _, _, _) if dd.symbol.owner == stateMachineClass =>
-        (ctx: analyzer.Context) =>
-          val changed = changeOwner(resumeBody, callSiteTyper.context.owner, dd.symbol)
-          val res = fixup(dd, changed, ctx)
-          res
     }
     result
   }

src/main/scala/scala/async/internal/ExprBuilder.scala

@@ -27,7 +27,7 @@ trait ExprBuilder {
 
     def nextStates: List[Int]
 
-    def mkHandlerCaseForState: CaseDef
+    def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef
 
     def mkOnCompleteHandler[T: WeakTypeTag]: Option[CaseDef] = None
 
@@ -51,8 +51,8 @@ trait ExprBuilder {
     def nextStates: List[Int] =
       List(nextState)
 
-    def mkHandlerCaseForState: CaseDef =
-      mkHandlerCase(state, stats :+ mkStateTree(nextState, symLookup) :+ mkResumeApply(symLookup))
+    def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef =
+      mkHandlerCase(state, stats :+ mkStateTree(nextState, symLookup))
 
     override val toString: String =
       s"AsyncState #$state, next = $nextState"
@@ -62,7 +62,7 @@ trait ExprBuilder {
     * a branch of an `if` or a `match`.
     */
   final class AsyncStateWithoutAwait(var stats: List[Tree], val state: Int, val nextStates: List[Int]) extends AsyncState {
-    override def mkHandlerCaseForState: CaseDef =
+    override def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef =
       mkHandlerCase(state, stats)
 
     override val toString: String =
@@ -72,45 +72,54 @@ trait ExprBuilder {
   /** A sequence of statements that concludes with an `await` call. The `onComplete`
     * handler will unconditionally transition to `nextState`.
     */
-  final class AsyncStateWithAwait(var stats: List[Tree], val state: Int, nextState: Int,
+  final class AsyncStateWithAwait(var stats: List[Tree], val state: Int, onCompleteState: Int, nextState: Int,
                                   val awaitable: Awaitable, symLookup: SymLookup)
     extends AsyncState {
 
     def nextStates: List[Int] =
       List(nextState)
 
-    override def mkHandlerCaseForState: CaseDef = {
-      val callOnComplete = futureSystemOps.onComplete(Expr(awaitable.expr),
-        Expr(This(tpnme.EMPTY)), Expr(Ident(name.execContext))).tree
-      mkHandlerCase(state, stats :+ callOnComplete)
+    override def mkHandlerCaseForState[T: WeakTypeTag]: CaseDef = {
+      val fun = This(tpnme.EMPTY)
+      val callOnComplete = futureSystemOps.onComplete[Any, Unit](Expr[futureSystem.Fut[Any]](awaitable.expr),
+        Expr[futureSystem.Tryy[Any] => Unit](fun), Expr[futureSystem.ExecContext](Ident(name.execContext))).tree
+      val tryGetOrCallOnComplete =
+        if (futureSystemOps.continueCompletedFutureOnSameThread)
+          If(futureSystemOps.isCompleted(Expr[futureSystem.Fut[_]](awaitable.expr)).tree,
+            Block(ifIsFailureTree[T](futureSystemOps.getCompleted[Any](Expr[futureSystem.Fut[Any]](awaitable.expr)).tree) :: Nil, literalUnit),
+            Block(callOnComplete :: Nil, Return(literalUnit)))
+        else
+          Block(callOnComplete :: Nil, Return(literalUnit))
+      mkHandlerCase(state, stats ++ List(mkStateTree(onCompleteState, symLookup), tryGetOrCallOnComplete))
     }
 
+    private def tryGetTree(tryReference: => Tree) =
+      Assign(
+        Ident(awaitable.resultName),
+        TypeApply(Select(futureSystemOps.tryyGet[Any](Expr[futureSystem.Tryy[Any]](tryReference)).tree, newTermName("asInstanceOf")), List(TypeTree(awaitable.resultType)))
+      )
+
+    /* if (tr.isFailure)
+     *   result.complete(tr.asInstanceOf[Try[T]])
+     * else {
+     *   <resultName> = tr.get.asInstanceOf[<resultType>]
+     *   <nextState>
+     *   <mkResumeApply>
+     * }
+     */
+    def ifIsFailureTree[T: WeakTypeTag](tryReference: => Tree) =
+      If(futureSystemOps.tryyIsFailure(Expr[futureSystem.Tryy[T]](tryReference)).tree,
+        Block(futureSystemOps.completeProm[T](
+          Expr[futureSystem.Prom[T]](symLookup.memberRef(name.result)),
+          Expr[futureSystem.Tryy[T]](
+            TypeApply(Select(tryReference, newTermName("asInstanceOf")),
+              List(TypeTree(futureSystemOps.tryType[T]))))).tree :: Nil,
+          Return(literalUnit)),
+        Block(List(tryGetTree(tryReference)), mkStateTree(nextState, symLookup))
+      )
+
     override def mkOnCompleteHandler[T: WeakTypeTag]: Option[CaseDef] = {
-      val tryGetTree =
-        Assign(
-          Ident(awaitable.resultName),
-          TypeApply(Select(futureSystemOps.tryyGet[T](Expr[futureSystem.Tryy[T]](Ident(symLookup.applyTrParam))).tree, newTermName("asInstanceOf")), List(TypeTree(awaitable.resultType)))
-        )
-
-      /* if (tr.isFailure)
-       *   result.complete(tr.asInstanceOf[Try[T]])
-       * else {
-       *   <resultName> = tr.get.asInstanceOf[<resultType>]
-       *   <nextState>
-       *   <mkResumeApply>
-       * }
-       */
-      val ifIsFailureTree =
-        If(futureSystemOps.tryyIsFailure(Expr[futureSystem.Tryy[T]](Ident(symLookup.applyTrParam))).tree,
-           futureSystemOps.completeProm[T](
-             Expr[futureSystem.Prom[T]](symLookup.memberRef(name.result)),
-             Expr[futureSystem.Tryy[T]](
-               TypeApply(Select(Ident(symLookup.applyTrParam), newTermName("asInstanceOf")),
-                         List(TypeTree(futureSystemOps.tryType[T]))))).tree,
-           Block(List(tryGetTree, mkStateTree(nextState, symLookup)), mkResumeApply(symLookup))
-         )
-
-      Some(mkHandlerCase(state, List(ifIsFailureTree)))
+      Some(mkHandlerCase(onCompleteState, List(ifIsFailureTree[T](Ident(symLookup.applyTrParam)))))
     }
 
     override val toString: String =
@@ -146,9 +155,10 @@ trait ExprBuilder {
     }
 
     def resultWithAwait(awaitable: Awaitable,
+                        onCompleteState: Int,
                         nextState: Int): AsyncState = {
       val effectiveNextState = nextJumpState.getOrElse(nextState)
-      new AsyncStateWithAwait(stats.toList, state, effectiveNextState, awaitable, symLookup)
+      new AsyncStateWithAwait(stats.toList, state, onCompleteState, effectiveNextState, awaitable, symLookup)
     }
 
     def resultSimple(nextState: Int): AsyncState = {
@@ -157,7 +167,7 @@ trait ExprBuilder {
     }
 
     def resultWithIf(condTree: Tree, thenState: Int, elseState: Int): AsyncState = {
-      def mkBranch(state: Int) = Block(mkStateTree(state, symLookup) :: Nil, mkResumeApply(symLookup))
+      def mkBranch(state: Int) = mkStateTree(state, symLookup)
       this += If(condTree, mkBranch(thenState), mkBranch(elseState))
       new AsyncStateWithoutAwait(stats.toList, state, List(thenState, elseState))
     }
@@ -177,15 +187,15 @@ trait ExprBuilder {
       val newCases = for ((cas, num) <- cases.zipWithIndex) yield cas match {
         case CaseDef(pat, guard, rhs) =>
           val bindAssigns = rhs.children.takeWhile(isSyntheticBindVal)
-          CaseDef(pat, guard, Block(bindAssigns :+ mkStateTree(caseStates(num), symLookup), mkResumeApply(symLookup)))
+          CaseDef(pat, guard, Block(bindAssigns, mkStateTree(caseStates(num), symLookup)))
       }
       // 2. insert changed match tree at the end of the current state
       this += Match(scrutTree, newCases)
       new AsyncStateWithoutAwait(stats.toList, state, caseStates)
     }
 
     def resultWithLabel(startLabelState: Int, symLookup: SymLookup): AsyncState = {
-      this += Block(mkStateTree(startLabelState, symLookup) :: Nil, mkResumeApply(symLookup))
+      this += mkStateTree(startLabelState, symLookup)
       new AsyncStateWithoutAwait(stats.toList, state, List(startLabelState))
     }
 
@@ -226,9 +236,10 @@ trait ExprBuilder {
     for (stat <- stats) stat match {
       // the val name = await(..) pattern
       case vd @ ValDef(mods, name, tpt, Apply(fun, arg :: Nil)) if isAwait(fun) =>
+        val onCompleteState = nextState()
         val afterAwaitState = nextState()
         val awaitable = Awaitable(arg, stat.symbol, tpt.tpe, vd)
-        asyncStates += stateBuilder.resultWithAwait(awaitable, afterAwaitState) // complete with await
+        asyncStates += stateBuilder.resultWithAwait(awaitable, onCompleteState, afterAwaitState) // complete with await
         currState = afterAwaitState
         stateBuilder = new AsyncStateBuilder(currState, symLookup)
 
@@ -296,8 +307,6 @@ trait ExprBuilder {
     def asyncStates: List[AsyncState]
 
     def onCompleteHandler[T: WeakTypeTag]: Tree
-
-    def resumeFunTree[T: WeakTypeTag]: DefDef
   }
 
   case class SymLookup(stateMachineClass: Symbol, applyTrParam: Symbol) {
@@ -330,13 +339,13 @@ trait ExprBuilder {
           val lastStateBody = Expr[T](lastState.body)
           val rhs = futureSystemOps.completeProm(
             Expr[futureSystem.Prom[T]](symLookup.memberRef(name.result)), futureSystemOps.tryySuccess[T](lastStateBody))
-          mkHandlerCase(lastState.state, rhs.tree)
+          mkHandlerCase(lastState.state, Block(rhs.tree, Return(literalUnit)))
         }
         asyncStates.toList match {
           case s :: Nil =>
             List(caseForLastState)
           case _        =>
-            val initCases = for (state <- asyncStates.toList.init) yield state.mkHandlerCaseForState
+            val initCases = for (state <- asyncStates.toList.init) yield state.mkHandlerCaseForState[T]
             initCases :+ caseForLastState
         }
       }
@@ -362,18 +371,23 @@ trait ExprBuilder {
        *       }
        *     }
        */
-      def resumeFunTree[T: WeakTypeTag]: DefDef =
-        DefDef(Modifiers(), name.resume, Nil, List(Nil), Ident(definitions.UnitClass),
+      private def resumeFunTree[T: WeakTypeTag]: Tree =
           Try(
-            Match(symLookup.memberRef(name.state), mkCombinedHandlerCases[T]),
+            Match(symLookup.memberRef(name.state), mkCombinedHandlerCases[T] ++ initStates.flatMap(_.mkOnCompleteHandler[T]) ),
             List(
               CaseDef(
                 Bind(name.t, Ident(nme.WILDCARD)),
                 Apply(Ident(defn.NonFatalClass), List(Ident(name.t))), {
                   val t = Expr[Throwable](Ident(name.t))
-                  futureSystemOps.completeProm[T](
+                  val complete = futureSystemOps.completeProm[T](
                     Expr[futureSystem.Prom[T]](symLookup.memberRef(name.result)), futureSystemOps.tryyFailure[T](t)).tree
-                })), EmptyTree))
+                  Block(complete :: Nil, Return(literalUnit))
+                })), EmptyTree)
+
+      def forever(t: Tree): Tree = {
+        val labelName = name.fresh("while$")
+        LabelDef(labelName, Nil, Block(t :: Nil, Apply(Ident(labelName), Nil)))
+      }
 
       /**
        * Builds a `match` expression used as an onComplete handler.
@@ -387,8 +401,12 @@ trait ExprBuilder {
        *         resume()
        *     }
        */
-      def onCompleteHandler[T: WeakTypeTag]: Tree =
-        Match(symLookup.memberRef(name.state), initStates.flatMap(_.mkOnCompleteHandler[T]).toList)
+      def onCompleteHandler[T: WeakTypeTag]: Tree = {
+        val onCompletes = initStates.flatMap(_.mkOnCompleteHandler[T]).toList
+        forever {
+          Block(resumeFunTree :: Nil, literalUnit)
+        }
+      }
     }
   }
 
@@ -399,9 +417,6 @@ trait ExprBuilder {
 
   case class Awaitable(expr: Tree, resultName: Symbol, resultType: Type, resultValDef: ValDef)
 
-  private def mkResumeApply(symLookup: SymLookup) =
-    Apply(symLookup.memberRef(name.resume), Nil)
-
   private def mkStateTree(nextState: Int, symLookup: SymLookup): Tree =
     Assign(symLookup.memberRef(name.state), Literal(Constant(nextState)))
 
@@ -411,5 +426,7 @@ trait ExprBuilder {
   private def mkHandlerCase(num: Int, rhs: Tree): CaseDef =
     CaseDef(Literal(Constant(num)), EmptyTree, rhs)
 
-  private def literalUnit = Literal(Constant(()))
+  def literalUnit = Literal(Constant(()))
+
+  def literalNull = Literal(Constant(null))
 }

src/main/scala/scala/async/internal/FutureSystem.scala

@@ -49,6 +49,12 @@ trait FutureSystem {
     def onComplete[A, U](future: Expr[Fut[A]], fun: Expr[Tryy[A] => U],
                          execContext: Expr[ExecContext]): Expr[Unit]
 
+    def continueCompletedFutureOnSameThread = false
+    def isCompleted(future: Expr[Fut[_]]): Expr[Boolean] =
+      throw new UnsupportedOperationException("isCompleted not supported by this FutureSystem")
+    def getCompleted[A: WeakTypeTag](future: Expr[Fut[A]]): Expr[Tryy[A]] =
+      throw new UnsupportedOperationException("getCompleted not supported by this FutureSystem")
+
     /** Complete a promise with a value */
     def completeProm[A](prom: Expr[Prom[A]], value: Expr[Tryy[A]]): Expr[Unit]
 
@@ -103,6 +109,15 @@ object ScalaConcurrentFutureSystem extends FutureSystem {
       future.splice.onComplete(fun.splice)(execContext.splice)
     }
 
+    override def continueCompletedFutureOnSameThread: Boolean = true
+
+    override def isCompleted(future: Expr[Fut[_]]): Expr[Boolean] = reify {
+      future.splice.isCompleted
+    }
+    override def getCompleted[A: WeakTypeTag](future: Expr[Fut[A]]): Expr[Tryy[A]] = reify {
+      future.splice.value.get
+    }
+
     def completeProm[A](prom: Expr[Prom[A]], value: Expr[scala.util.Try[A]]): Expr[Unit] = reify {
       prom.splice.complete(value.splice)
       Expr[Unit](Literal(Constant(()))).splice

src/main/scala/scala/async/internal/StateAssigner.scala

@@ -5,10 +5,12 @@
 package scala.async.internal
 
 private[async] final class StateAssigner {
-  private var current = -1
+  private var current = StateAssigner.Initial
 
-  def nextState(): Int = {
-    current += 1
-    current
-  }
+  def nextState(): Int =
+    try current finally current += 1
 }
+
+object StateAssigner {
+  final val Initial = 0
+}

src/main/scala/scala/async/internal/TransformUtils.scala

@@ -49,6 +49,7 @@ private[async] trait TransformUtils {
 
   private def isByName(fun: Tree): ((Int, Int) => Boolean) = {
     if (Boolean_ShortCircuits contains fun.symbol) (i, j) => true
+    else if (fun.tpe == null) (x, y) => false
     else {
       val paramss = fun.tpe.paramss
       val byNamess = paramss.map(_.map(_.isByNameParam))
@@ -72,10 +73,6 @@ private[async] trait TransformUtils {
       self.splice.contains(elem.splice)
     }
 
-    def mkFunction_apply[A, B](self: Expr[Function1[A, B]])(arg: Expr[A]) = reify {
-      self.splice.apply(arg.splice)
-    }
-
     def mkAny_==(self: Expr[Any])(other: Expr[Any]) = reify {
       self.splice == other.splice
     }