change intermidiate representation of for loop enumerators

Encode values into real trees rather than non-tree case classes.
This is needed for re-usability of desugaring code between quasiquotes
and parser.

7 files changed, 107 insertions, 55 deletions

diff --git a/src/compiler/scala/tools/nsc/ast/TreeGen.scala b/src/compiler/scala/tools/nsc/ast/TreeGen.scala
index 28b127698f..d4ac21a6b8 100644
--- a/src/compiler/scala/tools/nsc/ast/TreeGen.scala
+++ b/src/compiler/scala/tools/nsc/ast/TreeGen.scala
@@ -53,13 +53,6 @@ abstract class TreeGen extends scala.reflect.internal.TreeGen with TreeDSL {
     NewFromConstructor(constructor, expr)
   }
 
-  // annotate the expression with @unchecked
-  def mkUnchecked(expr: Tree): Tree = atPos(expr.pos) {
-    // This can't be "Annotated(New(UncheckedClass), expr)" because annotations
-    // are very picky about things and it crashes the compiler with "unexpected new".
-    Annotated(New(scalaDot(UncheckedClass.name), Nil), expr)
-  }
-
   // Builds a tree of the form "{ lhs = rhs ; lhs  }"
   def mkAssignAndReturn(lhs: Symbol, rhs: Tree): Tree = {
     def lhsRef = if (lhs.owner.isClass) Select(This(lhs.owner), lhs) else Ident(lhs)
@@ -263,7 +256,4 @@ abstract class TreeGen extends scala.reflect.internal.TreeGen with TreeDSL {
     val stats1 = if (stats.isEmpty) List(Literal(Constant(()))) else stats
     mkNew(Nil, noSelfType, stats1, NoPosition, NoPosition)
   }
-
-  def mkSyntheticParam(pname: TermName) =
-    ValDef(Modifiers(PARAM | SYNTHETIC), pname, TypeTree(), EmptyTree)
 }
diff --git a/src/compiler/scala/tools/nsc/ast/parser/Parsers.scala b/src/compiler/scala/tools/nsc/ast/parser/Parsers.scala
index 0bf4d5426e..4641422132 100644
--- a/src/compiler/scala/tools/nsc/ast/parser/Parsers.scala
+++ b/src/compiler/scala/tools/nsc/ast/parser/Parsers.scala
@@ -1394,7 +1394,7 @@ self =>
           newLinesOpt()
           if (in.token == YIELD) {
             in.nextToken()
-            makeForYield(enums, expr())
+            makeFor(enums, Yield(expr()))
           } else {
             makeFor(enums, expr())
           }
@@ -1700,22 +1700,25 @@ self =>
      *                |  val Pattern1 `=' Expr
      *  }}}
      */
-    def enumerators(): List[Enumerator] = {
-      val enums = new ListBuffer[Enumerator]
-      generator(enums, eqOK = false)
+    def enumerators(): List[Tree] = {
+      val enums = new ListBuffer[Tree]
+      enums ++= enumerator(isFirst = true)
       while (isStatSep) {
         in.nextToken()
-        if (in.token == IF) enums += makeFilter(in.offset, guard())
-        else generator(enums, eqOK = true)
+        enums ++= enumerator(isFirst = false)
       }
       enums.toList
     }
 
+    def enumerator(isFirst: Boolean, allowNestedIf: Boolean = true): List[Tree] =
+      if (in.token == IF && !isFirst) makeFilter(in.offset, guard()) :: Nil
+      else generator(!isFirst, allowNestedIf)
+
     /** {{{
      *  Generator ::= Pattern1 (`<-' | `=') Expr [Guard]
      *  }}}
      */
-    def generator(enums: ListBuffer[Enumerator], eqOK: Boolean) {
+    def generator(eqOK: Boolean, allowNestedIf: Boolean = true): List[Tree] = {
       val start  = in.offset
       val hasVal = in.token == VAL
       if (hasVal)
@@ -1733,10 +1736,19 @@ self =>
       if (hasEq && eqOK) in.nextToken()
       else accept(LARROW)
       val rhs = expr()
-      enums += makeGenerator(r2p(start, point, in.lastOffset max start), pat, hasEq, rhs)
-      // why max above? IDE stress tests have shown that lastOffset could be less than start,
+
+      def loop(): List[Tree] =
+        if (in.token != IF) Nil
+        else makeFilter(in.offset, guard()) :: loop()
+
+      val tail =
+        if (allowNestedIf) loop()
+        else Nil
+
+      // why max? IDE stress tests have shown that lastOffset could be less than start,
       // I guess this happens if instead if a for-expression we sit on a closing paren.
-      while (in.token == IF) enums += makeFilter(in.offset, guard())
+      val genPos = r2p(start, point, in.lastOffset max start)
+      makeGenerator(genPos, pat, hasEq, rhs) :: tail
     }
 
     def makeFilter(start: Offset, tree: Tree) = Filter(r2p(start, tree.pos.point, tree.pos.end), tree)
diff --git a/src/compiler/scala/tools/nsc/ast/parser/TreeBuilder.scala b/src/compiler/scala/tools/nsc/ast/parser/TreeBuilder.scala
index 5af279a62a..227d54f072 100644
--- a/src/compiler/scala/tools/nsc/ast/parser/TreeBuilder.scala
+++ b/src/compiler/scala/tools/nsc/ast/parser/TreeBuilder.scala
@@ -212,7 +212,7 @@ abstract class TreeBuilder {
   }
 
   /** Create tree for for-comprehension generator <val pat0 <- rhs0> */
-  def makeGenerator(pos: Position, pat: Tree, valeq: Boolean, rhs: Tree): Enumerator = {
+  def makeGenerator(pos: Position, pat: Tree, valeq: Boolean, rhs: Tree): Tree = {
     val pat1 = patvarTransformer.transform(pat)
     val rhs1 =
       if (valeq || treeInfo.isVarPatternDeep(pat)) rhs
@@ -228,10 +228,54 @@ abstract class TreeBuilder {
   def makeSyntheticTypeParam(pname: TypeName, bounds: Tree) =
     TypeDef(Modifiers(DEFERRED | SYNTHETIC), pname, Nil, bounds)
 
-  abstract class Enumerator { def pos: Position }
-  case class ValFrom(pos: Position, pat: Tree, rhs: Tree) extends Enumerator
-  case class ValEq(pos: Position, pat: Tree, rhs: Tree) extends Enumerator
-  case class Filter(pos: Position, test: Tree) extends Enumerator
+  object ValFrom {
+    def apply(pos: Position, pat: Tree, rhs: Tree): Tree =
+      Apply(Ident(nme.LARROWkw).updateAttachment(ForAttachment),
+        List(pat, rhs)).setPos(pos)
+
+    def unapply(tree: Tree): Option[(Position, Tree, Tree)] = tree match {
+      case app @ Apply(id @ Ident(nme.LARROWkw), List(pat, rhs))
+        if id.hasAttachment[ForAttachment.type] =>
+        Some((app.pos, pat, rhs))
+      case _ => None
+    }
+  }
+
+  object ValEq {
+    def apply(pos: Position, pat: Tree, rhs: Tree): Tree =
+      Assign(pat, rhs).updateAttachment(ForAttachment).setPos(pos)
+
+    def unapply(tree: Tree): Option[(Position, Tree, Tree)] = tree match {
+      case ass @ Assign(pat, rhs)
+        if tree.hasAttachment[ForAttachment.type] =>
+        Some((ass.pos, pat, rhs))
+      case _ => None
+    }
+  }
+
+  object Filter {
+    def apply(pos: Position, tree: Tree) =
+      Apply(Ident(nme.IFkw).updateAttachment(ForAttachment), List(tree)).setPos(pos)
+
+    def unapply(tree: Tree): Option[(Position, Tree)] = tree match {
+      case app @ Apply(id @ Ident(nme.IFkw), List(cond))
+        if id.hasAttachment[ForAttachment.type] =>
+        Some((app.pos, cond))
+      case _ => None
+    }
+  }
+
+  object Yield {
+    def apply(tree: Tree): Tree =
+      Apply(Ident(nme.YIELDkw).updateAttachment(ForAttachment), List(tree))
+
+    def unapply(tree: Tree): Option[Tree] = tree match {
+      case Apply(id @ Ident(nme.YIELDkw), List(tree))
+        if id.hasAttachment[ForAttachment.type] =>
+        Some(tree)
+      case _  => None
+    }
+  }
 
   /** Create tree for for-comprehension <for (enums) do body> or
   *   <for (enums) yield body> where mapName and flatMapName are chosen
@@ -281,7 +325,11 @@ abstract class TreeBuilder {
   *  @param enums        The enumerators in the for expression
   *  @param body          The body of the for expression
   */
-  private def makeFor(mapName: TermName, flatMapName: TermName, enums: List[Enumerator], body: Tree): Tree = {
+  def makeFor(enums: List[Tree], sugarBody: Tree)(implicit fresh: FreshNameCreator): Tree = {
+    val (mapName, flatMapName, body) = sugarBody match {
+      case Yield(tree) => (nme.map, nme.flatMap, tree)
+      case _           => (nme.foreach, nme.foreach, sugarBody)
+    }
 
     /* make a closure pat => body.
      * The closure is assigned a transparent position with the point at pos.point and
@@ -326,19 +374,16 @@ abstract class TreeBuilder {
       r2p(genpos.start, genpos.point, end)
     }
 
-//    val result =
     enums match {
       case ValFrom(pos, pat, rhs) :: Nil =>
         makeCombination(closurePos(pos), mapName, rhs, pat, body)
       case ValFrom(pos, pat, rhs) :: (rest @ (ValFrom(_,  _, _) :: _)) =>
         makeCombination(closurePos(pos), flatMapName, rhs, pat,
-                        makeFor(mapName, flatMapName, rest, body))
+                        makeFor(rest, sugarBody))
       case ValFrom(pos, pat, rhs) :: Filter(_, test) :: rest =>
-        makeFor(mapName, flatMapName,
-                ValFrom(pos, pat, makeCombination(rhs.pos union test.pos, nme.withFilter, rhs, pat.duplicate, test)) :: rest,
-                body)
+        makeFor(ValFrom(pos, pat, makeCombination(rhs.pos union test.pos, nme.withFilter, rhs, pat.duplicate, test)) :: rest, sugarBody)
       case ValFrom(pos, pat, rhs) :: rest =>
-        val valeqs = rest.take(definitions.MaxTupleArity - 1).takeWhile(_.isInstanceOf[ValEq])
+        val valeqs = rest.take(definitions.MaxTupleArity - 1).takeWhile { ValEq.unapply(_).nonEmpty }
         assert(!valeqs.isEmpty)
         val rest1 = rest.drop(valeqs.length)
         val pats = valeqs map { case ValEq(_, pat, _) => pat }
@@ -347,27 +392,17 @@ abstract class TreeBuilder {
         val defpats = pats map makeBind
         val pdefs = (defpats, rhss).zipped flatMap makePatDef
         val ids = (defpat1 :: defpats) map makeValue
-        val rhs1 = makeForYield(
+        val rhs1 = makeFor(
           List(ValFrom(pos, defpat1, rhs)),
-          Block(pdefs, atPos(wrappingPos(ids)) { makeTupleTerm(ids) }) setPos wrappingPos(pdefs))
+          Yield(Block(pdefs, atPos(wrappingPos(ids)) { makeTupleTerm(ids) }) setPos wrappingPos(pdefs)))
         val allpats = (pat :: pats) map (_.duplicate)
         val vfrom1 = ValFrom(r2p(pos.start, pos.point, rhs1.pos.end), atPos(wrappingPos(allpats)) { makeTupleTerm(allpats) } , rhs1)
-        makeFor(mapName, flatMapName, vfrom1 :: rest1, body)
+        makeFor(vfrom1 :: rest1, sugarBody)
       case _ =>
         EmptyTree //may happen for erroneous input
     }
-//    println("made for "+result)
-//    result
   }
 
-  /** Create tree for for-do comprehension <for (enums) body> */
-  def makeFor(enums: List[Enumerator], body: Tree): Tree =
-    makeFor(nme.foreach, nme.foreach, enums, body)
-
-  /** Create tree for for-yield comprehension <for (enums) yield body> */
-  def makeForYield(enums: List[Enumerator], body: Tree): Tree =
-    makeFor(nme.map, nme.flatMap, enums, body)
-
   /** Create tree for a pattern alternative */
   def makeAlternative(ts: List[Tree]): Tree = {
     def alternatives(t: Tree): List[Tree] = t match {
@@ -416,11 +451,11 @@ abstract class TreeBuilder {
   }
 
   /** Create tree for pattern definition <val pat0 = rhs> */
-  def makePatDef(pat: Tree, rhs: Tree): List[Tree] =
+  def makePatDef(pat: Tree, rhs: Tree)(implicit fresh: FreshNameCreator): List[Tree] =
     makePatDef(Modifiers(0), pat, rhs)
 
   /** Create tree for pattern definition <mods val pat0 = rhs> */
-  def makePatDef(mods: Modifiers, pat: Tree, rhs: Tree): List[Tree] = matchVarPattern(pat) match {
+  def makePatDef(mods: Modifiers, pat: Tree, rhs: Tree)(implicit fresh: FreshNameCreator): List[Tree] = matchVarPattern(pat) match {
     case Some((name, tpt)) =>
       List(atPos(pat.pos union rhs.pos) {
         ValDef(mods, name.toTermName, tpt, rhs)
@@ -460,7 +495,7 @@ abstract class TreeBuilder {
           rhs1,
           List(
             atPos(pat1.pos) {
-              CaseDef(pat1, EmptyTree, makeTupleTerm(vars map (_._1) map Ident.apply))
+              CaseDef(pat1, EmptyTree, gen.mkTuple(vars map (_._1) map Ident.apply))
             }
           ))
       }
diff --git a/src/reflect/scala/reflect/internal/StdAttachments.scala b/src/reflect/scala/reflect/internal/StdAttachments.scala
index fe4438015d..d37bacb462 100644
--- a/src/reflect/scala/reflect/internal/StdAttachments.scala
+++ b/src/reflect/scala/reflect/internal/StdAttachments.scala
@@ -31,4 +31,8 @@ trait StdAttachments {
   /** When present, indicates that the host `Ident` has been created from a backquoted identifier.
    */
   case object BackquotedIdentifierAttachment extends PlainAttachment
+
+  /** Identifies trees are either result or intermidiate value of for loop desugaring.
+   */
+  case object ForAttachment extends PlainAttachment
 }
diff --git a/src/reflect/scala/reflect/internal/StdNames.scala b/src/reflect/scala/reflect/internal/StdNames.scala
index 02f22a16f6..bc9c46a141 100644
--- a/src/reflect/scala/reflect/internal/StdNames.scala
+++ b/src/reflect/scala/reflect/internal/StdNames.scala
@@ -229,6 +229,7 @@ trait StdNames {
     final val Serializable: NameType    = "Serializable"
     final val Singleton: NameType       = "Singleton"
     final val Throwable: NameType       = "Throwable"
+    final val unchecked: NameType       = "unchecked"
 
     final val api: NameType                 = "api"
     final val Annotation: NameType          = "Annotation"
diff --git a/src/reflect/scala/reflect/internal/TreeGen.scala b/src/reflect/scala/reflect/internal/TreeGen.scala
index c87eeea8e0..b75368717f 100644
--- a/src/reflect/scala/reflect/internal/TreeGen.scala
+++ b/src/reflect/scala/reflect/internal/TreeGen.scala
@@ -4,6 +4,7 @@ package internal
 
 import Flags._
 import util._
+import scala.collection.mutable.ListBuffer
 
 abstract class TreeGen extends macros.TreeBuilder {
   val global: SymbolTable
@@ -311,7 +312,7 @@ abstract class TreeGen extends macros.TreeBuilder {
   }
 
   def mkSeqApply(arg: Tree): Apply = {
-    val factory = Select(gen.mkAttributedRef(SeqModule), nme.apply)
+    val factory = Select(mkAttributedRef(SeqModule), nme.apply)
     Apply(factory, List(arg))
   }
 
@@ -448,17 +449,15 @@ abstract class TreeGen extends macros.TreeBuilder {
     else Block(stats.init, stats.last)
 
   def mkTreeOrBlock(stats: List[Tree]) = stats match {
-    case Nil => EmptyTree
+    case Nil         => EmptyTree
     case head :: Nil => head
-    case _ => gen.mkBlock(stats)
+    case _           => mkBlock(stats)
   }
 
   /** Create a tree representing an assignment <lhs = rhs> */
   def mkAssign(lhs: Tree, rhs: Tree): Tree = lhs match {
-    case Apply(fn, args) =>
-      Apply(atPos(fn.pos)(Select(fn, nme.update)), args :+ rhs)
-    case _ =>
-      Assign(lhs, rhs)
+    case Apply(fn, args) => Apply(atPos(fn.pos)(Select(fn, nme.update)), args :+ rhs)
+    case _               => Assign(lhs, rhs)
   }
 
   def mkPackageObject(defn: ModuleDef, pidPos: Position = NoPosition, pkgPos: Position = NoPosition) = {
@@ -466,4 +465,14 @@ abstract class TreeGen extends macros.TreeBuilder {
     val pid    = atPos(pidPos)(Ident(defn.name))
     atPos(pkgPos)(PackageDef(pid, module :: Nil))
   }
+
+  // annotate the expression with @unchecked
+  def mkUnchecked(expr: Tree): Tree = atPos(expr.pos) {
+    // This can't be "Annotated(New(UncheckedClass), expr)" because annotations
+    // are very picky about things and it crashes the compiler with "unexpected new".
+    Annotated(New(scalaDot(tpnme.unchecked), Nil), expr)
+  }
+
+  def mkSyntheticParam(pname: TermName) =
+    ValDef(Modifiers(PARAM | SYNTHETIC), pname, TypeTree(), EmptyTree)
 }
diff --git a/src/reflect/scala/reflect/runtime/JavaUniverseForce.scala b/src/reflect/scala/reflect/runtime/JavaUniverseForce.scala
index 26091b84a1..711456f6c7 100644
--- a/src/reflect/scala/reflect/runtime/JavaUniverseForce.scala
+++ b/src/reflect/scala/reflect/runtime/JavaUniverseForce.scala
@@ -55,6 +55,7 @@ trait JavaUniverseForce { self: runtime.JavaUniverse  =>
     this.FixedMirrorTypeCreator
     this.CompoundTypeTreeOriginalAttachment
     this.BackquotedIdentifierAttachment
+    this.ForAttachment
     this.noPrint
     this.typeDebug
     // inaccessible: this.maxFree