Add initial support for raw docstrings in ASTs

3 files changed, 41 insertions, 9 deletions

diff --git a/src/dotty/tools/dotc/ast/Trees.scala b/src/dotty/tools/dotc/ast/Trees.scala
index d0197b443..eb5ab15e3 100644
--- a/src/dotty/tools/dotc/ast/Trees.scala
+++ b/src/dotty/tools/dotc/ast/Trees.scala
@@ -321,6 +321,10 @@ object Trees {
     private[ast] def rawMods: Modifiers[T] =
       if (myMods == null) genericEmptyModifiers else myMods
 
+    private[this] var myComment: Option[String] = None
+
+    def rawComment: Option[String] = myComment
+
     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 +333,11 @@ object Trees {
 
     def withFlags(flags: FlagSet): ThisTree[Untyped] = withMods(Modifiers(flags))
 
+    def withComment(cmt: Option[String]): ThisTree[Untyped] = {
+      myComment = cmt
+      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/parsing/Parsers.scala b/src/dotty/tools/dotc/parsing/Parsers.scala
index 44a70886e..a3f15d4b7 100644
--- a/src/dotty/tools/dotc/parsing/Parsers.scala
+++ b/src/dotty/tools/dotc/parsing/Parsers.scala
@@ -1913,10 +1913,14 @@ object Parsers {
         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)
+        .withComment(in.getDocString())
     }
 
     /** ConstrMods        ::=  AccessModifier
@@ -1935,7 +1939,10 @@ object Parsers {
     def objectDef(mods: Modifiers): ModuleDef = {
       val name = ident()
       val template = templateOpt(emptyConstructor())
-      ModuleDef(name, template).withMods(mods)
+
+      ModuleDef(name, template)
+        .withMods(mods)
+        .withComment(in.getDocString())
     }
 
 /* -------- TEMPLATES ------------------------------------------- */
diff --git a/src/dotty/tools/dotc/parsing/Scanners.scala b/src/dotty/tools/dotc/parsing/Scanners.scala
index 489038f1e..0529b5cfa 100644
--- a/src/dotty/tools/dotc/parsing/Scanners.scala
+++ b/src/dotty/tools/dotc/parsing/Scanners.scala
@@ -182,6 +182,18 @@ object Scanners {
      */
     var revComments: List[Comment] = Nil
 
+    /** The currently closest docstring, replaced every time a new docstring is
+     *  encountered
+     */
+    var closestDocString: Option[Comment] = None
+
+    /** Returns `closestDocString`'s raw string and sets it to `None` */
+    def getDocString(): Option[String] = {
+      val current = closestDocString.map(_.chrs)
+      closestDocString = None
+      current
+    }
+
     /** A buffer for comments */
     val commentBuf = new StringBuilder
 
@@ -528,8 +540,7 @@ object Scanners {
     }
 
     private def skipComment(): Boolean = {
-      def appendToComment(ch: Char) =
-        if (keepComments) commentBuf.append(ch)
+      def appendToComment(ch: Char) = commentBuf.append(ch)
       def nextChar() = {
         appendToComment(ch)
         Scanner.this.nextChar()
@@ -556,11 +567,16 @@ object Scanners {
       def nestedComment() = { nextChar(); skipComment() }
       val start = lastCharOffset
       def finishComment(): Boolean = {
-        if (keepComments) {
-          val pos = Position(start, charOffset, start)
-          nextChar()
-          revComments = Comment(pos, flushBuf(commentBuf)) :: revComments
-        }
+        val pos = Position(start, charOffset, start)
+        val comment = Comment(pos, flushBuf(commentBuf))
+
+        closestDocString =
+          if (comment.isDocComment) Option(comment)
+          else closestDocString
+
+        if (keepComments)
+          revComments = comment :: revComments
+
         true
       }
       nextChar()