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.

5 files changed, 98 insertions, 18 deletions

diff --git a/src/dotty/tools/dotc/ast/TreeInfo.scala b/src/dotty/tools/dotc/ast/TreeInfo.scala
index a36e2d600..5e04474cb 100644
--- a/src/dotty/tools/dotc/ast/TreeInfo.scala
+++ b/src/dotty/tools/dotc/ast/TreeInfo.scala
@@ -276,7 +276,11 @@ trait TreeInfo[T >: Untyped <: Type] { self: Trees.Instance[T] =>
   }
 }
 
-trait TypedTreeInfo extends TreeInfo[Type] {self: Trees.Instance[Type] =>
+trait UntypedTreeInfo extends TreeInfo[Untyped] { self: Trees.Instance[Untyped] =>
+  // todo: fill with methods from TreeInfo that only apply to untpd.Tree's
+}
+
+trait TypedTreeInfo extends TreeInfo[Type] { self: Trees.Instance[Type] =>
 
   /** Is tree a definition that has no side effects when
    *  evaluated as part of a block after the first time?
@@ -407,11 +411,67 @@ trait TypedTreeInfo extends TreeInfo[Type] {self: Trees.Instance[Type] =>
       false
   }
 
+  /** If `tree` is a DefTree, the symbol defined by it, otherwise NoSymbol */
+  def definedSym(tree: Tree)(implicit ctx: Context): Symbol =
+    if (tree.isDef) tree.symbol else NoSymbol
+
+  /** Going from child to parent, the path of tree nodes that starts
+   *  with a definition of symbol `sym` and ends with `root`, or Nil
+   *  if no such path exists.
+   *  Pre: `sym` must have a position.
+   */
+  def defPath(sym: Symbol, root: Tree)(implicit ctx: Context): List[Tree] = ctx.debugTraceIndented(s"defpath($sym with position ${sym.pos}, ${root.show})") {
+    def show(from: Any): String = from match {
+      case tree: Trees.Tree[_] => s"${tree.show} with attachments ${tree.allAttachments}"
+      case x: printing.Showable => x.show
+      case x => x.toString
+    }
+
+    def search(from: Any): List[Tree] = ctx.debugTraceIndented(s"search(${show(from)})") {
+      from match {
+        case tree: Tree @ unchecked =>
+          if (definedSym(tree) == sym) tree :: Nil
+          else if (tree.envelope.contains(sym.pos)) {
+            val p = search(tree.productIterator)
+            if (p.isEmpty) p else tree :: p
+          } else Nil
+        case xs: Iterable[_] =>
+          search(xs.iterator)
+        case xs: Iterator[_] =>
+          xs.map(search).find(_.nonEmpty).getOrElse(Nil)
+        case _ =>
+          Nil
+      }
+    }
+    require(sym.pos.exists)
+    search(root)
+  }
+
+  /** The statement sequence that contains a definition of `sym`, or Nil
+   *  if none was found.
+   *  For a tree to be found, The symbol must have a position and its definition
+   *  tree must be reachable from come tree stored in an enclosing context.
+   */
+  def definingStats(sym: Symbol)(implicit ctx: Context): List[Tree] =
+    if (!sym.pos.exists || (ctx eq NoContext) || ctx.compilationUnit == null) Nil
+    else defPath(sym, ctx.compilationUnit.tpdTree) match {
+      case defn :: encl :: _ =>
+        def verify(stats: List[Tree]) =
+          if (stats exists (definedSym(_) == sym)) stats else Nil
+        encl match {
+          case Block(stats, _) => verify(stats)
+          case Template(_, _, _, stats) => verify(stats)
+          case PackageDef(_, stats) => verify(stats)
+          case _ => Nil
+        }
+      case nil =>
+        Nil
+    }
+}
+
   /** a Match(Typed(_, tpt), _) must be translated into a switch if isSwitchAnnotation(tpt.tpe)
   def isSwitchAnnotation(tpe: Type) = tpe hasAnnotation defn.SwitchClass
   */
-}
-
 
   /** Does list of trees start with a definition of
    *  a class of module with given name (ignoring imports)
diff --git a/src/dotty/tools/dotc/ast/untpd.scala b/src/dotty/tools/dotc/ast/untpd.scala
index b6ae6661f..aea34e08a 100644
--- a/src/dotty/tools/dotc/ast/untpd.scala
+++ b/src/dotty/tools/dotc/ast/untpd.scala
@@ -9,7 +9,7 @@ import Decorators._
 import language.higherKinds
 import collection.mutable.ListBuffer
 
-object untpd extends Trees.Instance[Untyped] with TreeInfo[Untyped] {
+object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
 
 // ----- Tree cases that exist in untyped form only ------------------
 
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.
diff --git a/src/dotty/tools/dotc/typer/Applications.scala b/src/dotty/tools/dotc/typer/Applications.scala
index cf0d7d0a1..feaf678a2 100644
--- a/src/dotty/tools/dotc/typer/Applications.scala
+++ b/src/dotty/tools/dotc/typer/Applications.scala
@@ -22,7 +22,6 @@ import StdNames._
 import ProtoTypes._
 import EtaExpansion._
 import collection.mutable
-import reflect.ClassTag
 import config.Printers._
 import TypeApplications._
 import language.implicitConversions
diff --git a/src/dotty/tools/dotc/util/Attachment.scala b/src/dotty/tools/dotc/util/Attachment.scala
index d9e88a135..ec3019bab 100644
--- a/src/dotty/tools/dotc/util/Attachment.scala
+++ b/src/dotty/tools/dotc/util/Attachment.scala
@@ -74,10 +74,10 @@ object Attachment {
       else nx.removeAttachment(key)
     }
 
-    /** The list of all values attached to this container. */
-    final def allAttachments: List[Any] = {
+    /** The list of all keys and values attached to this container. */
+    final def allAttachments: List[(Key[_], Any)] = {
       val nx = next
-      if (nx == null) Nil else nx.value :: nx.allAttachments
+      if (nx == null) Nil else (nx.key, nx.value) :: nx.allAttachments
     }
   }