support merging companion objects in expanded trees

When there are multiple transform before typer, there might
be companion objects that only exist in the expanded trees.
This PR simplifies original companion object merging logic and
provides support for multi-level expansion.

Author: liufengyun

compiler/src/dotty/tools/dotc/typer/Namer.scala

@@ -484,46 +484,60 @@ 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 -> (tree, module class) (in the expanded tree of which the module appear)
+    val moduleDef = mutable.Map[TypeName, (Tree, TypeDef)]()
+    val moduleValDef = mutable.Set[TermName]()
+
+    /** 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]): Unit = {
+      val Thicket(trees) = expanded(mdef)
+      val merged = trees.map { tree =>
+        if (tree == modCls) {
+          val impl = modCls.rhs.asInstanceOf[Template]
+          cpy.TypeDef(modCls)(rhs = cpy.Template(impl)(body = stats ++ impl.body))
+        }
+        else tree
+      }
+
+      mdef.putAttachment(ExpandedTree, Thicket(merged))
+    }
 
     /** 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(Flags.Package) && mdef.mods.is(Flags.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 (moduleDef.contains(name)) {
+                  val (stat1, modCls) = moduleDef(name)
+                  removeInExpanded(stat, mcls)
+                  mergeModuleClass(stat1, modCls, impl.body)
+                }
+                else moduleDef(name) = (stat, mcls)
+
+              case vdef @ ValDef(name, _, _) if valid(vdef) =>
+                if (moduleValDef.contains(name)) removeInExpanded(stat, vdef)
+                else moduleValDef += name
+
               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 +546,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(Flags.Package)) return
+
+        if (cdef.mods.is(Flags.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 =>