An overdue overhaul of macro internals.

 - Avoid reset + retypecheck, instead hang onto the original types/symbols
 - Eliminated duplication between AsyncDefinitionUseAnalyzer and ExprBuilder
   - Instead, decide what do lift *after* running ExprBuilder
   - Account for transitive references local classes/objects and lift them
     as needed.
 - Make the execution context an regular implicit parameter of the macro
 - Fixes interaction with existential skolems and singleton types

Fixes #6, #13, #16, #17, #19, #21.

Author: retronym

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

@@ -7,6 +7,9 @@ package scala.async
 import scala.language.experimental.macros
 import scala.reflect.macros.Context
 import scala.reflect.internal.annotations.compileTimeOnly
+import scala.tools.nsc.Global
+import language.reflectiveCalls
+import scala.concurrent.ExecutionContext
 
 object Async extends AsyncBase {
 
@@ -15,18 +18,22 @@ object Async extends AsyncBase {
   lazy val futureSystem = ScalaConcurrentFutureSystem
   type FS = ScalaConcurrentFutureSystem.type
 
-  def async[T](body: T) = macro asyncImpl[T]
+  def async[T](body: T)(implicit execContext: ExecutionContext): Future[T] = macro asyncImpl[T]
 
-  override def asyncImpl[T: c.WeakTypeTag](c: Context)(body: c.Expr[T]): c.Expr[Future[T]] = super.asyncImpl[T](c)(body)
+  override def asyncImpl[T: c.WeakTypeTag](c: Context)
+                                          (body: c.Expr[T])
+                                          (execContext: c.Expr[futureSystem.ExecContext]): c.Expr[Future[T]] = {
+    super.asyncImpl[T](c)(body)(execContext)
+  }
 }
 
 object AsyncId extends AsyncBase {
   lazy val futureSystem = IdentityFutureSystem
   type FS = IdentityFutureSystem.type
 
-  def async[T](body: T) = macro asyncImpl[T]
+  def async[T](body: T) = macro asyncIdImpl[T]
 
-  override def asyncImpl[T: c.WeakTypeTag](c: Context)(body: c.Expr[T]): c.Expr[T] = super.asyncImpl[T](c)(body)
+  def asyncIdImpl[T: c.WeakTypeTag](c: Context)(body: c.Expr[T]): c.Expr[T] = asyncImpl[T](c)(body)(c.literalUnit)
 }
 
 /**
@@ -62,124 +69,26 @@ abstract class AsyncBase {
 
   protected[async] def fallbackEnabled = false
 
-  def asyncImpl[T: c.WeakTypeTag](c: Context)(body: c.Expr[T]): c.Expr[futureSystem.Fut[T]] = {
+  def asyncImpl[T: c.WeakTypeTag](c: Context)
+                                 (body: c.Expr[T])
+                                 (execContext: c.Expr[futureSystem.ExecContext]): c.Expr[futureSystem.Fut[T]] = {
     import c.universe._
 
-    val analyzer = AsyncAnalysis[c.type](c, this)
-    val utils = TransformUtils[c.type](c)
-    import utils.{name, defn}
-
-    analyzer.reportUnsupportedAwaits(body.tree)
-
-    // Transform to A-normal form:
-    //  - no await calls in qualifiers or arguments,
-    //  - if/match only used in statement position.
-    val anfTree: Block = {
-      val anf = AnfTransform[c.type](c)
-      val restored = utils.restorePatternMatchingFunctions(body.tree)
-      val stats1 :+ expr1 = anf(restored)
-      val block = Block(stats1, expr1)
-      c.typeCheck(block).asInstanceOf[Block]
-    }
-
-    // Analyze the block to find locals that will be accessed from multiple
-    // states of our generated state machine, e.g. a value assigned before
-    // an `await` and read afterwards.
-    val renameMap: Map[Symbol, TermName] = {
-      analyzer.defTreesUsedInSubsequentStates(anfTree).map {
-        vd =>
-          (vd.symbol, name.fresh(vd.name.toTermName))
-      }.toMap
-    }
-
-    val builder = ExprBuilder[c.type, futureSystem.type](c, self.futureSystem, anfTree)
-    import builder.futureSystemOps
-    val asyncBlock: builder.AsyncBlock = builder.build(anfTree, renameMap)
-    import asyncBlock.asyncStates
-    logDiagnostics(c)(anfTree, asyncStates.map(_.toString))
-
-    // Important to retain the original declaration order here!
-    val localVarTrees = anfTree.collect {
-      case vd@ValDef(_, _, tpt, _) if renameMap contains vd.symbol                            =>
-        utils.mkVarDefTree(tpt.tpe, renameMap(vd.symbol))
-      case dd@DefDef(mods, name, tparams, vparamss, tpt, rhs) if renameMap contains dd.symbol =>
-        DefDef(mods, renameMap(dd.symbol), tparams, vparamss, tpt, c.resetAllAttrs(utils.substituteNames(rhs, renameMap)))
-    }
-
-    val onCompleteHandler = {
-      Function(
-        List(ValDef(Modifiers(Flag.PARAM), name.tr, TypeTree(defn.TryAnyType), EmptyTree)),
-        asyncBlock.onCompleteHandler)
-    }
-    val resumeFunTree = asyncBlock.resumeFunTree[T]
-
-    val stateMachineType = utils.applied("scala.async.StateMachine", List(futureSystemOps.promType[T], futureSystemOps.execContextType))
-
-    lazy val stateMachine: ClassDef = {
-      val body: List[Tree] = {
-        val stateVar = ValDef(Modifiers(Flag.MUTABLE), name.state, TypeTree(definitions.IntTpe), Literal(Constant(0)))
-        val result = ValDef(NoMods, name.result, TypeTree(futureSystemOps.promType[T]), futureSystemOps.createProm[T].tree)
-        val execContext = ValDef(NoMods, name.execContext, TypeTree(), futureSystemOps.execContext.tree)
-        val applyDefDef: DefDef = {
-          val applyVParamss = List(List(ValDef(Modifiers(Flag.PARAM), name.tr, TypeTree(defn.TryAnyType), EmptyTree)))
-          val applyBody = asyncBlock.onCompleteHandler
-          DefDef(NoMods, name.apply, Nil, applyVParamss, TypeTree(definitions.UnitTpe), applyBody)
-        }
-        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
-          val applyVParamss = List(List(ValDef(Modifiers(Flag.PARAM), name.tr, TypeTree(defn.TryAnyType), EmptyTree)))
-          val applyBody = asyncBlock.onCompleteHandler
-          DefDef(NoMods, name.apply, Nil, Nil, TypeTree(definitions.UnitTpe), Apply(Ident(name.resume), Nil))
-        }
-        List(utils.emptyConstructor, stateVar, result, execContext) ++ localVarTrees ++ List(resumeFunTree, applyDefDef, apply0DefDef)
-      }
-      val template = {
-        Template(List(stateMachineType), emptyValDef, body)
-      }
-      ClassDef(NoMods, name.stateMachineT, Nil, template)
-    }
-
-    def selectStateMachine(selection: TermName) = Select(Ident(name.stateMachine), selection)
-
-    val code: c.Expr[futureSystem.Fut[T]] = {
-      val isSimple = asyncStates.size == 1
-      val tree =
-        if (isSimple)
-          Block(Nil, futureSystemOps.spawn(body.tree)) // generate lean code for the simple case of `async { 1 + 1 }`
-        else {
-          Block(List[Tree](
-            stateMachine,
-            ValDef(NoMods, name.stateMachine, stateMachineType, Apply(Select(New(Ident(name.stateMachineT)), nme.CONSTRUCTOR), Nil)),
-            futureSystemOps.spawn(Apply(selectStateMachine(name.apply), Nil))
-          ),
-          futureSystemOps.promiseToFuture(c.Expr[futureSystem.Prom[T]](selectStateMachine(name.result))).tree)
-        }
-      c.Expr[futureSystem.Fut[T]](tree)
-    }
-
-    AsyncUtils.vprintln(s"async state machine transform expands to:\n ${code.tree}")
-    code
-  }
+    val asyncMacro = AsyncMacro(c, futureSystem)
+
+    val code = asyncMacro.asyncTransform[T](
+      body.tree.asInstanceOf[asyncMacro.global.Tree],
+      execContext.tree.asInstanceOf[asyncMacro.global.Tree],
+      fallbackEnabled)(implicitly[c.WeakTypeTag[T]].asInstanceOf[asyncMacro.global.WeakTypeTag[T]])
 
-  def logDiagnostics(c: Context)(anfTree: c.Tree, states: Seq[String]) {
-    def location = try {
-      c.macroApplication.pos.source.path
-    } catch {
-      case _: UnsupportedOperationException =>
-        c.macroApplication.pos.toString
-    }
-
-    AsyncUtils.vprintln(s"In file '$location':")
-    AsyncUtils.vprintln(s"${c.macroApplication}")
-    AsyncUtils.vprintln(s"ANF transform expands to:\n $anfTree")
-    states foreach (s => AsyncUtils.vprintln(s))
+    AsyncUtils.vprintln(s"async state machine transform expands to:\n ${code}")
+    c.Expr[futureSystem.Fut[T]](code.asInstanceOf[Tree])
   }
 }
 
 /** Internal class used by the `async` macro; should not be manually extended by client code */
 abstract class StateMachine[Result, EC] extends (scala.util.Try[Any] => Unit) with (() => Unit) {
-  def result$async: Result
+  def result: Result
 
-  def execContext$async: EC
+  def execContext: EC
 }

src/main/scala/scala/async/Async.scala

@@ -7,60 +7,37 @@ package scala.async
 import scala.reflect.macros.Context
 import scala.collection.mutable
 
-private[async] final case class AsyncAnalysis[C <: Context](c: C, asyncBase: AsyncBase) {
-  import c.universe._
+trait AsyncAnalysis {
+  self: AsyncMacro =>
 
-  val utils = TransformUtils[c.type](c)
-
-  import utils._
+  import global._
 
   /**
    * Analyze the contents of an `async` block in order to:
    * - Report unsupported `await` calls under nested templates, functions, by-name arguments.
    *
    * Must be called on the original tree, not on the ANF transformed tree.
    */
-  def reportUnsupportedAwaits(tree: Tree): Boolean = {
-    val analyzer = new UnsupportedAwaitAnalyzer
+  def reportUnsupportedAwaits(tree: Tree, report: Boolean): Boolean = {
+    val analyzer = new UnsupportedAwaitAnalyzer(report)
     analyzer.traverse(tree)
     analyzer.hasUnsupportedAwaits
   }
 
-  /**
-   * Analyze the contents of an `async` block in order to:
-   * - Find which local `ValDef`-s need to be lifted to fields of the state machine, based
-   * on whether or not they are accessed only from a single state.
-   *
-   * Must be called on the ANF transformed tree.
-   */
-  def defTreesUsedInSubsequentStates(tree: Tree): List[DefTree] = {
-    val analyzer = new AsyncDefinitionUseAnalyzer
-    analyzer.traverse(tree)
-    val liftable: List[DefTree] = (analyzer.valDefsToLift ++ analyzer.nestedMethodsToLift).toList.distinct
-    liftable
-  }
-
-  private class UnsupportedAwaitAnalyzer extends AsyncTraverser {
+  private class UnsupportedAwaitAnalyzer(report: Boolean) extends AsyncTraverser {
     var hasUnsupportedAwaits = false
 
     override def nestedClass(classDef: ClassDef) {
-      val kind = if (classDef.symbol.asClass.isTrait) "trait" else "class"
-      if (!reportUnsupportedAwait(classDef, s"nested $kind")) {
-        // do not allow local class definitions, because of SI-5467 (specific to case classes, though)
-        if (classDef.symbol.asClass.isCaseClass)
-          c.error(classDef.pos, s"Local case class ${classDef.name.decoded} illegal within `async` block")
-      }
+      val kind = if (classDef.symbol.isTrait) "trait" else "class"
+      reportUnsupportedAwait(classDef, s"nested ${kind}")
     }
 
     override def nestedModule(module: ModuleDef) {
-      if (!reportUnsupportedAwait(module, "nested object")) {
-        // local object definitions lead to spurious type errors (because of resetAllAttrs?)
-        c.error(module.pos, s"Local object ${module.name.decoded} illegal within `async` block")
-      }
+      reportUnsupportedAwait(module, "nested object")
     }
 
-    override def nestedMethod(module: DefDef) {
-      reportUnsupportedAwait(module, "nested method")
+    override def nestedMethod(defDef: DefDef) {
+      reportUnsupportedAwait(defDef, "nested method")
     }
 
     override def byNameArgument(arg: Tree) {
@@ -82,9 +59,10 @@ private[async] final case class AsyncAnalysis[C <: Context](c: C, asyncBase: Asy
           reportUnsupportedAwait(tree, "try/catch")
           super.traverse(tree)
         case Return(_)                                        =>
-          c.abort(tree.pos, "return is illegal within a async block")
+          abort(tree.pos, "return is illegal within a async block")
         case ValDef(mods, _, _, _) if mods.hasFlag(Flag.LAZY) =>
-          c.abort(tree.pos, "lazy vals are illegal within an async block")
+          // TODO lift this restriction
+          abort(tree.pos, "lazy vals are illegal within an async block")
         case _                                                =>
           super.traverse(tree)
       }
@@ -106,87 +84,8 @@ private[async] final case class AsyncAnalysis[C <: Context](c: C, asyncBase: Asy
 
     private def reportError(pos: Position, msg: String) {
       hasUnsupportedAwaits = true
-      if (!asyncBase.fallbackEnabled)
-        c.error(pos, msg)
+      if (report)
+        abort(pos, msg)
     }
   }
-
-  private class AsyncDefinitionUseAnalyzer extends AsyncTraverser {
-    private var chunkId = 0
-
-    private def nextChunk() = chunkId += 1
-
-    private var valDefChunkId = Map[Symbol, (ValDef, Int)]()
-
-    val valDefsToLift      : mutable.Set[ValDef] = collection.mutable.Set()
-    val nestedMethodsToLift: mutable.Set[DefDef] = collection.mutable.Set()
-
-    override def nestedMethod(defDef: DefDef) {
-      nestedMethodsToLift += defDef
-      markReferencedVals(defDef)
-    }
-
-    override def function(function: Function) {
-      markReferencedVals(function)
-    }
-
-    override def patMatFunction(tree: Match) {
-      markReferencedVals(tree)
-    }
-
-    private def markReferencedVals(tree: Tree) {
-      tree foreach {
-        case rt: RefTree =>
-          valDefChunkId.get(rt.symbol) match {
-            case Some((vd, defChunkId)) =>
-              valDefsToLift += vd // lift all vals referred to by nested functions.
-            case _                      =>
-          }
-        case _           =>
-      }
-    }
-
-    override def traverse(tree: Tree) = {
-      tree match {
-        case If(cond, thenp, elsep) if tree exists isAwait     =>
-          traverseChunks(List(cond, thenp, elsep))
-        case Match(selector, cases) if tree exists isAwait     =>
-          traverseChunks(selector :: cases)
-        case LabelDef(name, params, rhs) if rhs exists isAwait =>
-          traverseChunks(rhs :: Nil)
-        case Apply(fun, args) if isAwait(fun)                  =>
-          super.traverse(tree)
-          nextChunk()
-        case vd: ValDef                                        =>
-          super.traverse(tree)
-          valDefChunkId += (vd.symbol -> (vd -> chunkId))
-          val isPatternBinder = vd.name.toString.contains(name.bindSuffix)
-          if (isAwait(vd.rhs) || isPatternBinder) valDefsToLift += vd
-        case as: Assign                                        =>
-          if (isAwait(as.rhs)) {
-            assert(as.lhs.symbol != null, "internal error: null symbol for Assign tree:" + as + " " + as.lhs.symbol)
-
-            // TODO test the orElse case, try to remove the restriction.
-            val (vd, defBlockId) = valDefChunkId.getOrElse(as.lhs.symbol, c.abort(as.pos, s"await may only be assigned to a var/val defined in the async block. ${as.lhs} ${as.lhs.symbol}"))
-            valDefsToLift += vd
-          }
-          super.traverse(tree)
-        case rt: RefTree                                       =>
-          valDefChunkId.get(rt.symbol) match {
-            case Some((vd, defChunkId)) if defChunkId != chunkId =>
-              valDefsToLift += vd
-            case _                                               =>
-          }
-          super.traverse(tree)
-        case _                                                 => super.traverse(tree)
-      }
-    }
-
-    private def traverseChunks(trees: List[Tree]) {
-      trees.foreach {
-        t => traverse(t); nextChunk()
-      }
-    }
-  }
-
 }

src/main/scala/scala/async/AsyncAnalysis.scala

@@ -0,0 +1,29 @@
+package scala.async
+
+import scala.tools.nsc.Global
+import scala.tools.nsc.transform.TypingTransformers
+
+object AsyncMacro {
+  def apply(c: reflect.macros.Context, futureSystem0: FutureSystem): AsyncMacro = {
+    import language.reflectiveCalls
+    val powerContext = c.asInstanceOf[c.type {val universe: Global; val callsiteTyper: universe.analyzer.Typer}]
+    new AsyncMacro {
+      val global: powerContext.universe.type = powerContext.universe
+      val callSiteTyper: global.analyzer.Typer = powerContext.callsiteTyper
+      val futureSystem: futureSystem0.type = futureSystem0
+      val futureSystemOps: futureSystem.Ops {val universe: global.type} = futureSystem0.mkOps(global)
+      val macroApplication: global.Tree = c.macroApplication.asInstanceOf[global.Tree]
+    }
+  }
+}
+
+private[async] trait AsyncMacro
+  extends TypingTransformers
+  with AnfTransform with TransformUtils with Lifter
+  with ExprBuilder with AsyncTransform with AsyncAnalysis {
+
+  val global: Global
+  val callSiteTyper: global.analyzer.Typer
+  val macroApplication: global.Tree
+
+}