Added methods to substitute trees and change their owners.

Required also some better integration with TypeMaps.

4 files changed, 127 insertions, 9 deletions

diff --git a/src/dotty/tools/dotc/core/SymDenotations.scala b/src/dotty/tools/dotc/core/SymDenotations.scala
index 9299e7e63..e179a11a7 100644
--- a/src/dotty/tools/dotc/core/SymDenotations.scala
+++ b/src/dotty/tools/dotc/core/SymDenotations.scala
@@ -255,9 +255,12 @@ object SymDenotations {
       isModuleClass && linkedClass.isNonBottomSubClass(base)
 
     /** Is this symbol a class that does not extend `AnyVal`? */
-    def isNonValueClass(implicit ctx: Context): Boolean =
+    final def isNonValueClass(implicit ctx: Context): Boolean =
       isClass && !isSubClass(defn.AnyValClass)
 
+    /** Are the contents of this denotation invariant under the type map `f`? */
+    def isInvariantUnder(f: Type => Type)(implicit ctx: Context) = info eq f(info)
+
     /** Is this definition accessible whenever `that` symbol is accessible?
      *  Does not take into account status of protected members.
      */
@@ -497,7 +500,7 @@ object SymDenotations {
     override protected def newLikeThis(s: Symbol, i: Type): SingleDenotation = new UniqueRefDenotation(s, i, validFor)
 
     /** Copy this denotation, overriding selective fields */
-    def copy(
+    def copySym(
         sym: Symbol,
         owner: Symbol = this.owner,
         name: Name = this.name,
@@ -798,6 +801,11 @@ object SymDenotations {
       decls.denotsNamed(cname).first.symbol
     }
 
+    override def isInvariantUnder(f: Type => Type)(implicit ctx: Context) =
+      (parents.mapConserve(f) eq parents) &&
+      (f(selfType) eq selfType) &&
+      (decls forall (_.isInvariantUnder(f)))
+
     def copyClass(
         sym: ClassSymbol,
         owner: Symbol = this.owner,
diff --git a/src/dotty/tools/dotc/core/Symbols.scala b/src/dotty/tools/dotc/core/Symbols.scala
index 88a4b21f4..edd610bca 100644
--- a/src/dotty/tools/dotc/core/Symbols.scala
+++ b/src/dotty/tools/dotc/core/Symbols.scala
@@ -116,7 +116,7 @@ trait Symbols { this: Context =>
     newSymbol(NoSymbol, nme.IMPORT, EmptyFlags, ImportType(expr), coord = coord)
 
   def newConstructor(cls: ClassSymbol, flags: FlagSet, paramNames: List[TermName], paramTypes: List[Type], privateWithin: Symbol = NoSymbol, coord: Coord = NoCoord) =
-    newSymbol(cls, nme.CONSTRUCTOR, flags, MethodType(paramNames, paramTypes)(_ => cls.typeConstructor), privateWithin, coord)
+    newSymbol(cls, nme.CONSTRUCTOR, flags | Method, MethodType(paramNames, paramTypes)(_ => cls.typeConstructor), privateWithin, coord)
 
   def newDefaultConstructor(cls: ClassSymbol) =
     newConstructor(cls, EmptyFlags, Nil, Nil)
@@ -142,7 +142,44 @@ trait Symbols { this: Context =>
     tparams
   }
 
-  private val reverseApply = (x: TypeSymbol, f: TypeSymbol => TypeBounds) => f(x)
+  type OwnerMap = Symbol => Symbol
+
+  /** Map given symbols, subjecting all types to given type map and owner map. Cross symbol
+   *  references are brought over from originals to copies.
+   *  Do not copy any symbols if all their attributes stay the same.
+   */
+  def mapSymbols(originals: List[Symbol], typeMap: TypeMap = IdentityTypeMap, ownerMap: OwnerMap = identity) = {
+    if (originals forall (sym =>
+        sym.isInvariantUnder(typeMap) && ownerMap(sym.owner) == sym.owner))
+      originals
+    else {
+      lazy val copies: List[Symbol] = originals map { orig =>
+        if (orig.isClass)
+          newLazyClassSymbol(ownerMap(orig.owner), orig.asClass.name, orig.flags, copyClassCompleter, orig.asClass.associatedFile, orig.coord)
+        else
+          newLazySymbol(ownerMap(orig.owner), orig.name, orig.flags, copySymCompleter, orig.coord)
+      }
+      lazy val prev = (copies zip originals).toMap
+      lazy val copyTypeMap = typeMap andThen ((tp: Type) => tp.substSym(originals, copies))
+      lazy val mapper = new TypedTrees.TreeMapper(typeMap, ownerMap)
+      def mapAnnotations(denot: isLazy[_]): Unit = {
+        denot.annotations = denot.annotations.mapConserve(mapper.apply)
+        denot.privateWithin = ownerMap(denot.privateWithin)
+      }
+      def copySymCompleter(denot: LazySymDenotation): Unit = {
+        denot.info = copyTypeMap(prev(denot.symbol).info)
+        mapAnnotations(denot)
+      }
+      def copyClassCompleter(denot: LazyClassDenotation): Unit = {
+        denot.parents = prev(denot.symbol).asClass.parents mapConserve (
+          tp => copyTypeMap(tp).asInstanceOf[TypeRef])
+        denot.selfType = copyTypeMap(prev(denot.symbol).asClass.selfType)
+        denot.decls = copyTypeMap.mapOver(prev(denot.symbol).asClass.decls)
+        mapAnnotations(denot)
+      }
+      copies
+    }
+  }
 
 // ----- Locating predefined symbols ----------------------------------------
 
diff --git a/src/dotty/tools/dotc/core/TypedTrees.scala b/src/dotty/tools/dotc/core/TypedTrees.scala
index 2bdbd66aa..42eb166ff 100644
--- a/src/dotty/tools/dotc/core/TypedTrees.scala
+++ b/src/dotty/tools/dotc/core/TypedTrees.scala
@@ -599,7 +599,7 @@ object TypedTrees {
       check(shared.isType || shared.isTerm)
   }
 
-  implicit class TreeOps[T <: tpd.Tree](val tree: T) extends AnyVal {
+  implicit class TreeOps[ThisTree <: tpd.Tree](val tree: ThisTree) extends AnyVal {
 
     def isValue(implicit ctx: Context): Boolean =
       tree.isTerm && tree.tpe.widen.isValueType
@@ -610,15 +610,64 @@ object TypedTrees {
     def isValueType: Boolean =
       tree.isType && tree.tpe.isValueType
 
-    def isInstantiation = tree match {
+    def isInstantiation: Boolean = tree match {
       case Apply(Select(New(_), nme.CONSTRUCTOR), _) => true
       case _ => false
     }
 
-    def checked(implicit ctx: Context): T = {
+    def checked(implicit ctx: Context): ThisTree = {
       if (ctx.settings.YcheckTypedTrees.value) checkType(tree)
       tree
     }
+
+    def shallowFold[T](z: T)(op: (T, tpd.Tree) => T) =
+      new ShallowFolder(op).apply(z, tree)
+
+    def deepFold[T](z: T)(op: (T, tpd.Tree) => T) =
+      new DeepFolder(op).apply(z, tree)
+
+    def subst(from: List[Symbol], to: List[Symbol])(implicit ctx: Context): ThisTree =
+      new TreeMapper(typeMap = new ctx.SubstSymMap(from, to)).apply(tree)
+
+    def changeOwner(from: Symbol, to: Symbol)(implicit ctx: Context): ThisTree =
+      new TreeMapper(ownerMap = (sym => if (sym == from) to else sym)).apply(tree)
+  }
+
+  class TreeMapper(typeMap: TypeMap = IdentityTypeMap, ownerMap: Symbol => Symbol = identity)(implicit ctx: Context) extends TreeTransformer[Type, Unit] {
+    override def transform(tree: tpd.Tree, c: Unit): tpd.Tree = {
+      val tree1 = tree.withType(typeMap(tree.tpe))
+      val tree2 = tree1 match {
+        case bind: tpd.Bind =>
+          val sym = bind.symbol
+          val newOwner = ownerMap(sym.owner)
+          val newInfo = typeMap(sym.info)
+          if ((newOwner ne sym.owner) || (newInfo ne sym.info))
+            bind.withType(tpd.refType(sym.copy(owner = newOwner, info = newInfo)))
+          else
+            tree1
+        case _ =>
+          tree1
+      }
+      super.transform(tree2, c)
+    }
+    override def transform(trees: List[tpd.Tree], c: Unit) = {
+      val locals = localSyms(trees)
+      val mapped = ctx.mapSymbols(locals, typeMap, ownerMap)
+      if (locals eq mapped)
+        super.transform(trees, c)
+      else
+        new TreeMapper(
+          typeMap andThen ((tp: Type) => tp.substSym(locals, mapped)),
+          ownerMap andThen (locals zip mapped).toMap)
+          .transform(trees, c)
+    }
+
+    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)
+    }
   }
 
   // ensure that constructors are fully applied?
diff --git a/src/dotty/tools/dotc/core/Types.scala b/src/dotty/tools/dotc/core/Types.scala
index 4b3271ed2..09a2585d0 100644
--- a/src/dotty/tools/dotc/core/Types.scala
+++ b/src/dotty/tools/dotc/core/Types.scala
@@ -209,7 +209,7 @@ object Types {
     final def forall(p: Type => Boolean): Boolean = !exists(!p(_))
 
     /** Substitute all types that refer in their symbol attribute to
-     *  one of the symbols in `from` by the corresponding types in `to`
+     *  one of the symbols in `from` by the corresponding types in `to`.
      */
     final def subst(from: List[Symbol], to: List[Type])(implicit ctx: Context): Type =
       if (from.isEmpty) this
@@ -237,6 +237,11 @@ object Types {
     final def substThis(rt: RefinedType, tp: Type)(implicit ctx: Context): Type =
       ctx.substThis(this, rt, tp, null)
 
+    /** Substitute all occurrences symbols in `from` by references to corresponding symbols in `to`
+     */
+    final def substSym(from: List[Symbol], to: List[Symbol])(implicit ctx: Context): Type =
+      ctx.substSym(this, from, to, null)
+
     /** For a ClassInfo type, its parents,
      *  Inherited by all type proxies. Empty for all other types.
      *  Overwritten in ClassInfo, where parents is cached.
@@ -963,6 +968,8 @@ object Types {
 
   abstract class GenericMethodType {
     def apply(paramNames: List[TermName], paramTypes: List[Type])(resultTypeExp: MethodType => Type)(implicit ctx: Context): MethodType
+    def apply(paramNames: List[TermName], paramTypes: List[Type], resultType: Type)(implicit ctx: Context): MethodType =
+      apply(paramNames, paramTypes)(_ => resultType)
     def fromSymbols(params: List[Symbol], resultType: Type)(implicit ctx: Context) = {
       def transResult(mt: MethodType) =
         resultType.subst(params, (0 until params.length).toList map (MethodParam(mt, _)))
@@ -1172,7 +1179,7 @@ object Types {
 
   // ----- TypeMaps --------------------------------------------------------------------
 
-  abstract class TypeMap(implicit ctx: Context) extends (Type => Type) {
+  abstract class TypeMap(implicit ctx: Context) extends (Type => Type) { thisMap =>
     def apply(tp: Type): Type
 
     def applyToBounds(tp: TypeBounds): TypeBounds =
@@ -1217,8 +1224,25 @@ object Types {
         tp
     }
 
+    def mapOver(syms: List[Symbol]): List[Symbol] =
+      ctx.mapSymbols(syms, this)
+
+    def mapOver(scope: Scope): Scope = {
+      val elems = scope.toList
+      val elems1 = mapOver(elems)
+      if (elems1 eq elems) scope
+      else newScopeWith(elems1: _*)
+    }
+
     def mapOverAnnotations(annots: List[Annotation]): List[Annotation] = ???
 
+    def andThen(f: Type => Type): TypeMap = new TypeMap {
+      def apply(tp: Type) = f(thisMap.apply(tp))
+    }
+  }
+
+  object IdentityTypeMap extends TypeMap()(NoContext) {
+    def apply(tp: Type) = tp
   }
 
   class InstMethodMap(mt: MethodType, argtypes: List[Type])(implicit ctx: Context) extends TypeMap {