support merging companion objects in expanded trees

Previous implementation is problematic when there are multiple
transforms before typer:

1. There might be objects needing merging that only exist in the
expanded trees, which cannot be handled by the previous algorithm.

2. There may be companion object and class defined only in the
expanded trees, the previous algorithm doesn't create links for them.

This PR simplifies the companion object merging logic and
fixes the problems above.

In fact, this PR supports multiple pre-typer transform during expansion.
The protocol is that the expansion of a MemberDef is either a flattened
thicket or a non-thicket tree.

1 files changed, 92 insertions, 31 deletions

diff --git a/compiler/src/dotty/tools/dotc/typer/Namer.scala b/compiler/src/dotty/tools/dotc/typer/Namer.scala
index 96660f15c..7d8c5f95f 100644
--- a/compiler/src/dotty/tools/dotc/typer/Namer.scala
+++ b/compiler/src/dotty/tools/dotc/typer/Namer.scala
@@ -484,46 +484,86 @@ class Namer { typer: Typer =>
   /** Create top-level symbols for statements and enter them into symbol table */
   def index(stats: List[Tree])(implicit ctx: Context): Context = {
 
-    val classDef  = mutable.Map[TypeName, TypeDef]()
-    val moduleDef = mutable.Map[TypeName, TypeDef]()
+    // module name -> (stat, moduleCls / moduleVal)
+    val moduleClsDef = mutable.Map[TypeName, (Tree, TypeDef)]()
+    val moduleValDef = mutable.Map[TermName, (Tree, ValDef)]()
+
+    /** Remove the subtree `tree` from the expanded tree of `mdef` */
+    def removeInExpanded(mdef: Tree, tree: Tree): Unit = {
+      val Thicket(trees) = expanded(mdef)
+      mdef.putAttachment(ExpandedTree, Thicket(trees.filter(_ != tree)))
+    }
+
+    /** Merge the module class `modCls` in the expanded tree of `mdef` with the given stats */
+    def mergeModuleClass(mdef: Tree, modCls: TypeDef, stats: List[Tree]): TypeDef = {
+      var res: TypeDef = null
+      val Thicket(trees) = expanded(mdef)
+      val merged = trees.map { tree =>
+        if (tree == modCls) {
+          val impl = modCls.rhs.asInstanceOf[Template]
+          res = cpy.TypeDef(modCls)(rhs = cpy.Template(impl)(body = stats ++ impl.body))
+          res
+        }
+        else tree
+      }
+
+      mdef.putAttachment(ExpandedTree, Thicket(merged))
+
+      res
+    }
 
     /** Merge the definitions of a synthetic companion generated by a case class
      *  and the real companion, if both exist.
      */
     def mergeCompanionDefs() = {
-      for (cdef @ TypeDef(name, _) <- stats)
-        if (cdef.isClassDef) {
-          classDef(name) = cdef
-          cdef.attachmentOrElse(ExpandedTree, cdef) match {
-            case Thicket(cls :: mval :: (mcls @ TypeDef(mname, _: Template)) :: crest)
-            if name.moduleClassName == mname =>
-              moduleDef(name) = mcls
-            case _ =>
-          }
-        }
-      for (mdef @ ModuleDef(name, _) <- stats if !mdef.mods.is(Flags.Package)) {
-        val typName = name.toTypeName
-        // Expansion of object is a flattened thicket with the first two elements being:
-        //    module val :: module class :: rest
-        val Thicket(vdef :: (mcls @ TypeDef(_, impl: Template)) :: rest) = expanded(mdef)
-        moduleDef(typName) = mcls
-        classDef get name.toTypeName match {
-          case Some(cdef) =>
-            cdef.attachmentOrElse(ExpandedTree, cdef) match {
-              case Thicket(cls :: mval :: TypeDef(mname, compimpl: Template) :: crest)
-              if name.moduleClassName == mname =>
-                val mcls1 = cpy.TypeDef(mcls)(
-                    rhs = cpy.Template(impl)(body = compimpl.body ++ impl.body))
-                mdef.putAttachment(ExpandedTree, Thicket(vdef :: mcls1 :: rest))
-                moduleDef(typName) = mcls1
-                cdef.putAttachment(ExpandedTree, Thicket(cls :: crest))
+      def valid(mdef: MemberDef): Boolean = !mdef.mods.is(Package) && mdef.mods.is(Module)
+
+      for (stat <- stats)
+        expanded(stat) match {
+          case Thicket(trees) => // companion object always expands to thickets
+            trees.map {
+              case mcls @ TypeDef(name, impl: Template) if valid(mcls) =>
+                if (moduleClsDef.contains(name)) {
+                  val (stat1, mcls1 @ TypeDef(_, impl1: Template)) = moduleClsDef(name)
+                  if (mcls.mods.is(Synthetic) && !mcls1.mods.is(Synthetic)) { // merge synthetic into user-defined module
+                    removeInExpanded(stat, mcls)
+                    val mcls2 = mergeModuleClass(stat1, mcls1, impl.body)
+                    moduleClsDef(name) = (stat1, mcls2)
+                  }
+                  else if (!mcls.mods.is(Synthetic) && mcls1.mods.is(Synthetic)) {
+                    removeInExpanded(stat1, mcls1)
+                    val mcls2 = mergeModuleClass(stat, mcls, impl1.body)
+                    moduleClsDef(name) = (stat, mcls2)
+                  }
+                  else {
+                    // redefinition of objects or case classes, handled elsewhere
+                  }
+                }
+                else moduleClsDef(name) = (stat, mcls)
+
+              case vdef @ ValDef(name, _, _) if valid(vdef) =>
+                if (moduleValDef.contains(name)) {
+                  val (stat1, vdef1) = moduleValDef(name)
+                  if (vdef.mods.is(Synthetic) && !vdef1.mods.is(Synthetic))  // merge synthetic into user-defined module
+                    removeInExpanded(stat, vdef)
+                  else if (!vdef.mods.is(Synthetic) && vdef1.mods.is(Synthetic)) {
+                    removeInExpanded(stat1, vdef1)
+                    moduleValDef(name) = (stat, vdef)
+                  }
+                  else {
+                    // redefinition of objects or case classes, handled elsewhere
+                  }
+                }
+                else moduleValDef(name) = (stat, vdef)
+
               case _ =>
             }
-          case none =>
+          case _ =>
+
         }
-      }
     }
 
+    /** Create links between companion object and companion class */
     def createLinks(classTree: TypeDef, moduleTree: TypeDef)(implicit ctx: Context) = {
       val claz = ctx.denotNamed(classTree.name.encode).symbol
       val modl = ctx.denotNamed(moduleTree.name.encode).symbol
@@ -532,8 +572,29 @@ class Namer { typer: Typer =>
     }
 
     def createCompanionLinks(implicit ctx: Context): Unit = {
+      val classDef  = mutable.Map[TypeName, TypeDef]()
+      val moduleDef = mutable.Map[TypeName, TypeDef]()
+
+      def updateCache(cdef: TypeDef): Unit = {
+        if (!cdef.isClassDef || cdef.mods.is(Package)) return
+
+        if (cdef.mods.is(ModuleClass)) moduleDef(cdef.name) = cdef
+        else classDef(cdef.name) = cdef
+      }
+
+      for (stat <- stats)
+        expanded(stat) match {
+          case cdef : TypeDef => updateCache(cdef)
+          case Thicket(trees) =>
+            trees.map {
+              case cdef: TypeDef => updateCache(cdef)
+              case _ =>
+            }
+          case _ =>
+        }
+
       for (cdef @ TypeDef(name, _) <- classDef.values) {
-        moduleDef.getOrElse(name, EmptyTree) match {
+        moduleDef.getOrElse(name.moduleClassName, EmptyTree) match {
           case t: TypeDef =>
             createLinks(cdef, t)
           case EmptyTree =>