Refactored TreeTypeMap

Goes into a separate source files. Several simplifying refactorings.

7 files changed, 155 insertions, 158 deletions

diff --git a/src/dotty/tools/dotc/ast/TreeInfo.scala b/src/dotty/tools/dotc/ast/TreeInfo.scala
index 881cc3f6e..51e1ff16f 100644
--- a/src/dotty/tools/dotc/ast/TreeInfo.scala
+++ b/src/dotty/tools/dotc/ast/TreeInfo.scala
@@ -420,6 +420,17 @@ trait TypedTreeInfo extends TreeInfo[Type] { self: Trees.Instance[Type] =>
     }
   }
 
+  /** The variables defined by a pattern, in reverse order of their appearance. */
+  def patVars(tree: Tree)(implicit ctx: Context): List[Symbol] = {
+    val acc = new TreeAccumulator[List[Symbol]] {
+      def apply(syms: List[Symbol], tree: Tree) = tree match {
+        case Bind(_, body) => apply(tree.symbol :: syms, body)
+        case _ => foldOver(syms, tree)
+      }
+    }
+    acc(Nil, tree)
+  }
+
   /** Is this pattern node a catch-all or type-test pattern? */
   def isCatchCase(cdef: CaseDef)(implicit ctx: Context) = cdef match {
     case CaseDef(Typed(Ident(nme.WILDCARD), tpt), EmptyTree, _) =>
@@ -438,6 +449,10 @@ trait TypedTreeInfo extends TreeInfo[Type] { self: Trees.Instance[Type] =>
       false
   }
 
+  /** The symbols defined locally in a statement list */
+  def localSyms(stats: List[Tree])(implicit ctx: Context): List[Symbol] =
+    for (stat <- stats if stat.isDef && stat.symbol.exists) yield stat.symbol
+
   /** If `tree` is a DefTree, the symbol defined by it, otherwise NoSymbol */
   def definedSym(tree: Tree)(implicit ctx: Context): Symbol =
     if (tree.isDef) tree.symbol else NoSymbol
diff --git a/src/dotty/tools/dotc/ast/TreeTypeMap.scala b/src/dotty/tools/dotc/ast/TreeTypeMap.scala
new file mode 100644
index 000000000..77ed24511
--- /dev/null
+++ b/src/dotty/tools/dotc/ast/TreeTypeMap.scala
@@ -0,0 +1,120 @@
+package dotty.tools
+package dotc
+package ast
+
+import core._
+import Types._, Contexts._, Constants._, Names._, Flags._
+import SymDenotations._, Symbols._, Annotations._, Trees._, Symbols._
+import Denotations._, Decorators._
+
+/** A map that applies three functions and a substitution together to a tree and
+ *  makes sure they are coordinated so that the result is well-typed. The functions are
+ *  @param  typeMap  A function from Type to Type that gets applied to the
+ *                   type of every tree node and to all locally defined symbols,
+ *                   followed by the substitution [substFrom := substTo].
+ *  @param ownerMap  A function that translates owners of top-level local symbols
+ *                   defined in the mapped tree.
+ *  @param treeMap   A transformer that translates all encountered subtrees in
+ *                   prefix traversal order.
+ *  @param substFrom The symbols that need to be substituted
+ *  @param substTo   The substitution targets
+ *
+ *  The reason the substitution is broken out from the rest of the type map is
+ *  that all symbols have to be substituted at the same time. If we do not do this,
+ *  we risk data races on named types. Example: Say we have `outer#1.inner#2` and we
+ *  have two substitutons S1 = [outer#1 := outer#3], S2 = [inner#2 := inner#4] where
+ *  hashtags precede symbol ids. If we do S1 first, we get outer#2.inner#3. If we then
+ *  do S2 we get outer#2.inner#4. But that means that the named type outer#2.inner
+ *  gets two different denotations in the same period. Hence, if -YnoDoubleBindings is
+ *  set, we would get a data race assertion error.
+ */
+final class TreeTypeMap(
+  val typeMap: Type => Type = IdentityTypeMap,
+  val ownerMap: Symbol => Symbol = identity _,
+  val treeMap: tpd.Tree => tpd.Tree = identity _,
+  val substFrom: List[Symbol] = Nil,
+  val substTo: List[Symbol] = Nil)(implicit ctx: Context) extends tpd.TreeMap {
+  import tpd._
+
+  def mapType(tp: Type) = typeMap(tp).substSym(substFrom, substTo)
+
+  override def transform(tree: tpd.Tree)(implicit ctx: Context): tpd.Tree = treeMap(tree) match {
+    case impl @ Template(constr, parents, self, body) =>
+      val tmap = withMappedSyms(impl.symbol :: impl.constr.symbol :: Nil)
+      val constr1 = tmap.transformSub(constr)
+      val parents1 = parents mapconserve transform
+      var self1 = transformDefs(self :: Nil)._2.head
+      val body1 = tmap.transformStats(body)
+      cpy.Template(impl)(constr1, parents1, self1, body1).withType(tmap.mapType(impl.tpe))
+    case tree1 =>
+      tree1.withType(mapType(tree1.tpe)) match {
+        case ddef @ DefDef(mods, name, tparams, vparamss, tpt, rhs) =>
+          val (tmap1, tparams1) = transformDefs(ddef.tparams)
+          val (tmap2, vparamss1) = tmap1.transformVParamss(vparamss)
+          cpy.DefDef(ddef)(mods, name, tparams1, vparamss1, tmap2.transform(tpt), tmap2.transform(rhs))
+        case blk @ Block(stats, expr) =>
+          val (tmap1, stats1) = transformDefs(stats)
+          val expr1 = tmap1.transform(expr)
+          cpy.Block(blk)(stats1, expr1)
+        case cdef @ CaseDef(pat, guard, rhs) =>
+          val tmap = withMappedSyms(patVars(pat))
+          val pat1 = tmap.transform(pat)
+          val guard1 = tmap.transform(guard)
+          val rhs1 = tmap.transform(rhs)
+          cpy.CaseDef(cdef)(pat1, guard1, rhs1)
+        case tree1 =>
+          super.transform(tree1)
+      }
+  }
+
+  override def transformStats(trees: List[tpd.Tree])(implicit ctx: Context) =
+    transformDefs(trees)._2
+
+  private def transformDefs[TT <: tpd.Tree](trees: List[TT])(implicit ctx: Context): (TreeTypeMap, List[TT]) = {
+    val tmap = withMappedSyms(tpd.localSyms(trees))
+    (tmap, tmap.transformSub(trees))
+  }
+
+  private def transformVParamss(vparamss: List[List[ValDef]]): (TreeTypeMap, List[List[ValDef]]) = vparamss match {
+    case vparams :: rest =>
+      val (tmap1, vparams1) = transformDefs(vparams)
+      val (tmap2, vparamss2) = tmap1.transformVParamss(rest)
+      (tmap2, vparams1 :: vparamss2)
+    case nil =>
+      (this, vparamss)
+  }
+
+  def apply[ThisTree <: tpd.Tree](tree: ThisTree): ThisTree = transform(tree).asInstanceOf[ThisTree]
+
+  def apply(annot: Annotation): Annotation = {
+    val tree1 = apply(annot.tree)
+    if (tree1 eq annot.tree) annot else ConcreteAnnotation(tree1)
+  }
+
+  /** The current tree map composed with a substitution [from -> to] */
+  def withSubstitution(from: List[Symbol], to: List[Symbol]): TreeTypeMap =
+    if (from eq to) this
+    else {
+      // assert that substitution stays idempotent, assuming its parts are
+      assert(!from.exists(substTo contains _))
+      assert(!to.exists(substFrom contains _))
+      new TreeTypeMap(
+        typeMap,
+        ownerMap andThen { sym =>
+          val idx = from.indexOf(sym)
+          if (idx >= 0) to(idx) else sym
+        },
+        treeMap,
+        from ++ substFrom,
+        to ++ substTo)
+    }
+
+  /** Apply `typeMap` and `ownerMap` to given symbols `syms`
+   *  and return a treemap that contains the substitution
+   *  between original and mapped symbols.
+   */
+  def withMappedSyms(syms: List[Symbol]): TreeTypeMap = {
+    val mapped = ctx.mapSymbols(syms, this)
+    withSubstitution(syms, mapped)
+  }
+}
diff --git a/src/dotty/tools/dotc/ast/tpd.scala b/src/dotty/tools/dotc/ast/tpd.scala
index 802743736..56cd35344 100644
--- a/src/dotty/tools/dotc/ast/tpd.scala
+++ b/src/dotty/tools/dotc/ast/tpd.scala
@@ -546,135 +546,6 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
     def tpes: List[Type] = xs map (_.tpe)
   }
 
-  /** A map that applies three functions together to a tree and makes sure
-   *  they are coordinated so that the result is well-typed. The functions are
-   *  @param  typeMap  A function from Type to type that gets applied to the
-   *                   type of every tree node and to all locally defined symbols,
-   *                   followed by the substitution [substFrom := substTo].
-   *  @param ownerMap  A function that translates owners of top-level local symbols
-   *                   defined in the mapped tree.
-   *  @param treeMap   A transformer that translates all encountered subtrees in
-   *                   prefix traversal order.
-   *  @param substFrom The symbols that need to be substituted
-   *  @param substTo   The substitution targets
-   *
-   *  The reason the substitution is broken out form the rest of the type map is
-   *  that all symbols have to be substituted at the same time. If we do not do this,
-   *  we risk data races on named types. Example: Say we have `outer#1.inner#2` and we
-   *  have two substitutons S1 = [outer#1 := outer#3], S2 = [inner#2 := inner#4] where
-   *  hashtags precede symbol ids. If we do S1 first, we get outer#2.inner#3. If we then
-   *  do S2 we get outer#2.inner#4. But that means that the named type outer#2.inner
-   *  gets two different denotations in the same period. Hence, if -YnoDoubleBindings is
-   *  set, we would get a data race assertion error.
-   */
-  final class TreeTypeMap(
-      val typeMap: Type => Type = IdentityTypeMap,
-      val ownerMap: Symbol => Symbol = identity _,
-      val treeMap: Tree => Tree = identity _,
-      val substFrom: List[Symbol] = Nil,
-      val substTo: List[Symbol] = Nil)(implicit ctx: Context) extends TreeMap {
-
-    def mapType(tp: Type) = typeMap(tp).substSym(substFrom, substTo)
-
-    override def transform(tree: tpd.Tree)(implicit ctx: Context): tpd.Tree = {
-      val tree1 = treeMap(tree)
-      tree1 match {
-        case impl @ Template(constr, parents, self, body) =>
-          val selfToMap = if (self.symbol.exists) self.symbol :: Nil else Nil
-          val symsToMap = impl.symbol :: impl.constr.symbol :: selfToMap
-          val mappedSyms = ctx.mapSymbols(symsToMap, typeMap, ownerMap, substFrom, substTo)
-          val tmap = withSubstitution(symsToMap, mappedSyms)
-          val constr1 = tmap.transformSub(constr)
-          val parents1 = parents mapconserve transform
-          var self1 = tmap.transformSub(self)
-          if (self.symbol.exists) self1 = self1.withType(mappedSyms.last.termRef)
-          val body1 = tmap.transformStats(body)
-          cpy.Template(impl)(constr1, parents1, self1, body1).withType(mappedSyms.head.nonMemberTermRef)
-        case _ =>
-          tree1.withType(mapType(tree1.tpe)) match {
-            case ddef @ DefDef(mods, name, tparams, vparamss, tpt, rhs) =>
-              val (tmap1, tparams1) = transformDefs(ddef.tparams)
-              val (tmap2, vparamss1) = tmap1.transformVParamss(vparamss)
-              cpy.DefDef(ddef)(mods, name, tparams1, vparamss1, tmap2.transform(tpt), tmap2.transform(rhs))
-            case blk @ Block(stats, expr) =>
-              val (tmap1, stats1) = transformDefs(stats)
-              val expr1 = tmap1.transform(expr)
-              cpy.Block(blk)(stats1, expr1)
-            case cdef @ CaseDef(pat, guard, rhs) =>
-              val tmap = withMappedSyms(patVars(pat))
-              val pat1 = tmap.transform(pat)
-              val guard1 = tmap.transform(guard)
-              val rhs1 = tmap.transform(rhs)
-              cpy.CaseDef(tree)(pat1, guard1, rhs1)
-            case tree1 =>
-              super.transform(tree1)
-          }
-      }
-    }
-
-    override def transformStats(trees: List[tpd.Tree])(implicit ctx: Context) =
-      transformDefs(trees)._2
-
-    private def transformDefs[TT <: tpd.Tree](trees: List[TT])(implicit ctx: Context): (TreeTypeMap, List[TT]) = {
-      val tmap = withMappedSyms(ta.localSyms(trees))
-      (tmap, tmap.transformSub(trees))
-    }
-
-    private def transformVParamss(vparamss: List[List[ValDef]]): (TreeTypeMap, List[List[ValDef]]) = vparamss match {
-      case vparams :: rest =>
-        val (tmap1, vparams1) = transformDefs(vparams)
-        val (tmap2, vparamss2) = tmap1.transformVParamss(rest)
-        (tmap2, vparams1 :: vparamss2)
-      case nil =>
-        (this, vparamss)
-    }
-
-    def apply[ThisTree <: tpd.Tree](tree: ThisTree): ThisTree = transform(tree).asInstanceOf[ThisTree]
-
-    def apply(annot: Annotation): Annotation = {
-      val tree1 = apply(annot.tree)
-      if (tree1 eq annot.tree) annot else ConcreteAnnotation(tree1)
-    }
-
-    /** The current tree map composed with a substitution [from -> to] */
-    def withSubstitution(from: List[Symbol], to: List[Symbol]): TreeTypeMap =
-      if (from eq to) this
-      else {
-        // assert that substitution stays idempotent, assuming its parts are
-        assert(!from.exists(substTo contains _))
-        assert(!to.exists(substFrom contains _))
-        new TreeTypeMap(
-          typeMap,
-          ownerMap andThen { sym =>
-            val idx = from.indexOf(sym)
-            if (idx >= 0) to(idx) else sym
-          },
-          treeMap,
-          from ++ substFrom,
-          to ++ substTo)
-      }
-
-    /** Apply `typeMap` and `ownerMap` to given symbols `syms`
-     *  and return a treemap that contains the substitution
-     *  between original and mapped symbols.
-     */
-    def withMappedSyms(syms: List[Symbol]): TreeTypeMap = {
-      val mapped = ctx.mapSymbols(syms, typeMap, ownerMap, substFrom, substTo)
-      withSubstitution(syms, mapped)
-    }
-  }
-
-  /** The variables defined by a pattern, in reverse order of their appearance. */
-  def patVars(tree: Tree)(implicit ctx: Context): List[Symbol] = {
-    val acc = new TreeAccumulator[List[Symbol]] {
-      def apply(syms: List[Symbol], tree: Tree) = tree match {
-        case Bind(_, body) => apply(tree.symbol :: syms, body)
-        case _ => foldOver(syms, tree)
-      }
-    }
-    acc(Nil, tree)
-  }
-
   // convert a numeric with a toXXX method
   def primitiveConversion(tree: Tree, numericCls: Symbol)(implicit ctx: Context): Tree = {
     val mname      = ("to" + numericCls.name).toTermName
diff --git a/src/dotty/tools/dotc/core/Symbols.scala b/src/dotty/tools/dotc/core/Symbols.scala
index 4156d59d1..96819f627 100644
--- a/src/dotty/tools/dotc/core/Symbols.scala
+++ b/src/dotty/tools/dotc/core/Symbols.scala
@@ -19,7 +19,8 @@ import util.Positions._
 import DenotTransformers._
 import StdNames._
 import NameOps._
-import ast.tpd.{TreeTypeMap, Tree}
+import ast.tpd.Tree
+import ast.TreeTypeMap
 import Denotations.{ Denotation, SingleDenotation, MultiDenotation }
 import collection.mutable
 import io.AbstractFile
@@ -261,22 +262,14 @@ trait Symbols { this: Context =>
     newSymbol(owner, name, SyntheticArtifact,
       if (name.isTypeName) TypeAlias(ErrorType) else ErrorType)
 
-  /** Map given symbols, subjecting all types to given type map and owner map,
-   *  as well as to the substition [substFrom := substTo].
+  /** Map given symbols, subjecting their attributes to the mappings
+   *  defined in the given TreeTypeMap `ttmap`.
    *  Cross symbol references are brought over from originals to copies.
    *  Do not copy any symbols if all attributes of all symbols stay the same.
    */
-  def mapSymbols(
-      originals: List[Symbol],
-      typeMap: Type => Type = IdentityTypeMap,
-      ownerMap: Symbol => Symbol = identity,
-      substFrom: List[Symbol] = Nil,
-      substTo: List[Symbol] = Nil)
-  =
+  def mapSymbols(originals: List[Symbol], ttmap: TreeTypeMap) =
     if (originals forall (sym =>
-        (typeMap(sym.info) eq sym.info) &&
-        (sym.info.substSym(substFrom, substTo) eq sym.info) &&
-        (ownerMap(sym.owner) eq sym.owner)))
+        (ttmap.mapType(sym.info) eq sym.info) && (ttmap.ownerMap(sym.owner) eq sym.owner)))
       originals
     else {
       val copies: List[Symbol] = for (original <- originals) yield
@@ -286,19 +279,18 @@ trait Symbols { this: Context =>
           case _ =>
             newNakedSymbol[original.ThisName](original.coord)
         }
-      val treeMap = new TreeTypeMap(typeMap, ownerMap, substFrom = substFrom, substTo = substTo)
-        .withSubstitution(originals, copies)
+      val ttmap1 = ttmap.withSubstitution(originals, copies)
       (originals, copies).zipped foreach {(original, copy) =>
-        copy.denot = original.denot // preliminar denotation, so that we can access symbols in subsequent transform
+        copy.denot = original.denot // preliminary denotation, so that we can access symbols in subsequent transform
       }
       (originals, copies).zipped foreach {(original, copy) =>
         val odenot = original.denot
         copy.denot = odenot.copySymDenotation(
           symbol = copy,
-          owner = treeMap.ownerMap(odenot.owner),
-          info = treeMap.mapType(odenot.info),
-          privateWithin = ownerMap(odenot.privateWithin), // since this refers to outer symbols, need not include copies (from->to) in ownermap here.
-          annotations = odenot.annotations.mapConserve(treeMap.apply))
+          owner = ttmap1.ownerMap(odenot.owner),
+          info = ttmap1.mapType(odenot.info),
+          privateWithin = ttmap1.ownerMap(odenot.privateWithin), // since this refers to outer symbols, need not include copies (from->to) in ownermap here.
+          annotations = odenot.annotations.mapConserve(ttmap1.apply))
       }
       copies
     }
diff --git a/src/dotty/tools/dotc/core/Types.scala b/src/dotty/tools/dotc/core/Types.scala
index 82b2fc71a..17e7f4fb5 100644
--- a/src/dotty/tools/dotc/core/Types.scala
+++ b/src/dotty/tools/dotc/core/Types.scala
@@ -17,7 +17,9 @@ import Periods._
 import util.Positions.Position
 import util.Stats._
 import util.{DotClass, SimpleMap}
-import ast.tpd._, printing.Texts._
+import ast.tpd._
+import ast.TreeTypeMap
+import printing.Texts._
 import ast.untpd
 import transform.Erasure
 import printing.Printer
@@ -2474,8 +2476,9 @@ object Types {
       }
     }
 
-    def mapOver(syms: List[Symbol]): List[Symbol] =
-      ctx.mapSymbols(syms, typeMap = this)
+    private def treeTypeMap = new TreeTypeMap(typeMap = this)
+
+    def mapOver(syms: List[Symbol]): List[Symbol] = ctx.mapSymbols(syms, treeTypeMap)
 
     def mapOver(scope: Scope): Scope = {
       val elems = scope.toList
@@ -2487,8 +2490,7 @@ object Types {
     def mapOver(annot: Annotation): Annotation =
       annot.derivedAnnotation(mapOver(annot.tree))
 
-    def mapOver(tree: Tree): Tree =
-      new TreeTypeMap(typeMap = this).apply(tree)
+    def mapOver(tree: Tree): Tree = treeTypeMap(tree)
 
     /** Can be overridden. By default, only the prefix is mapped. */
     protected def mapClassInfo(tp: ClassInfo): ClassInfo =
diff --git a/src/dotty/tools/dotc/typer/TypeAssigner.scala b/src/dotty/tools/dotc/typer/TypeAssigner.scala
index 69334f525..50b0fe8c1 100644
--- a/src/dotty/tools/dotc/typer/TypeAssigner.scala
+++ b/src/dotty/tools/dotc/typer/TypeAssigner.scala
@@ -75,9 +75,6 @@ trait TypeAssigner {
     widenMap(tp)
   }
 
-  def localSyms(stats: List[tpd.Tree])(implicit ctx: Context): List[Symbol] =
-    for (stat <- stats if stat.isDef) yield stat.symbol
-
   def seqToRepeated(tree: Tree)(implicit ctx: Context): Tree =
     Typed(tree, TypeTree(tree.tpe.widen.translateParameterized(defn.SeqClass, defn.RepeatedParamClass)))
 
diff --git a/src/dotty/tools/dotc/typer/Typer.scala b/src/dotty/tools/dotc/typer/Typer.scala
index 5c2b44877..e1f860589 100644
--- a/src/dotty/tools/dotc/typer/Typer.scala
+++ b/src/dotty/tools/dotc/typer/Typer.scala
@@ -399,7 +399,7 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
 
   def escapingRefs(block: Block)(implicit ctx: Context): collection.Set[NamedType] = {
     var hoisted: Set[Symbol] = Set()
-    lazy val locals = ctx.typeAssigner.localSyms(block.stats).toSet
+    lazy val locals = localSyms(block.stats).toSet
     def isLocal(sym: Symbol): Boolean =
       (locals contains sym) && !isHoistableClass(sym)
     def isHoistableClass(sym: Symbol) =