Bullet-proofing companion objects

Companion class/module computations now also work for local classes and modules. For this to work,
either one of two conditions must be met: (1) some enclosing context refers to a scope that
contains the companions. (2) the context's compilation unit has a typed tree that contains the
companions. (1) is usually true when type-checking, (2) when transforming trees.

Local companions are searched as follows: If (2) holds, we locate the statement sequence containing
the companions by searching down from the root stored in the compilation unit. Otherwise, we search
outwards in the enclosing contexts for a scope containing the
companions.

1 files changed, 31 insertions, 10 deletions

diff --git a/src/dotty/tools/dotc/core/SymDenotations.scala b/src/dotty/tools/dotc/core/SymDenotations.scala
index bea8576c6..7fd5cdacc 100644
--- a/src/dotty/tools/dotc/core/SymDenotations.scala
+++ b/src/dotty/tools/dotc/core/SymDenotations.scala
@@ -9,6 +9,7 @@ import collection.mutable
 import collection.immutable.BitSet
 import scala.reflect.io.AbstractFile
 import Decorators.SymbolIteratorDecorator
+import ast.tpd
 import annotation.tailrec
 import util.SimpleMap
 import util.Stats
@@ -601,22 +602,42 @@ object SymDenotations {
      *  NoSymbol if this module does not exist.
      */
     final def companionModule(implicit ctx: Context): Symbol =
-      if (owner.exists && name != tpnme.ANON_CLASS) // name test to avoid forcing, thereby causing cyclic reference errors
-        owner.info.decl(effectiveName.toTermName)
-          .suchThat(sym => (sym is Module) && sym.isCoDefinedWith(symbol))
-          .symbol
-      else NoSymbol
+      if (name == tpnme.ANON_CLASS)
+        NoSymbol // avoid forcing anon classes, this might cause cyclic reference errors
+      else
+        companionNamed(effectiveName.moduleClassName).sourceModule
 
     /** The class with the same (type-) name as this module or module class,
      *  and which is also defined in the same scope and compilation unit.
      *  NoSymbol if this class does not exist.
      */
     final def companionClass(implicit ctx: Context): Symbol =
-      if (owner.exists)
-        owner.info.decl(effectiveName.toTypeName)
-          .suchThat(sym => sym.isClass && sym.isCoDefinedWith(symbol))
-          .symbol
-      else NoSymbol
+      companionNamed(effectiveName.toTypeName)
+
+    /** Find companion class symbol with given name, or NoSymbol if none exists.
+     *  Three alternative strategies:
+     *  1. If owner is a class, look in its members, otherwise
+     *  2. If current compilation unit has a typed tree,
+     *     determine the definining statement sequence and search its trees, otherwise
+     *  3. If context has an enclosing scope which defines this symbol,
+     *     lookup its companion in the same scope.
+     */
+    private def companionNamed(name: TypeName)(implicit ctx: Context): Symbol =
+      if (owner.isClass)
+        owner.info.decl(name).suchThat(_.isCoDefinedWith(symbol)).symbol
+      else if (!owner.exists || ctx.compilationUnit == null)
+        NoSymbol
+      else if (!ctx.compilationUnit.tpdTree.isEmpty)
+        tpd.definingStats(symbol).iterator
+          .map(tpd.definedSym)
+          .find(_.name == name)
+          .getOrElse(NoSymbol)
+      else if (ctx.scope == null)
+        NoSymbol
+      else if (ctx.scope.lookup(this.name) == symbol)
+        ctx.scope.lookup(name)
+      else
+        companionNamed(name)(ctx.outersIterator.dropWhile(_.scope eq ctx.scope).next)
 
     /** If this is a class, the module class of its companion object.
      *  If this is a module class, its companion class.