Merge pull request #1151 from felixmulder/topic/wip-docstrings

Add support for raw docstrings in ASTs

11 files changed, 605 insertions, 46 deletions

diff --git a/src/dotty/tools/dotc/ast/Trees.scala b/src/dotty/tools/dotc/ast/Trees.scala
index d0197b443..be3e2eb5f 100644
--- a/src/dotty/tools/dotc/ast/Trees.scala
+++ b/src/dotty/tools/dotc/ast/Trees.scala
@@ -29,6 +29,9 @@ object Trees {
   /** The total number of created tree nodes, maintained if Stats.enabled */
   @sharable var ntrees = 0
 
+  /** Attachment key for trees with documentation strings attached */
+  val DocComment = new Attachment.Key[String]
+
   /** Modifiers and annotations for definitions
    *  @param flags          The set flags
    *  @param privateWithin  If a private or protected has is followed by a
@@ -321,6 +324,8 @@ object Trees {
     private[ast] def rawMods: Modifiers[T] =
       if (myMods == null) genericEmptyModifiers else myMods
 
+    def rawComment: Option[String] = getAttachment(DocComment)
+
     def withMods(mods: Modifiers[Untyped]): ThisTree[Untyped] = {
       val tree = if (myMods == null || (myMods == mods)) this else clone.asInstanceOf[MemberDef[Untyped]]
       tree.setMods(mods)
@@ -329,6 +334,11 @@ object Trees {
 
     def withFlags(flags: FlagSet): ThisTree[Untyped] = withMods(Modifiers(flags))
 
+    def setComment(comment: Option[String]): ThisTree[Untyped] = {
+      comment.map(putAttachment(DocComment, _))
+      asInstanceOf[ThisTree[Untyped]]
+    }
+
     protected def setMods(mods: Modifiers[T @uncheckedVariance]) = myMods = mods
 
     override def envelope: Position = rawMods.pos.union(pos).union(initialPos)
diff --git a/src/dotty/tools/dotc/config/Printers.scala b/src/dotty/tools/dotc/config/Printers.scala
index 21147fe6f..fa36ad12c 100644
--- a/src/dotty/tools/dotc/config/Printers.scala
+++ b/src/dotty/tools/dotc/config/Printers.scala
@@ -13,6 +13,7 @@ object Printers {
   }
 
   val default: Printer = new Printer
+  val dottydoc: Printer = noPrinter
   val core: Printer = noPrinter
   val typr: Printer = noPrinter
   val constr: Printer = noPrinter
diff --git a/src/dotty/tools/dotc/config/ScalaSettings.scala b/src/dotty/tools/dotc/config/ScalaSettings.scala
index 07a23fdb6..f9c10442a 100644
--- a/src/dotty/tools/dotc/config/ScalaSettings.scala
+++ b/src/dotty/tools/dotc/config/ScalaSettings.scala
@@ -159,6 +159,7 @@ class ScalaSettings extends Settings.SettingGroup {
   val YprintSyms = BooleanSetting("-Yprint-syms", "when printing trees print info in symbols instead of corresponding info in trees.")
   val YtestPickler = BooleanSetting("-Ytest-pickler", "self-test for pickling functionality; should be used with -Ystop-after:pickler")
   val YcheckReentrant = BooleanSetting("-Ycheck-reentrant", "check that compiled program does not contain vars that can be accessed from a global root.")
+  val YkeepComments = BooleanSetting("-Ykeep-comments", "Keep comments when scanning source files.")
   def stop = YstopAfter
 
   /** Area-specific debug output.
diff --git a/src/dotty/tools/dotc/core/Contexts.scala b/src/dotty/tools/dotc/core/Contexts.scala
index a0bb03e50..ad3a0057d 100644
--- a/src/dotty/tools/dotc/core/Contexts.scala
+++ b/src/dotty/tools/dotc/core/Contexts.scala
@@ -550,6 +550,14 @@ object Contexts {
     def squashed(p: Phase): Phase = {
       allPhases.find(_.period.containsPhaseId(p.id)).getOrElse(NoPhase)
     }
+
+    val _docstrings: mutable.Map[Symbol, String] =
+      mutable.Map.empty
+
+    def docstring(sym: Symbol): Option[String] = _docstrings.get(sym)
+
+    def addDocstring(sym: Symbol, doc: Option[String]): Unit =
+      doc.map(d => _docstrings += (sym -> d))
   }
 
   /** The essential mutable state of a context base, collected into a common class */
diff --git a/src/dotty/tools/dotc/parsing/Parsers.scala b/src/dotty/tools/dotc/parsing/Parsers.scala
index 44a70886e..cdfc366a7 100644
--- a/src/dotty/tools/dotc/parsing/Parsers.scala
+++ b/src/dotty/tools/dotc/parsing/Parsers.scala
@@ -1752,13 +1752,13 @@ object Parsers {
      */
     def defOrDcl(start: Int, mods: Modifiers): Tree = in.token match {
       case VAL =>
-        patDefOrDcl(posMods(start, mods))
+        patDefOrDcl(posMods(start, mods), in.getDocString(start))
       case VAR =>
-        patDefOrDcl(posMods(start, addFlag(mods, Mutable)))
+        patDefOrDcl(posMods(start, addFlag(mods, Mutable)), in.getDocString(start))
       case DEF =>
-        defDefOrDcl(posMods(start, mods))
+        defDefOrDcl(posMods(start, mods), in.getDocString(start))
       case TYPE =>
-        typeDefOrDcl(posMods(start, mods))
+        typeDefOrDcl(posMods(start, mods), in.getDocString(start))
       case _ =>
         tmplDef(start, mods)
     }
@@ -1768,7 +1768,7 @@ object Parsers {
      *  ValDcl ::= Id {`,' Id} `:' Type
      *  VarDcl ::= Id {`,' Id} `:' Type
      */
-    def patDefOrDcl(mods: Modifiers): Tree = {
+    def patDefOrDcl(mods: Modifiers, docstring: Option[String] = None): Tree = {
       val lhs = commaSeparated(pattern2)
       val tpt = typedOpt()
       val rhs =
@@ -1782,8 +1782,10 @@ object Parsers {
           }
         } else EmptyTree
       lhs match {
-        case (id @ Ident(name: TermName)) :: Nil => cpy.ValDef(id)(name, tpt, rhs).withMods(mods)
-        case _ => PatDef(mods, lhs, tpt, rhs)
+        case (id @ Ident(name: TermName)) :: Nil => {
+          cpy.ValDef(id)(name, tpt, rhs).withMods(mods).setComment(docstring)
+        } case _ =>
+          PatDef(mods, lhs, tpt, rhs)
       }
     }
 
@@ -1792,7 +1794,7 @@ object Parsers {
      *  DefDcl ::= DefSig `:' Type
      *  DefSig ::= id [DefTypeParamClause] ParamClauses
      */
-    def defDefOrDcl(mods: Modifiers): Tree = atPos(tokenRange) {
+    def defDefOrDcl(mods: Modifiers, docstring: Option[String] = None): Tree = atPos(tokenRange) {
       def scala2ProcedureSyntax(resultTypeStr: String) = {
         val toInsert =
           if (in.token == LBRACE) s"$resultTypeStr ="
@@ -1833,7 +1835,7 @@ object Parsers {
             accept(EQUALS)
             expr()
           }
-        DefDef(name, tparams, vparamss, tpt, rhs).withMods(mods1)
+        DefDef(name, tparams, vparamss, tpt, rhs).withMods(mods1).setComment(docstring)
       }
     }
 
@@ -1867,7 +1869,7 @@ object Parsers {
     /** TypeDef ::= type Id [TypeParamClause] `=' Type
      *  TypeDcl ::= type Id [TypeParamClause] TypeBounds
      */
-    def typeDefOrDcl(mods: Modifiers): Tree = {
+    def typeDefOrDcl(mods: Modifiers, docstring: Option[String] = None): Tree = {
       newLinesOpt()
       atPos(tokenRange) {
         val name = ident().toTypeName
@@ -1875,7 +1877,7 @@ object Parsers {
         in.token match {
           case EQUALS =>
             in.nextToken()
-            TypeDef(name, tparams, typ()).withMods(mods)
+            TypeDef(name, tparams, typ()).withMods(mods).setComment(docstring)
           case SUPERTYPE | SUBTYPE | SEMI | NEWLINE | NEWLINES | COMMA | RBRACE | EOF =>
             TypeDef(name, tparams, typeBounds()).withMods(mods)
           case _ =>
@@ -1888,35 +1890,40 @@ object Parsers {
     /** TmplDef ::= ([`case'] `class' | `trait') ClassDef
      *            |  [`case'] `object' ObjectDef
      */
-    def tmplDef(start: Int, mods: Modifiers): Tree = in.token match {
-      case TRAIT =>
-        classDef(posMods(start, addFlag(mods, Trait)))
-      case CLASS =>
-        classDef(posMods(start, mods))
-      case CASECLASS =>
-        classDef(posMods(start, mods | Case))
-      case OBJECT =>
-        objectDef(posMods(start, mods | Module))
-      case CASEOBJECT =>
-        objectDef(posMods(start, mods | Case | Module))
-      case _ =>
-        syntaxErrorOrIncomplete("expected start of definition")
-        EmptyTree
+    def tmplDef(start: Int, mods: Modifiers): Tree = {
+      val docstring = in.getDocString(start)
+      in.token match {
+        case TRAIT =>
+          classDef(posMods(start, addFlag(mods, Trait)), docstring)
+        case CLASS =>
+          classDef(posMods(start, mods), docstring)
+        case CASECLASS =>
+          classDef(posMods(start, mods | Case), docstring)
+        case OBJECT =>
+          objectDef(posMods(start, mods | Module), docstring)
+        case CASEOBJECT =>
+          objectDef(posMods(start, mods | Case | Module), docstring)
+        case _ =>
+          syntaxErrorOrIncomplete("expected start of definition")
+          EmptyTree
+      }
     }
 
     /** ClassDef ::= Id [ClsTypeParamClause]
      *               [ConstrMods] ClsParamClauses TemplateOpt
      */
-    def classDef(mods: Modifiers): TypeDef = atPos(tokenRange) {
+    def classDef(mods: Modifiers, docstring: Option[String]): TypeDef = atPos(tokenRange) {
       val name = ident().toTypeName
       val constr = atPos(in.offset) {
         val tparams = typeParamClauseOpt(ParamOwner.Class)
         val cmods = constrModsOpt()
         val vparamss = paramClauses(name, mods is Case)
+
         makeConstructor(tparams, vparamss).withMods(cmods)
       }
       val templ = templateOpt(constr)
-      TypeDef(name, templ).withMods(mods)
+
+      TypeDef(name, templ).withMods(mods).setComment(docstring)
     }
 
     /** ConstrMods        ::=  AccessModifier
@@ -1932,10 +1939,11 @@ object Parsers {
 
     /** ObjectDef       ::= Id TemplateOpt
      */
-    def objectDef(mods: Modifiers): ModuleDef = {
+    def objectDef(mods: Modifiers, docstring: Option[String] = None): ModuleDef = {
       val name = ident()
       val template = templateOpt(emptyConstructor())
-      ModuleDef(name, template).withMods(mods)
+
+      ModuleDef(name, template).withMods(mods).setComment(docstring)
     }
 
 /* -------- TEMPLATES ------------------------------------------- */
@@ -2160,7 +2168,8 @@ object Parsers {
         if (in.token == PACKAGE) {
           in.nextToken()
           if (in.token == OBJECT) {
-            ts += objectDef(atPos(start, in.skipToken()) { Modifiers(Package) })
+            val docstring = in.getDocString(start)
+            ts += objectDef(atPos(start, in.skipToken()) { Modifiers(Package) }, docstring)
             if (in.token != EOF) {
               acceptStatSep()
               ts ++= topStatSeq()
diff --git a/src/dotty/tools/dotc/parsing/Scanners.scala b/src/dotty/tools/dotc/parsing/Scanners.scala
index 489038f1e..7ebe63397 100644
--- a/src/dotty/tools/dotc/parsing/Scanners.scala
+++ b/src/dotty/tools/dotc/parsing/Scanners.scala
@@ -175,12 +175,39 @@ object Scanners {
   }
 
   class Scanner(source: SourceFile, override val startFrom: Offset = 0)(implicit ctx: Context) extends ScannerCommon(source)(ctx) {
-    var keepComments = false
+    val keepComments = ctx.settings.YkeepComments.value
 
-    /** All comments in the reverse order of their position in the source.
-     *  set only when `keepComments` is true.
+    /** All doc comments as encountered, each list contains doc comments from
+     *  the same block level. Starting with the deepest level and going upward
      */
-    var revComments: List[Comment] = Nil
+    private[this] var docsPerBlockStack: List[List[Comment]] = List(List())
+
+    /** Adds level of nesting to docstrings */
+    def enterBlock(): Unit =
+      docsPerBlockStack = Nil ::: docsPerBlockStack
+
+    /** Removes level of nesting for docstrings */
+    def exitBlock(): Unit = docsPerBlockStack = docsPerBlockStack match {
+      case x :: xs => List(List())
+      case _ => docsPerBlockStack.tail
+    }
+
+    /** Returns the closest docstring preceding the position supplied */
+    def getDocString(pos: Int): Option[String] = {
+      def closest(c: Comment, docstrings: List[Comment]): Comment = docstrings match {
+        case x :: xs if (c.pos.end < x.pos.end && x.pos.end <= pos) => closest(x, xs)
+        case Nil => c
+      }
+
+      docsPerBlockStack match {
+        case (list @ (x :: xs)) :: _ => {
+          val c = closest(x, xs)
+          docsPerBlockStack = list.dropWhile(_ != c).tail :: docsPerBlockStack.tail
+          Some(c.chrs)
+        }
+        case _ => None
+      }
+    }
 
     /** A buffer for comments */
     val commentBuf = new StringBuilder
@@ -487,13 +514,13 @@ object Scanners {
         case ',' =>
           nextChar(); token = COMMA
         case '(' =>
-          nextChar(); token = LPAREN
+          enterBlock(); nextChar(); token = LPAREN
         case '{' =>
-          nextChar(); token = LBRACE
+          enterBlock(); nextChar(); token = LBRACE
         case ')' =>
-          nextChar(); token = RPAREN
+          exitBlock(); nextChar(); token = RPAREN
         case '}' =>
-          nextChar(); token = RBRACE
+          exitBlock(); nextChar(); token = RBRACE
         case '[' =>
           nextChar(); token = LBRACKET
         case ']' =>
@@ -558,9 +585,12 @@ object Scanners {
       def finishComment(): Boolean = {
         if (keepComments) {
           val pos = Position(start, charOffset, start)
-          nextChar()
-          revComments = Comment(pos, flushBuf(commentBuf)) :: revComments
+          val comment = Comment(pos, flushBuf(commentBuf))
+
+          if (comment.isDocComment)
+            docsPerBlockStack = (docsPerBlockStack.head :+ comment) :: docsPerBlockStack.tail
         }
+
         true
       }
       nextChar()
diff --git a/src/dotty/tools/dotc/typer/Namer.scala b/src/dotty/tools/dotc/typer/Namer.scala
index de27333d5..82b3b56e9 100644
--- a/src/dotty/tools/dotc/typer/Namer.scala
+++ b/src/dotty/tools/dotc/typer/Namer.scala
@@ -401,19 +401,29 @@ class Namer { typer: Typer =>
       val pkg = createPackageSymbol(pcl.pid)
       index(pcl.stats)(ctx.fresh.setOwner(pkg.moduleClass))
       invalidateCompanions(pkg, Trees.flatten(pcl.stats map expanded))
+      setDocstring(pkg, stat)
       ctx
     case imp: Import =>
       importContext(createSymbol(imp), imp.selectors)
     case mdef: DefTree =>
-      enterSymbol(createSymbol(mdef))
+      val sym = enterSymbol(createSymbol(mdef))
+      setDocstring(sym, stat)
       ctx
     case stats: Thicket =>
-      for (tree <- stats.toList) enterSymbol(createSymbol(tree))
+      for (tree <- stats.toList) {
+        val sym = enterSymbol(createSymbol(tree))
+        setDocstring(sym, stat)
+      }
       ctx
     case _ =>
       ctx
   }
 
+  def setDocstring(sym: Symbol, tree: Tree)(implicit ctx: Context) = tree match {
+    case t: MemberDef => ctx.base.addDocstring(sym, t.rawComment)
+    case _ => ()
+  }
+
   /** Create top-level symbols for statements and enter them into symbol table */
   def index(stats: List[Tree])(implicit ctx: Context): Context = {
 
@@ -859,7 +869,7 @@ class Namer { typer: Typer =>
         WildcardType
     }
     paramFn(typedAheadType(mdef.tpt, tptProto).tpe)
- }
+  }
 
   /** The type signature of a DefDef with given symbol */
   def defDefSig(ddef: DefDef, sym: Symbol)(implicit ctx: Context) = {
diff --git a/src/dotty/tools/dotc/typer/Typer.scala b/src/dotty/tools/dotc/typer/Typer.scala
index 3b8ada2a8..84abf85e0 100644
--- a/src/dotty/tools/dotc/typer/Typer.scala
+++ b/src/dotty/tools/dotc/typer/Typer.scala
@@ -586,7 +586,7 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
         case _ => false
       }
 
-      /** The funcion body to be returned in the closure. Can become a TypedSplice
+      /** The function body to be returned in the closure. Can become a TypedSplice
        *  of a typed expression if this is necessary to infer a parameter type.
        */
       var fnBody = tree.body
diff --git a/test/test/DottyDocParsingTests.scala b/test/test/DottyDocParsingTests.scala
new file mode 100644
index 000000000..32cfaaddf
--- /dev/null
+++ b/test/test/DottyDocParsingTests.scala
@@ -0,0 +1,459 @@
+package test
+
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.ast.Trees._
+
+import org.junit.Assert._
+import org.junit.Test
+
+class DottyDocParsingTests extends DottyDocTest {
+
+  @Test def noComment = {
+    import dotty.tools.dotc.ast.untpd._
+    val source = "class Class"
+
+    checkFrontend(source) {
+      case PackageDef(_, Seq(c: TypeDef)) =>
+        assert(c.rawComment == None, "Should not have a comment, mainly used for exhaustive tests")
+    }
+  }
+
+  @Test def singleClassInPackage = {
+    val source =
+      """
+      |package a
+      |
+      |/** Hello world! */
+      |class Class(val x: String)
+      """.stripMargin
+
+    checkFrontend(source) {
+      case PackageDef(_, Seq(t @ TypeDef(name, _))) if name.toString == "Class" =>
+        checkDocString(t.rawComment, "/** Hello world! */")
+    }
+  }
+
+  @Test def multipleOpenedOnSingleClassInPackage = {
+    val source =
+      """
+      |package a
+      |
+      |/** Hello /* multiple open */ world! */
+      |class Class(val x: String)
+      """.stripMargin
+
+    checkFrontend(source) {
+      case PackageDef(_, Seq(t @ TypeDef(name, _))) if name.toString == "Class" =>
+        checkDocString(t.rawComment, "/** Hello /* multiple open */ world! */")
+    }
+  }
+  @Test def multipleClassesInPackage = {
+    val source =
+      """
+      |package a
+      |
+      |/** Class1 docstring */
+      |class Class1(val x: String)
+      |
+      |/** Class2 docstring */
+      |class Class2(val x: String)
+      """.stripMargin
+
+    checkCompile("frontend", source) { (_, ctx) =>
+      ctx.compilationUnit.untpdTree match {
+        case PackageDef(_, Seq(c1 @ TypeDef(_,_), c2 @ TypeDef(_,_))) => {
+          checkDocString(c1.rawComment, "/** Class1 docstring */")
+          checkDocString(c2.rawComment, "/** Class2 docstring */")
+        }
+      }
+    }
+  }
+
+  @Test def singleCaseClassWithoutPackage = {
+    val source =
+      """
+      |/** Class without package */
+      |case class Class(val x: Int)
+      """.stripMargin
+
+    checkFrontend(source) {
+      case PackageDef(_, Seq(t @ TypeDef(_,_))) => checkDocString(t.rawComment, "/** Class without package */")
+    }
+  }
+
+  @Test def SingleTraitWihoutPackage = {
+    val source = "/** Trait docstring */\ntrait Trait"
+
+    checkFrontend(source) {
+      case PackageDef(_, Seq(t @ TypeDef(_,_))) => checkDocString(t.rawComment, "/** Trait docstring */")
+    }
+  }
+
+  @Test def multipleTraitsWithoutPackage = {
+    val source =
+      """
+      |/** Trait1 docstring */
+      |trait Trait1
+      |
+      |/** Trait2 docstring */
+      |trait Trait2
+      """.stripMargin
+
+    checkFrontend(source) {
+      case PackageDef(_, Seq(t1 @ TypeDef(_,_), t2 @ TypeDef(_,_))) => {
+        checkDocString(t1.rawComment, "/** Trait1 docstring */")
+        checkDocString(t2.rawComment, "/** Trait2 docstring */")
+      }
+    }
+  }
+
+  @Test def multipleMixedEntitiesWithPackage = {
+    val source =
+      """
+      |/** Trait1 docstring */
+      |trait Trait1
+      |
+      |/** Class2 docstring */
+      |class Class2(val x: Int)
+      |
+      |/** CaseClass3 docstring */
+      |case class CaseClass3()
+      |
+      |case class NoComment()
+      |
+      |/** AbstractClass4 docstring */
+      |abstract class AbstractClass4(val x: Int)
+      """.stripMargin
+
+    checkFrontend(source) {
+      case PackageDef(_, Seq(t1 @ TypeDef(_,_), c2 @ TypeDef(_,_), cc3 @ TypeDef(_,_), _, ac4 @ TypeDef(_,_))) => {
+        checkDocString(t1.rawComment, "/** Trait1 docstring */")
+        checkDocString(c2.rawComment, "/** Class2 docstring */")
+        checkDocString(cc3.rawComment, "/** CaseClass3 docstring */")
+        checkDocString(ac4.rawComment, "/** AbstractClass4 docstring */")
+      }
+    }
+  }
+
+  @Test def nestedClass = {
+    val source =
+      """
+      |/** Outer docstring */
+      |class Outer {
+      |  /** Inner docstring */
+      |  class Inner(val x: Int)
+      |}
+      """.stripMargin
+
+    checkFrontend(source) {
+      case PackageDef(_, Seq(outer @ TypeDef(_, tpl @ Template(_,_,_,_)))) => {
+        checkDocString(outer.rawComment, "/** Outer docstring */")
+        tpl.body match {
+          case (inner @ TypeDef(_,_)) :: _ => checkDocString(inner.rawComment, "/** Inner docstring */")
+          case _ => assert(false, "Couldn't find inner class")
+        }
+      }
+    }
+  }
+
+  @Test def nestedClassThenOuter = {
+    val source =
+      """
+      |/** Outer1 docstring */
+      |class Outer1 {
+      |  /** Inner docstring */
+      |  class Inner(val x: Int)
+      |}
+      |
+      |/** Outer2 docstring */
+      |class Outer2
+      """.stripMargin
+
+    checkFrontend(source) {
+      case PackageDef(_, Seq(o1 @ TypeDef(_, tpl @ Template(_,_,_,_)), o2 @ TypeDef(_,_))) => {
+        checkDocString(o1.rawComment, "/** Outer1 docstring */")
+        checkDocString(o2.rawComment, "/** Outer2 docstring */")
+        tpl.body match {
+          case (inner @ TypeDef(_,_)) :: _ => checkDocString(inner.rawComment, "/** Inner docstring */")
+          case _ => assert(false, "Couldn't find inner class")
+        }
+      }
+    }
+  }
+
+  @Test def objects = {
+    val source =
+      """
+      |package p
+      |
+      |/** Object1 docstring */
+      |object Object1
+      |
+      |/** Object2 docstring */
+      |object Object2
+      """.stripMargin
+
+    checkFrontend(source) {
+      case p @ PackageDef(_, Seq(o1: MemberDef[Untyped], o2: MemberDef[Untyped])) => {
+        assertEquals(o1.name.toString, "Object1")
+        checkDocString(o1.rawComment, "/** Object1 docstring */")
+        assertEquals(o2.name.toString, "Object2")
+        checkDocString(o2.rawComment, "/** Object2 docstring */")
+      }
+    }
+  }
+
+  @Test def objectsNestedClass = {
+    val source =
+      """
+      |package p
+      |
+      |/** Object1 docstring */
+      |object Object1
+      |
+      |/** Object2 docstring */
+      |object Object2 {
+      |  class A1
+      |  /** Inner docstring */
+      |  class Inner
+      |}
+      """.stripMargin
+
+    import dotty.tools.dotc.ast.untpd._
+    checkFrontend(source) {
+      case p @ PackageDef(_, Seq(o1: ModuleDef, o2: ModuleDef)) => {
+        assert(o1.name.toString == "Object1")
+        checkDocString(o1.rawComment, "/** Object1 docstring */")
+        assert(o2.name.toString == "Object2")
+        checkDocString(o2.rawComment, "/** Object2 docstring */")
+
+        o2.impl.body match {
+          case _ :: (inner @ TypeDef(_,_)) :: _ => checkDocString(inner.rawComment, "/** Inner docstring */")
+          case _ => assert(false, "Couldn't find inner class")
+        }
+      }
+    }
+  }
+
+  @Test def packageObject = {
+    val source =
+      """
+      |/** Package object docstring */
+      |package object foo {
+      |  /** Boo docstring */
+      |  case class Boo()
+      |
+      |  /** Trait docstring */
+      |  trait Trait
+      |
+      |  /** InnerObject docstring */
+      |  object InnerObject {
+      |    /** InnerClass docstring */
+      |    class InnerClass
+      |  }
+      |}
+      """.stripMargin
+
+    import dotty.tools.dotc.ast.untpd._
+    checkFrontend(source) {
+      case PackageDef(_, Seq(p: ModuleDef)) => {
+        checkDocString(p.rawComment, "/** Package object docstring */")
+
+        p.impl.body match {
+          case (b: TypeDef) :: (t: TypeDef) :: (o: ModuleDef) :: Nil => {
+            checkDocString(b.rawComment, "/** Boo docstring */")
+            checkDocString(t.rawComment, "/** Trait docstring */")
+            checkDocString(o.rawComment, "/** InnerObject docstring */")
+            checkDocString(o.impl.body.head.asInstanceOf[TypeDef].rawComment, "/** InnerClass docstring */")
+          }
+          case _ => assert(false, "Incorrect structure inside package object")
+        }
+      }
+    }
+  }
+
+  @Test def multipleDocStringsBeforeEntity = {
+    val source =
+      """
+      |/** First comment */
+      |/** Second comment */
+      |/** Real comment */
+      |class Class
+      """.stripMargin
+
+    import dotty.tools.dotc.ast.untpd._
+    checkFrontend(source) {
+      case PackageDef(_, Seq(c: TypeDef)) =>
+        checkDocString(c.rawComment, "/** Real comment */")
+    }
+  }
+
+  @Test def multipleDocStringsBeforeAndAfter = {
+    val source =
+      """
+      |/** First comment */
+      |/** Second comment */
+      |/** Real comment */
+      |class Class
+      |/** Following comment 1 */
+      |/** Following comment 2 */
+      |/** Following comment 3 */
+      """.stripMargin
+
+    import dotty.tools.dotc.ast.untpd._
+    checkFrontend(source) {
+      case PackageDef(_, Seq(c: TypeDef)) =>
+        checkDocString(c.rawComment, "/** Real comment */")
+    }
+  }
+
+  @Test def valuesWithDocString = {
+    val source =
+      """
+      |object Object {
+      |  /** val1 */
+      |  val val1 = 1
+      |
+      |  /** val2 */
+      |  val val2: Int = 2
+      |  /** bogus docstring */
+      |
+      |  /** bogus docstring */
+      |  /** val3 */
+      |  val val3: List[Int] = 1 :: 2 :: 3 :: Nil
+      |}
+      """.stripMargin
+
+    import dotty.tools.dotc.ast.untpd._
+    checkFrontend(source) {
+      case PackageDef(_, Seq(o: ModuleDef)) => {
+        o.impl.body match {
+          case (v1: MemberDef) :: (v2: MemberDef) :: (v3: MemberDef) :: Nil => {
+            checkDocString(v1.rawComment, "/** val1 */")
+            checkDocString(v2.rawComment, "/** val2 */")
+            checkDocString(v3.rawComment, "/** val3 */")
+          }
+          case _ => assert(false, "Incorrect structure inside object")
+        }
+      }
+    }
+  }
+
+  @Test def varsWithDocString = {
+    val source =
+      """
+      |object Object {
+      |  /** var1 */
+      |  var var1 = 1
+      |
+      |  /** var2 */
+      |  var var2: Int = 2
+      |  /** bogus docstring */
+      |
+      |  /** bogus docstring */
+      |  /** var3 */
+      |  var var3: List[Int] = 1 :: 2 :: 3 :: Nil
+      |}
+      """.stripMargin
+
+    import dotty.tools.dotc.ast.untpd._
+    checkFrontend(source) {
+      case PackageDef(_, Seq(o: ModuleDef)) => {
+        o.impl.body match {
+          case (v1: MemberDef) :: (v2: MemberDef) :: (v3: MemberDef) :: Nil => {
+            checkDocString(v1.rawComment, "/** var1 */")
+            checkDocString(v2.rawComment, "/** var2 */")
+            checkDocString(v3.rawComment, "/** var3 */")
+          }
+          case _ => assert(false, "Incorrect structure inside object")
+        }
+      }
+    }
+  }
+
+  @Test def defsWithDocString = {
+    val source =
+      """
+      |object Object {
+      |  /** def1 */
+      |  def def1 = 1
+      |
+      |  /** def2 */
+      |  def def2: Int = 2
+      |  /** bogus docstring */
+      |
+      |  /** bogus docstring */
+      |  /** def3 */
+      |  def def3: List[Int] = 1 :: 2 :: 3 :: Nil
+      |}
+      """.stripMargin
+
+    import dotty.tools.dotc.ast.untpd._
+    checkFrontend(source) {
+      case PackageDef(_, Seq(o: ModuleDef)) => {
+        o.impl.body match {
+          case (v1: MemberDef) :: (v2: MemberDef) :: (v3: MemberDef) :: Nil => {
+            checkDocString(v1.rawComment, "/** def1 */")
+            checkDocString(v2.rawComment, "/** def2 */")
+            checkDocString(v3.rawComment, "/** def3 */")
+          }
+          case _ => assert(false, "Incorrect structure inside object")
+        }
+      }
+    }
+  }
+
+  @Test def typesWithDocString = {
+    val source =
+      """
+      |object Object {
+      |  /** type1 */
+      |  type T1 = Int
+      |
+      |  /** type2 */
+      |  type T2 = String
+      |  /** bogus docstring */
+      |
+      |  /** bogus docstring */
+      |  /** type3 */
+      |  type T3 = T2
+      |}
+      """.stripMargin
+
+    import dotty.tools.dotc.ast.untpd._
+    checkFrontend(source) {
+      case PackageDef(_, Seq(o: ModuleDef)) => {
+        o.impl.body match {
+          case (v1: MemberDef) :: (v2: MemberDef) :: (v3: MemberDef) :: Nil => {
+            checkDocString(v1.rawComment, "/** type1 */")
+            checkDocString(v2.rawComment, "/** type2 */")
+            checkDocString(v3.rawComment, "/** type3 */")
+          }
+          case _ => assert(false, "Incorrect structure inside object")
+        }
+      }
+    }
+  }
+
+  @Test def defInnerClass = {
+    val source =
+      """
+      |object Foo {
+      |  def foo() = {
+      |    /** Innermost */
+      |    class Innermost
+      |  }
+      |}
+      """.stripMargin
+
+    import dotty.tools.dotc.ast.untpd._
+    checkFrontend(source) {
+      case PackageDef(_, Seq(o: ModuleDef)) =>
+        o.impl.body match {
+          case (foo: MemberDef) :: Nil =>
+            expectNoDocString(foo.rawComment)
+          case _ => assert(false, "Incorrect structure inside object")
+        }
+    }
+  }
+} /* End class */
diff --git a/test/test/DottyDocTest.scala b/test/test/DottyDocTest.scala
new file mode 100644
index 000000000..34269ec64
--- /dev/null
+++ b/test/test/DottyDocTest.scala
@@ -0,0 +1,32 @@
+package test
+
+import dotty.tools.dotc.ast.Trees._
+import dotty.tools.dotc.core.Contexts.Context
+
+trait DottyDocTest extends DottyTest {
+  ctx = ctx.fresh.setSetting(ctx.settings.YkeepComments, true)
+
+  def checkDocString(actual: Option[String], expected: String): Unit = actual match {
+    case Some(str) =>
+      assert(str == expected, s"""Docstring: "$str" didn't match expected "$expected"""")
+    case None =>
+      assert(false, s"""No docstring found, expected: "$expected"""")
+  }
+
+  def expectNoDocString(doc: Option[String]): Unit =
+    doc.fold(()) { d => assert(false, s"""Expected not to find a docstring, but found: "$d"""") }
+
+  def defaultAssertion: PartialFunction[Any, Unit] = {
+    case t: Tree[Untyped] =>
+      assert(false, s"Couldn't match resulting AST to expected AST in: ${t.show}")
+    case x =>
+      assert(false, s"Couldn't match resulting AST to expected AST in: $x")
+  }
+
+  def checkFrontend(source: String)(docAssert: PartialFunction[Tree[Untyped], Unit]) = {
+    checkCompile("frontend", source) { (_, ctx) =>
+      implicit val c = ctx
+      (docAssert orElse defaultAssertion)(ctx.compilationUnit.untpdTree)
+    }
+  }
+}
diff --git a/test/test/DottyTest.scala b/test/test/DottyTest.scala
index 77642561a..15d82c208 100644
--- a/test/test/DottyTest.scala
+++ b/test/test/DottyTest.scala
@@ -32,7 +32,6 @@ class DottyTest /*extends ContextEscapeDetection*/ {
   }
 */
   private def compilerWithChecker(phase: String)(assertion:(tpd.Tree, Context) => Unit) = new Compiler {
-    self =>
     override def phases = {
       val allPhases = super.phases
       val targetPhase = allPhases.flatten.find(p => p.phaseName == phase).get
@@ -48,7 +47,7 @@ class DottyTest /*extends ContextEscapeDetection*/ {
     }
   }
 
-  def checkCompile(checkAfterPhase: String, source:String)(assertion:(tpd.Tree, Context) => Unit): Unit = {
+  def checkCompile(checkAfterPhase: String, source: String)(assertion: (tpd.Tree, Context) => Unit): Unit = {
     val c = compilerWithChecker(checkAfterPhase)(assertion)
     c.rootContext(ctx)
     val run = c.newRun