Late expansion fixes

src/main/scala/scala/async/internal/AnfTransform.scala

@@ -27,6 +27,27 @@ private[async] trait AnfTransform {
     val tree1 = adjustTypeOfTranslatedPatternMatches(block, owner)
 
     var mode: AnfMode = Anf
+
+    object trace {
+      private var indent = -1
+
+      private def indentString = "  " * indent
+
+      def apply[T](args: Any)(t: => T): T = {
+        def prefix = mode.toString.toLowerCase
+        indent += 1
+        def oneLine(s: Any) = s.toString.replaceAll("""\n""", "\\\\n").take(127)
+        try {
+          AsyncUtils.trace(s"${indentString}$prefix(${oneLine(args)})")
+          val result = t
+          AsyncUtils.trace(s"${indentString}= ${oneLine(result)}")
+          result
+        } finally {
+          indent -= 1
+        }
+      }
+    }
+
     typingTransform(tree1, owner)((tree, api) => {
       def blockToList(tree: Tree): List[Tree] = tree match {
         case Block(stats, expr) => stats :+ expr
@@ -97,8 +118,11 @@ private[async] trait AnfTransform {
                 val ifWithAssign = treeCopy.If(tree, cond, branchWithAssign(thenp), branchWithAssign(elsep)).setType(definitions.UnitTpe)
                 stats :+ varDef :+ ifWithAssign :+ atPos(tree.pos)(gen.mkAttributedStableRef(varDef.symbol)).setType(tree.tpe)
               }
-            case LabelDef(name, params, rhs) =>
-              statsExprUnit
+            case ld @ LabelDef(name, params, rhs) =>
+              if (ld.symbol.info.resultType.typeSymbol == definitions.UnitClass)
+                statsExprUnit
+              else
+                stats :+ expr
 
             case Match(scrut, cases) =>
               // if type of match is Unit don't introduce assignment,
@@ -134,26 +158,6 @@ private[async] trait AnfTransform {
         }
       }
 
-      object trace {
-        private var indent = -1
-
-        private def indentString = "  " * indent
-
-        def apply[T](args: Any)(t: => T): T = {
-          def prefix = mode.toString.toLowerCase
-          indent += 1
-          def oneLine(s: Any) = s.toString.replaceAll("""\n""", "\\\\n").take(127)
-          try {
-            AsyncUtils.trace(s"${indentString}$prefix(${oneLine(args)})")
-            val result = t
-            AsyncUtils.trace(s"${indentString}= ${oneLine(result)}")
-            result
-          } finally {
-            indent -= 1
-          }
-        }
-      }
-
       def defineVal(prefix: String, lhs: Tree, pos: Position): ValDef = {
         val sym = api.currentOwner.newTermSymbol(name.fresh(prefix), pos, SYNTHETIC).setInfo(uncheckedBounds(lhs.tpe))
         internal.valDef(sym, internal.changeOwner(lhs, api.currentOwner, sym)).setType(NoType).setPos(pos)
@@ -219,8 +223,29 @@ private[async] trait AnfTransform {
               funStats ++ argStatss.flatten.flatten :+ typedNewApply
 
             case Block(stats, expr)                                    =>
-              val trees = stats.flatMap(linearize.transformToList).filterNot(isLiteralUnit) ::: linearize.transformToList(expr)
-              eliminateMatchEndLabelParameter(trees)
+              val stats1 = stats.flatMap(linearize.transformToList).filterNot(isLiteralUnit)
+              val exprs1 = linearize.transformToList(expr)
+              val trees = stats1 ::: exprs1
+              def isMatchEndLabel(t: Tree): Boolean = t match {
+                case ValDef(_, _, _, t) if isMatchEndLabel(t) => true
+                case ld: LabelDef if ld.name.toString.startsWith("matchEnd") => true
+                case _ => false
+              }
+              def groupsEndingWith[T](ts: List[T])(f: T => Boolean): List[List[T]] = if (ts.isEmpty) Nil else {
+                ts.indexWhere(f) match {
+                  case -1 => List(ts)
+                  case i =>
+                    val (ts1, ts2) = ts.splitAt(i + 1)
+                    ts1 :: groupsEndingWith(ts2)(f)
+                }
+              }
+              val matchGroups = groupsEndingWith(trees)(isMatchEndLabel)
+              val trees1 = matchGroups.flatMap(eliminateMatchEndLabelParameter)
+              val result = trees1 flatMap {
+                case Block(stats, expr) => stats :+ expr
+                case t => t :: Nil
+              }
+              result
 
             case ValDef(mods, name, tpt, rhs) =>
               if (containsAwait(rhs)) {
@@ -260,7 +285,10 @@ private[async] trait AnfTransform {
               scrutStats :+ treeCopy.Match(tree, scrutExpr, caseDefs)
 
             case LabelDef(name, params, rhs) =>
-              List(LabelDef(name, params, newBlock(linearize.transformToList(rhs), Literal(Constant(())))).setSymbol(tree.symbol))
+              if (tree.symbol.info.typeSymbol == definitions.UnitClass)
+                List(treeCopy.LabelDef(tree, name, params, api.typecheck(newBlock(linearize.transformToList(rhs), Literal(Constant(()))))).setSymbol(tree.symbol))
+              else
+                List(treeCopy.LabelDef(tree, name, params, api.typecheck(listToBlock(linearize.transformToList(rhs)))).setSymbol(tree.symbol))
 
             case TypeApply(fun, targs) =>
               val funStats :+ simpleFun = linearize.transformToList(fun)
@@ -274,7 +302,7 @@ private[async] trait AnfTransform {
 
       // Replace the label parameters on `matchEnd` with use of a `matchRes` temporary variable
       //
-      // CaseDefs are translated to labels without parmeters. A terminal label, `matchEnd`, accepts
+      // CaseDefs are translated to labels without parameters. 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:
@@ -286,34 +314,71 @@ private[async] trait AnfTransform {
         val caseDefToMatchResult = collection.mutable.Map[Symbol, Symbol]()
 
         val matchResults = collection.mutable.Buffer[Tree]()
-        val statsExpr0 = statsExpr.reverseMap {
-          case ld @ LabelDef(_, param :: Nil, body) =>
+        def modifyLabelDef(ld: LabelDef): (Tree, Tree) = {
+          val symTab = c.universe.asInstanceOf[reflect.internal.SymbolTable]
+          val param = ld.params.head
+          val ld2 = if (ld.params.head.tpe.typeSymbol == definitions.UnitClass) {
+            // Unit typed match: eliminate the label def parameter, but don't create a matchres temp variable to
+            // store the result for cleaner generated code.
+            caseDefToMatchResult(ld.symbol) = NoSymbol
+            val rhs2 = substituteTrees(ld.rhs, param.symbol :: Nil, api.typecheck(literalUnit) :: Nil)
+            (treeCopy.LabelDef(ld, ld.name, Nil, api.typecheck(literalUnit)), rhs2)
+          } else {
+            // Otherwise, create the matchres var. We'll callers of the label def below.
+            // Remember: we're iterating through the statement sequence in reverse, so we'll get
+            // to the LabelDef and mutate `matchResults` before we'll get to its callers.
             val matchResult = linearize.defineVar(name.matchRes, param.tpe, ld.pos)
             matchResults += matchResult
             caseDefToMatchResult(ld.symbol) = matchResult.symbol
-            val ld2 = treeCopy.LabelDef(ld, ld.name, Nil, body.substituteSymbols(param.symbol :: Nil, matchResult.symbol :: Nil))
-            setInfo(ld.symbol, methodType(Nil, ld.symbol.info.resultType))
-            ld2
+            val rhs2 = ld.rhs.substituteSymbols(param.symbol :: Nil, matchResult.symbol :: Nil)
+            (treeCopy.LabelDef(ld, ld.name, Nil, api.typecheck(literalUnit)), rhs2)
+          }
+          setInfo(ld.symbol, methodType(Nil, definitions.UnitTpe))
+          ld2
+        }
+        val statsExpr0 = statsExpr.reverse.flatMap {
+          case ld @ LabelDef(_, param :: Nil, _) =>
+            val (ld1, after) = modifyLabelDef(ld)
+            List(after, ld1)
+          case a @ ValDef(mods, name, tpt, ld @ LabelDef(_, param :: Nil, _)) =>
+            val (ld1, after) = modifyLabelDef(ld)
+            List(treeCopy.ValDef(a, mods, name, tpt, after), ld1)
           case t =>
-            if (caseDefToMatchResult.isEmpty) t
-            else typingTransform(t)((tree, api) =>
+            if (caseDefToMatchResult.isEmpty) t :: Nil
+            else typingTransform(t)((tree, api) => {
+              def typedPos(pos: Position)(t: Tree): Tree =
+                api.typecheck(atPos(pos)(t))
               tree match {
                 case Apply(fun, arg :: Nil) if isLabel(fun.symbol) && caseDefToMatchResult.contains(fun.symbol) =>
-                  api.typecheck(atPos(tree.pos)(newBlock(Assign(Ident(caseDefToMatchResult(fun.symbol)), api.recur(arg)) :: Nil, treeCopy.Apply(tree, fun, Nil))))
-                case Block(stats, expr) =>
+                  val temp = caseDefToMatchResult(fun.symbol)
+                  if (temp == NoSymbol)
+                    typedPos(tree.pos)(newBlock(api.recur(arg) :: Nil, treeCopy.Apply(tree, fun, Nil)))
+                  else
+                    // setType needed for LateExpansion.shadowingRefinedType test case. There seems to be an inconsistency
+                    // in the trees after pattern matcher.
+                    // TODO miminize the problem in patmat and fix in scalac.
+                    typedPos(tree.pos)(newBlock(Assign(Ident(temp), api.recur(internal.setType(arg, fun.tpe.paramLists.head.head.info))) :: Nil, treeCopy.Apply(tree, fun, Nil)))
+                case Block(stats, expr: Apply) if isLabel(expr.symbol) =>
                   api.default(tree) match {
-                    case Block(stats, Block(stats1, expr)) =>
-                      treeCopy.Block(tree, stats ::: stats1, expr)
+                    case Block(stats0, Block(stats1, expr1)) =>
+                      // flatten the block returned by `case Apply` above into the enclosing block for
+                      // cleaner generated code.
+                      treeCopy.Block(tree, stats0 ::: stats1, expr1)
                     case t => t
                   }
                 case _ =>
                   api.default(tree)
               }
-            )
+            }) :: Nil
         }
         matchResults.toList match {
-          case Nil => statsExpr
-          case r1 :: Nil => (r1 +: statsExpr0.reverse) :+ atPos(tree.pos)(gen.mkAttributedIdent(r1.symbol))
+          case _ if caseDefToMatchResult.isEmpty =>
+            statsExpr // return the original trees if nothing changed
+          case Nil =>
+            statsExpr0.reverse :+ literalUnit // must have been a unit-typed match, no matchRes variable to definne or refer to
+          case r1 :: Nil =>
+            // { var matchRes = _; ....; matchRes }
+            (r1 +: statsExpr0.reverse) :+ atPos(tree.pos)(gen.mkAttributedIdent(r1.symbol))
           case _ => c.error(macroPos, "Internal error: unexpected tree encountered during ANF transform " + statsExpr); statsExpr
         }
       }

src/main/scala/scala/async/internal/AsyncBase.scala

@@ -55,12 +55,14 @@ abstract class AsyncBase {
 
   protected[async] def asyncMethod(u: Universe)(asyncMacroSymbol: u.Symbol): u.Symbol = {
     import u._
-    asyncMacroSymbol.owner.typeSignature.member(newTermName("async"))
+    if (asyncMacroSymbol == null) NoSymbol
+    else asyncMacroSymbol.owner.typeSignature.member(newTermName("async"))
   }
 
   protected[async] def awaitMethod(u: Universe)(asyncMacroSymbol: u.Symbol): u.Symbol = {
     import u._
-    asyncMacroSymbol.owner.typeSignature.member(newTermName("await"))
+    if (asyncMacroSymbol == null) NoSymbol
+    else asyncMacroSymbol.owner.typeSignature.member(newTermName("await"))
   }
 
   protected[async] def nullOut(u: Universe)(name: u.Expr[String], v: u.Expr[Any]): u.Expr[Unit] =

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

@@ -49,8 +49,16 @@ trait AsyncTransform {
         List(emptyConstructor, stateVar) ++ resultAndAccessors ++ List(execContextValDef) ++ List(applyDefDefDummyBody, apply0DefDef)
       }
 
-      val tryToUnit = appliedType(definitions.FunctionClass(1), futureSystemOps.tryType[Any], typeOf[Unit])
-      val template = Template(List(tryToUnit, typeOf[() => Unit]).map(TypeTree(_)), emptyValDef, body)
+      val customParents = futureSystemOps.stateMachineClassParents
+      val tycon = if (customParents.exists(!_.typeSymbol.asClass.isTrait)) {
+        // prefer extending a class to reduce the class file size of the state machine.
+        symbolOf[scala.runtime.AbstractFunction1[Any, Any]]
+      } else {
+        // ... unless a custom future system already extends some class
+        symbolOf[scala.Function1[Any, Any]]
+      }
+      val tryToUnit = appliedType(tycon, futureSystemOps.tryType[Any], typeOf[Unit])
+      val template = Template((futureSystemOps.stateMachineClassParents ::: List(tryToUnit, typeOf[() => Unit])).map(TypeTree(_)), emptyValDef, body)
 
       val t = ClassDef(NoMods, name.stateMachineT, Nil, template)
       typecheckClassDef(t)

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 =
@@ -99,10 +97,10 @@ trait ExprBuilder {
         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),
-            Block(toList(callOnComplete), Return(literalUnit)))
+            Block(toList(callOnComplete), Return(literalUnit))) :: Nil
         else
-          Block(toList(callOnComplete), Return(literalUnit))
-      mkHandlerCase(state, stats ++ List(mkStateTree(onCompleteState, symLookup), tryGetOrCallOnComplete))
+          toList(callOnComplete) ::: Return(literalUnit) :: Nil
+      mkHandlerCase(state, stats ++ List(mkStateTree(onCompleteState, symLookup)) ++ tryGetOrCallOnComplete)
     }
 
     private def tryGetTree(tryReference: => Tree) =
@@ -251,12 +249,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()
@@ -315,10 +318,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
   }
@@ -357,8 +363,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 {
@@ -392,7 +398,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(
@@ -462,13 +471,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)
     }

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

@@ -33,6 +33,7 @@ trait FutureSystem {
     def promType[A: WeakTypeTag]: Type
     def tryType[A: WeakTypeTag]: Type
     def execContextType: Type
+    def stateMachineClassParents: List[Type] = Nil
 
     /** Create an empty promise */
     def createProm[A: WeakTypeTag]: Expr[Prom[A]]
@@ -55,6 +56,7 @@ trait FutureSystem {
 
     /** Complete a promise with a value */
     def completeProm[A](prom: Expr[Prom[A]], value: Expr[Tryy[A]]): Expr[Unit]
+    def completeWithSuccess[A: WeakTypeTag](prom: Expr[Prom[A]], value: Expr[A]): Expr[Unit] = completeProm(prom, tryySuccess(value))
 
     def spawn(tree: Tree, execContext: Tree): Tree =
       future(c.Expr[Unit](tree))(c.Expr[ExecContext](execContext)).tree