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.

1 files changed, 213 insertions, 237 deletions

diff --git a/src/main/scala/scala/async/AnfTransform.scala b/src/main/scala/scala/async/AnfTransform.scala
index 5b9901d..275bc49 100644
--- a/src/main/scala/scala/async/AnfTransform.scala
+++ b/src/main/scala/scala/async/AnfTransform.scala
@@ -5,270 +5,246 @@
 
 package scala.async
 
-import scala.reflect.macros.Context
+import scala.tools.nsc.Global
 
-private[async] final case class AnfTransform[C <: Context](c: C) {
+private[async] trait AnfTransform {
+  self: AsyncMacro =>
 
-  import c.universe._
+  import global._
+  import reflect.internal.Flags._
 
-  val utils = TransformUtils[c.type](c)
-
-  import utils._
-
-  def apply(tree: Tree): List[Tree] = {
-    val unique = uniqueNames(tree)
+  def anfTransform(tree: Tree): Block = {
     // Must prepend the () for issue #31.
-    anf.transformToList(Block(List(c.literalUnit.tree), unique))
-  }
+    val block = callSiteTyper.typedPos(tree.pos)(Block(List(Literal(Constant(()))), tree)).setType(tree.tpe)
 
-  private def uniqueNames(tree: Tree): Tree = {
-    new UniqueNames(tree).transform(tree)
+    new SelectiveAnfTransform().transform(block)
   }
 
-  /** Assigns unique names to all definitions in a tree, and adjusts references to use the new name.
-    * Only modifies names that appear more than once in the tree.
-    *
-    * This step is needed to allow us to safely merge blocks during the `inline` transform below.
-    */
-  private final class UniqueNames(tree: Tree) extends Transformer {
-    val repeatedNames: Set[Symbol] = {
-      class DuplicateNameTraverser extends AsyncTraverser {
-        val result = collection.mutable.Buffer[Symbol]()
-
-        override def traverse(tree: Tree) {
-          tree match {
-            case dt: DefTree => result += dt.symbol
-            case _           => super.traverse(tree)
-          }
-        }
-      }
-      val dupNameTraverser = new DuplicateNameTraverser
-      dupNameTraverser.traverse(tree)
-      dupNameTraverser.result.groupBy(x => x.name).filter(_._2.size > 1).values.flatten.toSet[Symbol]
-    }
+  sealed abstract class AnfMode
+
+  case object Anf extends AnfMode
 
-    /** Stepping outside of the public Macro API to call [[scala.reflect.internal.Symbols.Symbol.name_=]] */
-    val symtab = c.universe.asInstanceOf[reflect.internal.SymbolTable]
+  case object Linearizing extends AnfMode
 
-    val renamed = collection.mutable.Set[Symbol]()
+  final class SelectiveAnfTransform extends MacroTypingTransformer {
+    var mode: AnfMode = Anf
 
-    override def transform(tree: Tree): Tree = {
+    def blockToList(tree: Tree): List[Tree] = tree match {
+      case Block(stats, expr) => stats :+ expr
+      case t                  => t :: Nil
+    }
+
+    def listToBlock(trees: List[Tree]): Block = trees match {
+      case trees @ (init :+ last) =>
+        val pos = trees.map(_.pos).reduceLeft(_ union _)
+        Block(init, last).setType(last.tpe).setPos(pos)
+    }
+
+    override def transform(tree: Tree): Block = {
+      def anfLinearize: Block = {
+        val trees: List[Tree] = mode match {
+          case Anf         => anf._transformToList(tree)
+          case Linearizing => linearize._transformToList(tree)
+        }
+        listToBlock(trees)
+      }
       tree match {
-        case defTree: DefTree if repeatedNames(defTree.symbol) =>
-          val trans = super.transform(defTree)
-          val origName = defTree.symbol.name
-          val sym = defTree.symbol.asInstanceOf[symtab.Symbol]
-          val fresh = name.fresh(sym.name.toString)
-          sym.name = origName match {
-            case _: TermName => symtab.newTermName(fresh)
-            case _: TypeName => symtab.newTypeName(fresh)
-          }
-          renamed += trans.symbol
-          val newName = trans.symbol.name
-          trans match {
-            case ValDef(mods, name, tpt, rhs)                    =>
-              treeCopy.ValDef(trans, mods, newName, tpt, rhs)
-            case Bind(name, body)                                =>
-              treeCopy.Bind(trans, newName, body)
-            case DefDef(mods, name, tparams, vparamss, tpt, rhs) =>
-              treeCopy.DefDef(trans, mods, newName, tparams, vparamss, tpt, rhs)
-            case TypeDef(mods, name, tparams, rhs)               =>
-              treeCopy.TypeDef(tree, mods, newName, tparams, transform(rhs))
-            // If we were to allow local classes / objects, we would need to rename here.
-            case ClassDef(mods, name, tparams, impl) =>
-              treeCopy.ClassDef(tree, mods, newName, tparams, transform(impl).asInstanceOf[Template])
-            case ModuleDef(mods, name, impl)         =>
-              treeCopy.ModuleDef(tree, mods, newName, transform(impl).asInstanceOf[Template])
-            case x                                   => super.transform(x)
-          }
-        case Ident(name)                                       =>
-          if (renamed(tree.symbol)) treeCopy.Ident(tree, tree.symbol.name)
-          else tree
-        case Select(fun, name)                                 =>
-          if (renamed(tree.symbol)) {
-            treeCopy.Select(tree, transform(fun), tree.symbol.name)
-          } else super.transform(tree)
-        case tt: TypeTree =>
-          val tt1 = tt.asInstanceOf[symtab.TypeTree]
-          val orig = tt1.original
-          if (orig != null) tt1.setOriginal(transform(orig.asInstanceOf[Tree]).asInstanceOf[symtab.Tree])
-          super.transform(tt)
-        case _                                                 => super.transform(tree)
+        case _: ValDef | _: DefDef | _: Function | _: ClassDef | _: TypeDef =>
+          atOwner(tree.symbol)(anfLinearize)
+        case _: ModuleDef                                                   =>
+          atOwner(tree.symbol.moduleClass orElse tree.symbol)(anfLinearize)
+        case _                                                              =>
+          anfLinearize
       }
     }
-  }
 
-  private object trace {
-    private var indent = -1
-
-    def indentString = "  " * indent
-
-    def apply[T](prefix: String, args: Any)(t: => T): T = {
-      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
+    private object linearize {
+      def transformToList(tree: Tree): List[Tree] = {
+        mode = Linearizing; blockToList(transform(tree))
       }
-    }
-  }
 
-  private object inline {
-    def transformToList(tree: Tree): List[Tree] = trace("inline", tree) {
-      val stats :+ expr = anf.transformToList(tree)
-      expr match {
-        case Apply(fun, args) if isAwait(fun) =>
-          val valDef = defineVal(name.await, expr, tree.pos)
-          stats :+ valDef :+ Ident(valDef.name)
-
-        case If(cond, thenp, elsep) =>
-          // 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(()))
-          } else {
-            val varDef = defineVar(name.ifRes, expr.tpe, tree.pos)
-            def branchWithAssign(orig: Tree) = orig match {
-              case Block(thenStats, thenExpr) => Block(thenStats, Assign(Ident(varDef.name), thenExpr))
-              case _                          => Assign(Ident(varDef.name), orig)
+      def transformToBlock(tree: Tree): Block = listToBlock(transformToList(tree))
+
+      def _transformToList(tree: Tree): List[Tree] = trace(tree) {
+        val stats :+ expr = anf.transformToList(tree)
+        expr match {
+          case Apply(fun, args) if isAwait(fun) =>
+            val valDef = defineVal(name.await, expr, tree.pos)
+            stats :+ valDef :+ gen.mkAttributedStableRef(valDef.symbol)
+
+          case If(cond, thenp, elsep) =>
+            // 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 :+ localTyper.typedPos(expr.pos)(Literal(Constant(())))
+            } else {
+              val varDef = defineVar(name.ifRes, expr.tpe, tree.pos)
+              def branchWithAssign(orig: Tree) = localTyper.typedPos(orig.pos)(
+                orig match {
+                  case Block(thenStats, thenExpr) => Block(thenStats, Assign(Ident(varDef.symbol), thenExpr))
+                  case _                          => Assign(Ident(varDef.symbol), orig)
+                }
+              ).setType(orig.tpe)
+              val ifWithAssign = treeCopy.If(tree, cond, branchWithAssign(thenp), branchWithAssign(elsep))
+              stats :+ varDef :+ ifWithAssign :+ gen.mkAttributedStableRef(varDef.symbol)
+            }
+
+          case Match(scrut, cases) =>
+            // if type of match 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 :+ localTyper.typedPos(expr.pos)(Literal(Constant(())))
             }
-            val ifWithAssign = If(cond, branchWithAssign(thenp), branchWithAssign(elsep))
-            stats :+ varDef :+ ifWithAssign :+ Ident(varDef.name)
-          }
-
-        case Match(scrut, cases) =>
-          // if type of match 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(()))
-          }
-          else {
-            val varDef = defineVar(name.matchRes, expr.tpe, tree.pos)
-            val casesWithAssign = cases map {
-              case cd@CaseDef(pat, guard, Block(caseStats, caseExpr)) =>
-                attachCopy(cd)(CaseDef(pat, guard, Block(caseStats, Assign(Ident(varDef.name), caseExpr))))
-              case cd@CaseDef(pat, guard, body)                       =>
-                attachCopy(cd)(CaseDef(pat, guard, Assign(Ident(varDef.name), body)))
+            else {
+              val varDef = defineVar(name.matchRes, expr.tpe, tree.pos)
+              def typedAssign(lhs: Tree) =
+                localTyper.typedPos(lhs.pos)(Assign(Ident(varDef.symbol), lhs))
+              val casesWithAssign = cases map {
+                case cd@CaseDef(pat, guard, body) =>
+                  val newBody = body match {
+                    case b@Block(caseStats, caseExpr) => treeCopy.Block(b, caseStats, typedAssign(caseExpr))
+                    case _                            => typedAssign(body)
+                  }
+                  treeCopy.CaseDef(cd, pat, guard, newBody)
+              }
+              val matchWithAssign = treeCopy.Match(tree, scrut, casesWithAssign)
+              require(matchWithAssign.tpe != null, matchWithAssign)
+              stats :+ varDef :+ matchWithAssign :+ gen.mkAttributedStableRef(varDef.symbol)
             }
-            val matchWithAssign = attachCopy(tree)(Match(scrut, casesWithAssign))
-            stats :+ varDef :+ matchWithAssign :+ Ident(varDef.name)
-          }
-        case _                   =>
-          stats :+ expr
+          case _                   =>
+            stats :+ expr
+        }
       }
-    }
 
-    def transformToList(trees: List[Tree]): List[Tree] = trees flatMap transformToList
+      private def defineVar(prefix: String, tp: Type, pos: Position): ValDef = {
+        val sym = currOwner.newTermSymbol(name.fresh(prefix), pos, MUTABLE | SYNTHETIC).setInfo(tp)
+        ValDef(sym, gen.mkZero(tp)).setType(NoType).setPos(pos)
+      }
+    }
 
-    def transformToBlock(tree: Tree): Block = transformToList(tree) match {
-      case stats :+ expr => Block(stats, expr)
+    private object trace {
+      private var indent = -1
+
+      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
+        }
+      }
     }
 
-    private def defineVar(prefix: String, tp: Type, pos: Position): ValDef = {
-      val vd = ValDef(Modifiers(Flag.MUTABLE), name.fresh(prefix), TypeTree(tp), defaultValue(tp))
-      vd.setPos(pos)
-      vd
+    private def defineVal(prefix: String, lhs: Tree, pos: Position): ValDef = {
+      val sym = currOwner.newTermSymbol(name.fresh(prefix), pos, SYNTHETIC).setInfo(lhs.tpe)
+      changeOwner(lhs, currentOwner, sym)
+      ValDef(sym, changeOwner(lhs, currentOwner, sym)).setType(NoType).setPos(pos)
     }
-  }
 
-  private def defineVal(prefix: String, lhs: Tree, pos: Position): ValDef = {
-    val vd = ValDef(NoMods, name.fresh(prefix), TypeTree(), lhs)
-    vd.setPos(pos)
-    vd
-  }
+    private object anf {
+      def transformToList(tree: Tree): List[Tree] = {
+        mode = Anf; blockToList(transform(tree))
+      }
 
-  private object anf {
-
-    private[AnfTransform] def transformToList(tree: Tree): List[Tree] = trace("anf", tree) {
-      val containsAwait = tree exists isAwait
-      if (!containsAwait) {
-        List(tree)
-      } else tree match {
-        case Select(qual, sel) =>
-          val stats :+ expr = inline.transformToList(qual)
-          stats :+ attachCopy(tree)(Select(expr, sel).setSymbol(tree.symbol))
-
-        case Applied(fun, targs, argss) if argss.nonEmpty =>
-          // we an assume that no await call appears in a by-name argument position,
-          // this has already been checked.
-          val funStats :+ simpleFun = inline.transformToList(fun)
-          def isAwaitRef(name: Name) = name.toString.startsWith(utils.name.await + "$")
-          val (argStatss, argExprss): (List[List[List[Tree]]], List[List[Tree]]) =
-            mapArgumentss[List[Tree]](fun, argss) {
-              case Arg(expr, byName, _) if byName || isSafeToInline(expr) => (Nil, expr)
-              case Arg(expr@Ident(name), _, _) if isAwaitRef(name)        => (Nil, expr) // not typed, so it eludes the check in `isSafeToInline`
-              case Arg(expr, _, argName)                                 =>
-                inline.transformToList(expr) match {
-                  case stats :+ expr1 =>
-                    val valDef = defineVal(argName, expr1, expr.pos)
-                    (stats :+ valDef, Ident(valDef.name))
-                }
-            }
-            val core = if (targs.isEmpty) simpleFun else TypeApply(simpleFun, targs)
-            val newApply = argExprss.foldLeft(core)(Apply(_, _)).setSymbol(tree.symbol)
-          funStats ++ argStatss.flatten.flatten :+ attachCopy(tree)(newApply)
-        case Block(stats, expr) =>
-          inline.transformToList(stats :+ expr)
-
-        case ValDef(mods, name, tpt, rhs) =>
-          if (rhs exists isAwait) {
-            val stats :+ expr = inline.transformToList(rhs)
-            stats :+ attachCopy(tree)(ValDef(mods, name, tpt, expr).setSymbol(tree.symbol))
-          } else List(tree)
-
-        case Assign(lhs, rhs) =>
-          val stats :+ expr = inline.transformToList(rhs)
-          stats :+ attachCopy(tree)(Assign(lhs, expr))
-
-        case If(cond, thenp, elsep) =>
-          val condStats :+ condExpr = inline.transformToList(cond)
-          val thenBlock = inline.transformToBlock(thenp)
-          val elseBlock = inline.transformToBlock(elsep)
-          // Typechecking with `condExpr` as the condition fails if the condition
-          // contains an await. `ifTree.setType(tree.tpe)` also fails; it seems
-          // we rely on this call to `typeCheck` descending into the branches.
-          // But, we can get away with typechecking a throwaway `If` tree with the
-          // original scrutinee and the new branches, and setting that type on
-          // the real `If` tree.
-          val ifType = c.typeCheck(If(cond, thenBlock, elseBlock)).tpe
-          condStats :+
-            attachCopy(tree)(If(condExpr, thenBlock, elseBlock)).setType(ifType)
-
-        case Match(scrut, cases) =>
-          val scrutStats :+ scrutExpr = inline.transformToList(scrut)
-          val caseDefs = cases map {
-            case CaseDef(pat, guard, body) =>
-              // extract local variables for all names bound in `pat`, and rewrite `body`
-              // to refer to these.
-              // TODO we can move this into ExprBuilder once we get rid of `AsyncDefinitionUseAnalyzer`.
-              val block = inline.transformToBlock(body)
-              val (valDefs, mappings) = (pat collect {
-                case b@Bind(name, _) =>
-                  val newName = newTermName(utils.name.fresh(name.toTermName + utils.name.bindSuffix))
-                  val vd = ValDef(NoMods, newName, TypeTree(), Ident(b.symbol))
-                  (vd, (b.symbol, newName))
-              }).unzip
-              val Block(stats1, expr1) = utils.substituteNames(block, mappings.toMap).asInstanceOf[Block]
-              attachCopy(tree)(CaseDef(pat, guard, Block(valDefs ++ stats1, expr1)))
-          }
-          // Refer to comments the translation of `If` above.
-          val matchType = c.typeCheck(Match(scrut, caseDefs)).tpe
-          val typedMatch = attachCopy(tree)(Match(scrutExpr, caseDefs)).setType(tree.tpe)
-          scrutStats :+ typedMatch
-
-        case LabelDef(name, params, rhs) =>
-          List(LabelDef(name, params, Block(inline.transformToList(rhs), Literal(Constant(())))).setSymbol(tree.symbol))
-
-        case TypeApply(fun, targs) =>
-          val funStats :+ simpleFun = inline.transformToList(fun)
-          funStats :+ attachCopy(tree)(TypeApply(simpleFun, targs).setSymbol(tree.symbol))
-
-        case _ =>
+      def _transformToList(tree: Tree): List[Tree] = trace(tree) {
+        val containsAwait = tree exists isAwait
+        if (!containsAwait) {
           List(tree)
+        } else tree match {
+          case Select(qual, sel) =>
+            val stats :+ expr = linearize.transformToList(qual)
+            stats :+ treeCopy.Select(tree, expr, sel)
+
+          case treeInfo.Applied(fun, targs, argss) if argss.nonEmpty =>
+            // we an assume that no await call appears in a by-name argument position,
+            // this has already been checked.
+            val funStats :+ simpleFun = linearize.transformToList(fun)
+            def isAwaitRef(name: Name) = name.toString.startsWith(AnfTransform.this.name.await + "$")
+            val (argStatss, argExprss): (List[List[List[Tree]]], List[List[Tree]]) =
+              mapArgumentss[List[Tree]](fun, argss) {
+                case Arg(expr, byName, _) if byName /*|| isPure(expr) TODO */ => (Nil, expr)
+                case Arg(expr@Ident(name), _, _) if isAwaitRef(name)          => (Nil, expr) // TODO needed? // not typed, so it eludes the check in `isSafeToInline`
+                case Arg(expr, _, argName)                                    =>
+                  linearize.transformToList(expr) match {
+                    case stats :+ expr1 =>
+                      val valDef = defineVal(argName, expr1, expr1.pos)
+                      require(valDef.tpe != null, valDef)
+                      val stats1 = stats :+ valDef
+                      //stats1.foreach(changeOwner(_, currentOwner, currentOwner.owner))
+                      (stats1, gen.stabilize(gen.mkAttributedIdent(valDef.symbol)))
+                  }
+              }
+            val applied = treeInfo.dissectApplied(tree)
+            val core = if (targs.isEmpty) simpleFun else treeCopy.TypeApply(applied.callee, simpleFun, targs)
+            val newApply = argExprss.foldLeft(core)(Apply(_, _)).setSymbol(tree.symbol)
+            val typedNewApply = localTyper.typedPos(tree.pos)(newApply).setType(tree.tpe)
+            funStats ++ argStatss.flatten.flatten :+ typedNewApply
+          case Block(stats, expr)                                    =>
+            (stats :+ expr).flatMap(linearize.transformToList)
+
+          case ValDef(mods, name, tpt, rhs) =>
+            if (rhs exists isAwait) {
+              val stats :+ expr = atOwner(currOwner.owner)(linearize.transformToList(rhs))
+              stats.foreach(changeOwner(_, currOwner, currOwner.owner))
+              stats :+ treeCopy.ValDef(tree, mods, name, tpt, expr)
+            } else List(tree)
+
+          case Assign(lhs, rhs) =>
+            val stats :+ expr = linearize.transformToList(rhs)
+            stats :+ treeCopy.Assign(tree, lhs, expr)
+
+          case If(cond, thenp, elsep) =>
+            val condStats :+ condExpr = linearize.transformToList(cond)
+            val thenBlock = linearize.transformToBlock(thenp)
+            val elseBlock = linearize.transformToBlock(elsep)
+            // Typechecking with `condExpr` as the condition fails if the condition
+            // contains an await. `ifTree.setType(tree.tpe)` also fails; it seems
+            // we rely on this call to `typeCheck` descending into the branches.
+            // But, we can get away with typechecking a throwaway `If` tree with the
+            // original scrutinee and the new branches, and setting that type on
+            // the real `If` tree.
+            val iff = treeCopy.If(tree, condExpr, thenBlock, elseBlock)
+            condStats :+ iff
+
+          case Match(scrut, cases) =>
+            val scrutStats :+ scrutExpr = linearize.transformToList(scrut)
+            val caseDefs = cases map {
+              case CaseDef(pat, guard, body) =>
+                // extract local variables for all names bound in `pat`, and rewrite `body`
+                // to refer to these.
+                // TODO we can move this into ExprBuilder once we get rid of `AsyncDefinitionUseAnalyzer`.
+                val block = linearize.transformToBlock(body)
+                val (valDefs, mappings) = (pat collect {
+                  case b@Bind(name, _) =>
+                    val vd = defineVal(name.toTermName + AnfTransform.this.name.bindSuffix, gen.mkAttributedStableRef(b.symbol), b.pos)
+                    (vd, (b.symbol, vd.symbol))
+                }).unzip
+                val (from, to) = mappings.unzip
+                val b@Block(stats1, expr1) = block.substituteSymbols(from, to).asInstanceOf[Block]
+                val newBlock = treeCopy.Block(b, valDefs ++ stats1, expr1)
+                treeCopy.CaseDef(tree, pat, guard, newBlock)
+            }
+            // Refer to comments the translation of `If` above.
+            val typedMatch = treeCopy.Match(tree, scrutExpr, caseDefs)
+            scrutStats :+ typedMatch
+
+          case LabelDef(name, params, rhs) =>
+            List(LabelDef(name, params, Block(linearize.transformToList(rhs), Literal(Constant(())))).setSymbol(tree.symbol))
+
+          case TypeApply(fun, targs) =>
+            val funStats :+ simpleFun = linearize.transformToList(fun)
+            funStats :+ treeCopy.TypeApply(tree, simpleFun, targs)
+
+          case _ =>
+            List(tree)
+        }
       }
     }
   }