support merging companion objects in expanded trees

compiler/src/dotty/tools/dotc/ast/untpd.scala

@@ -149,6 +149,7 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
 
     def is(fs: FlagSet): Boolean = flags is fs
     def is(fc: FlagConjunction): Boolean = flags is fc
+    def is(fc: FlagSet, butNot: FlagSet): Boolean = flags.is(fc, butNot = butNot)
 
     def | (fs: FlagSet): Modifiers = withFlags(flags | fs)
     def & (fs: FlagSet): Modifiers = withFlags(flags & fs)

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

@@ -484,46 +484,91 @@ 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 `fromCls` of `fromStat` into `toCls` of `toStat`
+     *  if the former is synthetic and the latter not.
+     *
+     *  Note:
+     *    1. `fromStat` and `toStat` could be the same stat
+     *    2. `fromCls` and `toCls` are necessarily different
+     */
+    def mergeIfSynthetic(fromStat: Tree, fromCls: TypeDef, toStat: Tree, toCls: TypeDef): Unit =
+      if (fromCls.mods.is(Synthetic) && !toCls.mods.is(Synthetic)) {
+        removeInExpanded(fromStat, fromCls)
+        val mcls = mergeModuleClass(toStat, toCls, fromCls.rhs.asInstanceOf[Template].body)
+        moduleClsDef(fromCls.name) = (toStat, mcls)
+      }
 
     /** 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(Module, butNot = Package)
+
+      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) =>
+                moduleClsDef.get(name) match {
+                  case Some((stat1, mcls1@TypeDef(_, impl1: Template))) =>
+                    mergeIfSynthetic(stat, mcls, stat1, mcls1)
+                    mergeIfSynthetic(stat1, mcls1, stat, mcls)
+                  case None =>
+                    moduleClsDef(name) = (stat, mcls)
+                }
+
+              case vdef @ ValDef(name, _, _) if valid(vdef) =>
+                moduleValDef.get(name) match {
+                  case Some((stat1, vdef1)) =>
+                    if (vdef.mods.is(Synthetic) && !vdef1.mods.is(Synthetic))
+                      removeInExpanded(stat, vdef)
+                    else if (!vdef.mods.is(Synthetic) && vdef1.mods.is(Synthetic)) {
+                      removeInExpanded(stat1, vdef1)
+                      moduleValDef(name) = (stat, vdef)
+                    }
+                    else {
+                      // double definition of objects or case classes, handled elsewhere
+                    }
+                  case None =>
+                    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 +577,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 =>