From 2c0d7d904b5d018e47a6d5ff78ef528c30e01d97 Mon Sep 17 00:00:00 2001
From: Martin Duhem <martin.duhem@gmail.com>
Date: Fri, 4 May 2018 11:55:59 +0200
Subject: [PATCH 1/9] Add `Comments` section in TASTY

The comments are pickled when `-Ykeep-comments` is set.
---
 compiler/src/dotty/tools/dotc/Driver.scala    |  22 ++-
 .../src/dotty/tools/dotc/core/Comments.scala  |   2 +-
 compiler/src/dotty/tools/dotc/core/Mode.scala |   4 +
 .../dotc/core/quoted/TastyUnpickler.scala     |  10 +-
 .../dotc/core/tasty/CommentPickler.scala      |  42 +++++
 .../dotc/core/tasty/CommentUnpickler.scala    |  32 ++++
 .../dotc/core/tasty/DottyUnpickler.scala      |  16 +-
 .../tools/dotc/core/tasty/TastyFormat.scala   |   4 +
 .../tools/dotc/core/tasty/TreeUnpickler.scala |  19 +-
 .../dotty/tools/dotc/transform/Pickler.scala  |   4 +
 .../tools/dotc/CommentPicklingTest.scala      | 178 ++++++++++++++++++
 11 files changed, 316 insertions(+), 17 deletions(-)
 create mode 100644 compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
 create mode 100644 compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala
 create mode 100644 compiler/test/dotty/tools/dotc/CommentPicklingTest.scala

diff --git a/compiler/src/dotty/tools/dotc/Driver.scala b/compiler/src/dotty/tools/dotc/Driver.scala
index 7349c4615938..4f3244a7b59a 100644
--- a/compiler/src/dotty/tools/dotc/Driver.scala
+++ b/compiler/src/dotty/tools/dotc/Driver.scala
@@ -2,7 +2,9 @@ package dotty.tools.dotc
 
 import dotty.tools.FatalError
 import config.CompilerCommand
+import core.Comments.{ContextDoc, ContextDocstrings}
 import core.Contexts.{Context, ContextBase}
+import core.Mode
 import util.DotClass
 import reporting._
 import scala.util.control.NonFatal
@@ -40,10 +42,24 @@ class Driver extends DotClass {
 
   protected def sourcesRequired = true
 
+  /**
+   * Should the `ContextDocstrings` be set for this context? The `ContextDocstrings` is used
+   * to store doc comments when `-Ykeep-comments` is set, or when TASTY is configured to
+   * unpickle the doc comments.
+   */
+  protected def shouldAddDocContext(implicit ctx: Context): Boolean = {
+    ctx.settings.YkeepComments.value || ctx.mode.is(Mode.ReadComments)
+  }
+
   def setup(args: Array[String], rootCtx: Context): (List[String], Context) = {
-    val ctx = rootCtx.fresh
-    val summary = CompilerCommand.distill(args)(ctx)
-    ctx.setSettings(summary.sstate)
+    val ctx0 = rootCtx.fresh
+    val summary = CompilerCommand.distill(args)(ctx0)
+    ctx0.setSettings(summary.sstate)
+
+    val ctx =
+      if (shouldAddDocContext(ctx0)) ctx0.setProperty(ContextDoc, new ContextDocstrings)
+      else ctx0
+
     val fileNames = CompilerCommand.checkUsage(summary, sourcesRequired)(ctx)
     (fileNames, ctx)
   }
diff --git a/compiler/src/dotty/tools/dotc/core/Comments.scala b/compiler/src/dotty/tools/dotc/core/Comments.scala
index c95e56bf525c..0d38791efacc 100644
--- a/compiler/src/dotty/tools/dotc/core/Comments.scala
+++ b/compiler/src/dotty/tools/dotc/core/Comments.scala
@@ -99,7 +99,7 @@ object Comments {
   }
 
   object Comment {
-    def apply(pos: Position, raw: String, expanded: Boolean = false, usc: List[UseCase] = Nil)(implicit ctx: Context): Comment =
+    def apply(pos: Position, raw: String, expanded: Boolean = false, usc: List[UseCase] = Nil): Comment =
       new Comment(pos, raw) {
         val isExpanded = expanded
         val usecases = usc
diff --git a/compiler/src/dotty/tools/dotc/core/Mode.scala b/compiler/src/dotty/tools/dotc/core/Mode.scala
index c97e20a88d49..020b884b22b3 100644
--- a/compiler/src/dotty/tools/dotc/core/Mode.scala
+++ b/compiler/src/dotty/tools/dotc/core/Mode.scala
@@ -93,4 +93,8 @@ object Mode {
 
   /** We are in the IDE */
   val Interactive = newMode(20, "Interactive")
+
+  /** Read comments from definitions when unpickling from TASTY */
+  val ReadComments = newMode(21, "ReadComments")
+
 }
diff --git a/compiler/src/dotty/tools/dotc/core/quoted/TastyUnpickler.scala b/compiler/src/dotty/tools/dotc/core/quoted/TastyUnpickler.scala
index d47294ad7d76..4bbe18662a4e 100644
--- a/compiler/src/dotty/tools/dotc/core/quoted/TastyUnpickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/quoted/TastyUnpickler.scala
@@ -4,10 +4,10 @@ import dotty.tools.dotc.core.tasty._
 import dotty.tools.dotc.core.tasty.TastyUnpickler.NameTable
 
 object TastyUnpickler {
-  class QuotedTreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], splices: Seq[Any])
-    extends DottyUnpickler.TreeSectionUnpickler(posUnpickler) {
+  class QuotedTreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], commentUnpickler: Option[CommentUnpickler], splices: Seq[Any])
+    extends DottyUnpickler.TreeSectionUnpickler(posUnpickler, commentUnpickler) {
     override def unpickle(reader: TastyReader, nameAtRef: NameTable) =
-      new TreeUnpickler(reader, nameAtRef, posUnpickler, splices)
+      new TreeUnpickler(reader, nameAtRef, posUnpickler, commentUnpickler, splices)
   }
 }
 
@@ -19,6 +19,6 @@ class TastyUnpickler(bytes: Array[Byte], splices: Seq[Any]) extends DottyUnpickl
   import DottyUnpickler._
   import TastyUnpickler._
 
-  protected override def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler]): TreeSectionUnpickler =
-    new QuotedTreeSectionUnpickler(posUnpicklerOpt, splices)
+  protected override def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler =
+    new QuotedTreeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt, splices)
 }
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
new file mode 100644
index 000000000000..3e1f55d49a78
--- /dev/null
+++ b/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
@@ -0,0 +1,42 @@
+package dotty.tools.dotc.core.tasty
+
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Comments.{Comment, CommentsContext}
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.core.tasty.TastyBuffer.Addr
+
+import java.nio.charset.Charset
+
+class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr])(implicit ctx: Context) {
+  val buf = new TastyBuffer(5000)
+  pickler.newSection("Comments", buf)
+
+  def pickleComment(root: tpd.Tree): Unit =
+    new Traverser().traverse(root)
+
+  def pickleComment(addrOfTree: Option[Addr], comment: Option[Comment]): Unit = (addrOfTree, comment) match {
+    case (Some(addr), Some(cmt)) =>
+      val bytes = cmt.raw.getBytes(Charset.forName("UTF-8"))
+      val length = bytes.length
+      buf.writeAddr(addr)
+      buf.writeNat(length)
+      buf.writeBytes(bytes, length)
+    case other =>
+      ()
+  }
+
+  private class Traverser extends tpd.TreeTraverser {
+    override def traverse(tree: tpd.Tree)(implicit ctx: Context): Unit =
+      tree match {
+        case md: tpd.MemberDef =>
+          ctx.docCtx.foreach { docCtx =>
+            val comment = docCtx.docstring(md.symbol)
+            pickleComment(addrOfTree(md), comment)
+          }
+          traverseChildren(md)
+        case _ =>
+          traverseChildren(tree)
+      }
+  }
+
+}
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala
new file mode 100644
index 000000000000..bdea47e386e4
--- /dev/null
+++ b/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala
@@ -0,0 +1,32 @@
+package dotty.tools.dotc.core.tasty
+
+import dotty.tools.dotc.core.Comments.Comment
+import dotty.tools.dotc.core.Symbols.Symbol
+import dotty.tools.dotc.core.tasty.TastyBuffer.Addr
+import dotty.tools.dotc.util.Positions
+
+import scala.collection.mutable.HashMap
+
+import java.nio.charset.Charset
+
+class CommentUnpickler(reader: TastyReader) {
+  import reader._
+
+  private[tasty] lazy val comments = {
+    val comments = new HashMap[Addr, Comment]
+    while (!isAtEnd) {
+      val addr = readAddr()
+      val length = readNat()
+      if (length > 0) {
+        val bytes = readBytes(length)
+        val rawComment = new String(bytes, Charset.forName("UTF-8"))
+        comments(addr) = Comment(Positions.NoPosition, rawComment)
+      }
+    }
+    comments.toMap
+  }
+
+  def commentAt(addr: Addr): Option[Comment] =
+    comments.get(addr)
+
+}
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala
index 91832b997877..8f2aa8710a7e 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala
@@ -17,16 +17,21 @@ object DottyUnpickler {
   /** Exception thrown if classfile is corrupted */
   class BadSignature(msg: String) extends RuntimeException(msg)
 
-  class TreeSectionUnpickler(posUnpickler: Option[PositionUnpickler])
+  class TreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], commentUnpickler: Option[CommentUnpickler])
   extends SectionUnpickler[TreeUnpickler]("ASTs") {
     def unpickle(reader: TastyReader, nameAtRef: NameTable) =
-      new TreeUnpickler(reader, nameAtRef, posUnpickler, Seq.empty)
+      new TreeUnpickler(reader, nameAtRef, posUnpickler, commentUnpickler, Seq.empty)
   }
 
   class PositionsSectionUnpickler extends SectionUnpickler[PositionUnpickler]("Positions") {
     def unpickle(reader: TastyReader, nameAtRef: NameTable) =
       new PositionUnpickler(reader)
   }
+
+  class CommentsSectionUnpickler extends SectionUnpickler[CommentUnpickler]("Comments") {
+    def unpickle(reader: TastyReader, nameAtRef: NameTable): CommentUnpickler =
+      new CommentUnpickler(reader)
+  }
 }
 
 /** A class for unpickling Tasty trees and symbols.
@@ -38,7 +43,8 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded with t
 
   val unpickler = new TastyUnpickler(bytes)
   private val posUnpicklerOpt = unpickler.unpickle(new PositionsSectionUnpickler)
-  private val treeUnpickler = unpickler.unpickle(treeSectionUnpickler(posUnpicklerOpt)).get
+  private val commentUnpicklerOpt = unpickler.unpickle(new CommentsSectionUnpickler)
+  private val treeUnpickler = unpickler.unpickle(treeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt)).get
 
   /** Enter all toplevel classes and objects into their scopes
    *  @param roots          a set of SymDenotations that should be overwritten by unpickling
@@ -52,8 +58,8 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded with t
   def unpickleTypeTree()(implicit ctx: Context): Tree =
     treeUnpickler.unpickleTypeTree()
 
-  protected def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler]): TreeSectionUnpickler = {
-    new TreeSectionUnpickler(posUnpicklerOpt)
+  protected def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler = {
+    new TreeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt)
   }
 
   protected def computeTrees(implicit ctx: Context) = treeUnpickler.unpickle()
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala b/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
index 1d659828abfe..a557c375e5f4 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
@@ -222,6 +222,10 @@ Standard Section: "Positions" Assoc*
                                             // same offset in the previously recorded node (or 0 for the first recorded node)
   Delta         = Int                       // Difference between consecutive offsets,
 
+Standard Section: "Comments" Comment*
+
+  Comment       = Length Bytes              // Raw comment's bytes encoded as UTF-8
+
 **************************************************************************************/
 
 object TastyFormat {
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
index ce955cee96c3..04333602e9a4 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
@@ -3,6 +3,7 @@ package dotc
 package core
 package tasty
 
+import Comments.CommentsContext
 import Contexts._, Symbols._, Types._, Scopes._, SymDenotations._, Names._, NameOps._
 import StdNames._, Denotations._, Flags._, Constants._, Annotations._
 import NameKinds._
@@ -25,13 +26,15 @@ import scala.quoted.Types.TreeType
 import scala.quoted.Exprs.TastyTreeExpr
 
 /** Unpickler for typed trees
- *  @param reader          the reader from which to unpickle
- *  @param tastyName       the nametable
- *  @param posUNpicklerOpt the unpickler for positions, if it exists
+ *  @param reader              the reader from which to unpickle
+ *  @param posUnpicklerOpt     the unpickler for positions, if it exists
+ *  @param commentUnpicklerOpt the unpickler for comments, if it exists
+ *  @param splices
  */
 class TreeUnpickler(reader: TastyReader,
                     nameAtRef: NameRef => TermName,
                     posUnpicklerOpt: Option[PositionUnpickler],
+                    commentUnpicklerOpt: Option[CommentUnpickler],
                     splices: Seq[Any]) {
   import TastyFormat._
   import TreeUnpickler._
@@ -820,6 +823,16 @@ class TreeUnpickler(reader: TastyReader,
         // Child annotations for local classes and enum values are not pickled, so
         // need to be re-established here.
         sym.registerIfChild(late = true)
+
+      if (ctx.mode.is(Mode.ReadComments)) {
+        for { docCtx <- ctx.docCtx
+              commentUnpickler <- commentUnpicklerOpt } {
+          val comment = commentUnpickler.commentAt(start)
+          docCtx.addDocstring(tree.symbol, comment)
+          tree.setComment(comment)
+        }
+      }
+
       tree
     }
 
diff --git a/compiler/src/dotty/tools/dotc/transform/Pickler.scala b/compiler/src/dotty/tools/dotc/transform/Pickler.scala
index b36386b9cf1a..7868ca188eff 100644
--- a/compiler/src/dotty/tools/dotc/transform/Pickler.scala
+++ b/compiler/src/dotty/tools/dotc/transform/Pickler.scala
@@ -60,6 +60,9 @@ class Pickler extends Phase {
       if (tree.pos.exists)
         new PositionPickler(pickler, treePkl.buf.addrOfTree).picklePositions(tree :: Nil)
 
+      if (ctx.settings.YkeepComments.value)
+        new CommentPickler(pickler, treePkl.buf.addrOfTree).pickleComment(tree)
+
       // other pickle sections go here.
       val pickled = pickler.assembleParts()
       unit.pickled += (cls -> pickled)
@@ -84,6 +87,7 @@ class Pickler extends Phase {
             .setPeriod(Period(ctx.runId + 1, FirstPhaseId))
             .setReporter(new ThrowingReporter(ctx.reporter))
             .addMode(Mode.ReadPositions)
+            .addMode(Mode.ReadComments)
             .addMode(Mode.PrintShowExceptions))
     result
   }
diff --git a/compiler/test/dotty/tools/dotc/CommentPicklingTest.scala b/compiler/test/dotty/tools/dotc/CommentPicklingTest.scala
new file mode 100644
index 000000000000..ec01a0f19b15
--- /dev/null
+++ b/compiler/test/dotty/tools/dotc/CommentPicklingTest.scala
@@ -0,0 +1,178 @@
+package dotty.tools.dotc
+
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Comments.CommentsContext
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.core.Decorators.PreNamedString
+import dotty.tools.dotc.core.Mode
+import dotty.tools.dotc.core.Names.Name
+import dotty.tools.dotc.core.tasty.DottyUnpickler
+import dotty.tools.dotc.interfaces.Diagnostic.ERROR
+import dotty.tools.dotc.reporting.TestReporter
+
+import dotty.tools.vulpix.{TestConfiguration, ParallelTesting}
+
+import java.io.IOException
+import java.nio.file.{FileSystems, FileVisitOption, FileVisitResult, FileVisitor, Files, Path}
+import java.nio.file.attribute.BasicFileAttributes
+import java.util.EnumSet
+
+import scala.concurrent.duration.{Duration, DurationInt}
+
+import org.junit.Test
+import org.junit.Assert.{assertEquals, assertFalse, fail}
+
+class CommentPicklingTest extends ParallelTesting {
+
+  override def isInteractive: Boolean = false
+  override def testFilter: Option[String] = None
+  override def maxDuration: Duration = 30.seconds
+  override def numberOfSlaves: Int = 5
+  override def safeMode: Boolean = false
+
+  val compileOptions = TestConfiguration.defaultOptions and "-Ykeep-comments" and "-Yemit-tasty"
+  val unpickleOptions = TestConfiguration.defaultOptions
+
+  @Test def commentOnDef: Unit = {
+    val sources = "object A { /** foo */ def bar = 2 }" :: Nil
+    compileAndCheckComment(sources, "bar".toTermName, Some("/** foo */"))
+  }
+
+  @Test def commentOnVal: Unit = {
+    val sources = "object A { /** foo */ val bar = 2 }" :: Nil
+    compileAndCheckComment(sources, "bar".toTermName, Some("/** foo */"))
+  }
+
+  @Test def commentLocalVal: Unit = {
+    val sources = "object A { def buzz = { /** foo */ val bar = 3 } }" :: Nil
+    compileAndCheckComment(sources, "bar".toTermName, Some("/** foo */"))
+  }
+
+  @Test def commentLocalDef: Unit = {
+    val sources = "object A { def buzz = { /** foo */ def bar = 5 } }" :: Nil
+    compileAndCheckComment(sources, "bar".toTermName, Some("/** foo */"))
+  }
+
+  @Test def commentOnClass: Unit = {
+    val sources = "/** foo */ class A" :: Nil
+    compileAndCheckComment(sources, "A".toTypeName, Some("/** foo */"))
+  }
+
+  @Test def commentOnObject: Unit = {
+    val sources = "/** foo */ object A" :: Nil
+    compileAndCheckComment(sources, "A".toTermName, Some("/** foo */"))
+  }
+
+  @Test def commentOnlazyVal: Unit = {
+    val sources = "class A { /** foo */ lazy val buzz = 2 }" :: Nil
+    compileAndCheckComment(sources, "buzz".toTermName, Some("/** foo */"))
+  }
+
+  private def compileAndCheckComment(sources: Seq[String], treeName: Name, expectedComment: Option[String]): Unit = {
+    compileAndUnpickle(sources) { (trees, ctx) =>
+      findTreeNamed(treeName)(trees, ctx) match {
+        case Some(md: tpd.MemberDef) =>
+          val symbol = md.symbol(ctx)
+          val comment = for { docCtx <- ctx.docCtx
+                              comment <- docCtx.docstring(symbol) } yield comment.raw
+          assertEquals(expectedComment, comment)
+        case other =>
+          fail(s"Unexpected: $other")
+      }
+    }
+  }
+
+  private def findTreeNamed(name: Name)(trees: Seq[tpd.Tree], ctx: Context): Option[tpd.NameTree] = {
+    val acc = new tpd.TreeAccumulator[Option[tpd.NameTree]] {
+      override def apply(x: Option[tpd.NameTree], tree: tpd.Tree)(implicit ctx: Context): Option[tpd.NameTree] = {
+        x.orElse(tree match {
+          case md: tpd.NameTree if md.name == name => Some(md)
+          case md: tpd.NameTree => foldOver(None, md)
+          case other => foldOver(None, other)
+        })
+      }
+    }
+    acc(None, trees)(ctx)
+  }
+
+  private def compileAndUnpickle[T](sources: Seq[String])(fn: (Seq[tpd.Tree], Context) => T) = {
+    inTempDirectory { tmp =>
+      val sourceFiles = sources.zipWithIndex.map {
+        case (src, id) =>
+          val path = tmp.resolve(s"Src$id.scala").toAbsolutePath
+          Files.write(path, src.getBytes("UTF-8"))
+          path.toString
+      }
+
+      val out = tmp.resolve("out")
+      Files.createDirectories(out)
+
+      val options = compileOptions.and("-d", out.toAbsolutePath.toString).and(sourceFiles: _*)
+      val driver = new Driver
+      val reporter = TestReporter.reporter(System.out, logLevel = ERROR)
+      driver.process(options.all, reporter)
+      assertFalse("Compilation failed.", reporter.hasErrors)
+
+      val tastyFiles = getAll(tmp, "glob:**.tasty")
+      val unpicklingOptions = unpickleOptions
+        .withClasspath(out.toAbsolutePath.toString)
+        .and("dummy") // Need to pass a dummy source file name
+      val unpicklingDriver = new UnpicklingDriver
+      unpicklingDriver.unpickle(unpicklingOptions.all, tastyFiles)(fn)
+    }
+  }
+
+  private class UnpicklingDriver extends Driver {
+    override def initCtx = super.initCtx.addMode(Mode.ReadComments)
+    def unpickle[T](args: Array[String], paths: Seq[Path])(fn: (Seq[tpd.Tree], Context) => T): T = {
+      implicit val (_, ctx: Context) = setup(args, initCtx)
+      ctx.initialize()
+      val trees = paths.flatMap { p =>
+        val bytes = Files.readAllBytes(p)
+        val unpickler = new DottyUnpickler(bytes)
+        unpickler.enter(roots = Set.empty)
+        unpickler.trees(ctx)
+      }
+      fn(trees, ctx)
+    }
+  }
+
+  private def inTempDirectory[T](fn: Path => T): T = {
+    val temp = Files.createTempDirectory("temp")
+    try fn(temp)
+    finally {
+      val allFiles = getAll(temp, "glob:**").sortBy(_.toAbsolutePath.toString).reverse
+      allFiles.foreach(Files.delete(_))
+    }
+  }
+
+  private def getAll(base: Path,
+             pattern: String,
+             maxDepth: Int = Int.MaxValue): Seq[Path] = {
+    val out = collection.mutable.ListBuffer.empty[Path]
+    val matcher = FileSystems.getDefault.getPathMatcher(pattern)
+    val visitor = new FileVisitor[Path] {
+      override def preVisitDirectory(directory: Path, attributes: BasicFileAttributes): FileVisitResult = {
+        if (matcher.matches(directory)) out += directory
+        FileVisitResult.CONTINUE
+      }
+
+      override def postVisitDirectory(directory: Path, exception: IOException): FileVisitResult =
+        FileVisitResult.CONTINUE
+
+      override def visitFile(file: Path, attributes: BasicFileAttributes): FileVisitResult = {
+        if (matcher.matches(file)) out += file
+        FileVisitResult.CONTINUE
+      }
+
+      override def visitFileFailed(file: Path, exception: IOException): FileVisitResult =
+        FileVisitResult.CONTINUE
+    }
+    Files.walkFileTree(base,
+                       EnumSet.of(FileVisitOption.FOLLOW_LINKS),
+                       maxDepth,
+                       visitor)
+    out
+  }
+
+}

From dbd674983dabe5c33942fdab242b48874d0644ca Mon Sep 17 00:00:00 2001
From: Martin Duhem <martin.duhem@gmail.com>
Date: Fri, 4 May 2018 14:48:08 +0200
Subject: [PATCH 2/9] Store whether comment is expanded in TASTY

---
 compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala | 1 +
 .../src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala    | 3 ++-
 compiler/src/dotty/tools/dotc/core/tasty/TastyBuffer.scala    | 4 ++++
 compiler/src/dotty/tools/dotc/core/tasty/TastyReader.scala    | 4 ++++
 4 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
index 3e1f55d49a78..626e89cae314 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
@@ -21,6 +21,7 @@ class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr]
       buf.writeAddr(addr)
       buf.writeNat(length)
       buf.writeBytes(bytes, length)
+      buf.writeBoolean(cmt.isExpanded)
     case other =>
       ()
   }
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala
index bdea47e386e4..3baa5d17a3e3 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala
@@ -19,8 +19,9 @@ class CommentUnpickler(reader: TastyReader) {
       val length = readNat()
       if (length > 0) {
         val bytes = readBytes(length)
+        val expanded = readBoolean()
         val rawComment = new String(bytes, Charset.forName("UTF-8"))
-        comments(addr) = Comment(Positions.NoPosition, rawComment)
+        comments(addr) = Comment(Positions.NoPosition, rawComment, expanded = expanded)
       }
     }
     comments.toMap
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TastyBuffer.scala b/compiler/src/dotty/tools/dotc/core/tasty/TastyBuffer.scala
index 4c092c22e56d..e263f732d7b4 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TastyBuffer.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TastyBuffer.scala
@@ -48,6 +48,10 @@ class TastyBuffer(initialSize: Int) {
 
   // -- Output routines --------------------------------------------
 
+  /** Write a boolean value. */
+  def writeBoolean(b: Boolean): Unit =
+    writeByte(if (b) 1 else 0)
+
   /** Write a byte of data. */
   def writeByte(b: Int): Unit = {
     if (length >= bytes.length)
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TastyReader.scala b/compiler/src/dotty/tools/dotc/core/tasty/TastyReader.scala
index 938c85e10d1c..9b9886c970be 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TastyReader.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TastyReader.scala
@@ -41,6 +41,10 @@ class TastyReader(val bytes: Array[Byte], start: Int, end: Int, val base: Int =
   def subReader(start: Addr, end: Addr): TastyReader =
     new TastyReader(bytes, index(start), index(end), base)
 
+  /** Read a boolean value. */
+  def readBoolean(): Boolean =
+    readByte() == 1
+
   /** Read a byte of data. */
   def readByte(): Int = {
     val result = bytes(bp) & 0xff

From 905a178edb0b162a4ea126aa7d721fc9069d9460 Mon Sep 17 00:00:00 2001
From: Martin Duhem <martin.duhem@gmail.com>
Date: Wed, 23 May 2018 10:04:48 +0200
Subject: [PATCH 3/9] Cleanup tests

---
 .../tasty}/CommentPicklingTest.scala          | 53 +++++--------------
 1 file changed, 13 insertions(+), 40 deletions(-)
 rename compiler/test/dotty/tools/dotc/{ => core/tasty}/CommentPicklingTest.scala (71%)

diff --git a/compiler/test/dotty/tools/dotc/CommentPicklingTest.scala b/compiler/test/dotty/tools/dotc/core/tasty/CommentPicklingTest.scala
similarity index 71%
rename from compiler/test/dotty/tools/dotc/CommentPicklingTest.scala
rename to compiler/test/dotty/tools/dotc/core/tasty/CommentPicklingTest.scala
index ec01a0f19b15..461f6a610d8e 100644
--- a/compiler/test/dotty/tools/dotc/CommentPicklingTest.scala
+++ b/compiler/test/dotty/tools/dotc/core/tasty/CommentPicklingTest.scala
@@ -1,12 +1,13 @@
-package dotty.tools.dotc
+package dotty.tools.dotc.core.tasty
 
 import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.ast.tpd.TreeOps
+import dotty.tools.dotc.Driver
 import dotty.tools.dotc.core.Comments.CommentsContext
 import dotty.tools.dotc.core.Contexts.Context
 import dotty.tools.dotc.core.Decorators.PreNamedString
 import dotty.tools.dotc.core.Mode
 import dotty.tools.dotc.core.Names.Name
-import dotty.tools.dotc.core.tasty.DottyUnpickler
 import dotty.tools.dotc.interfaces.Diagnostic.ERROR
 import dotty.tools.dotc.reporting.TestReporter
 
@@ -17,18 +18,13 @@ import java.nio.file.{FileSystems, FileVisitOption, FileVisitResult, FileVisitor
 import java.nio.file.attribute.BasicFileAttributes
 import java.util.EnumSet
 
+import scala.collection.JavaConverters.asScalaIteratorConverter
 import scala.concurrent.duration.{Duration, DurationInt}
 
 import org.junit.Test
 import org.junit.Assert.{assertEquals, assertFalse, fail}
 
-class CommentPicklingTest extends ParallelTesting {
-
-  override def isInteractive: Boolean = false
-  override def testFilter: Option[String] = None
-  override def maxDuration: Duration = 30.seconds
-  override def numberOfSlaves: Int = 5
-  override def safeMode: Boolean = false
+class CommentPicklingTest {
 
   val compileOptions = TestConfiguration.defaultOptions and "-Ykeep-comments" and "-Yemit-tasty"
   val unpickleOptions = TestConfiguration.defaultOptions
@@ -68,7 +64,7 @@ class CommentPicklingTest extends ParallelTesting {
     compileAndCheckComment(sources, "buzz".toTermName, Some("/** foo */"))
   }
 
-  private def compileAndCheckComment(sources: Seq[String], treeName: Name, expectedComment: Option[String]): Unit = {
+  private def compileAndCheckComment(sources: List[String], treeName: Name, expectedComment: Option[String]): Unit = {
     compileAndUnpickle(sources) { (trees, ctx) =>
       findTreeNamed(treeName)(trees, ctx) match {
         case Some(md: tpd.MemberDef) =>
@@ -82,12 +78,11 @@ class CommentPicklingTest extends ParallelTesting {
     }
   }
 
-  private def findTreeNamed(name: Name)(trees: Seq[tpd.Tree], ctx: Context): Option[tpd.NameTree] = {
+  private def findTreeNamed(name: Name)(trees: List[tpd.Tree], ctx: Context): Option[tpd.NameTree] = {
     val acc = new tpd.TreeAccumulator[Option[tpd.NameTree]] {
       override def apply(x: Option[tpd.NameTree], tree: tpd.Tree)(implicit ctx: Context): Option[tpd.NameTree] = {
         x.orElse(tree match {
           case md: tpd.NameTree if md.name == name => Some(md)
-          case md: tpd.NameTree => foldOver(None, md)
           case other => foldOver(None, other)
         })
       }
@@ -95,7 +90,7 @@ class CommentPicklingTest extends ParallelTesting {
     acc(None, trees)(ctx)
   }
 
-  private def compileAndUnpickle[T](sources: Seq[String])(fn: (Seq[tpd.Tree], Context) => T) = {
+  private def compileAndUnpickle[T](sources: List[String])(fn: (List[tpd.Tree], Context) => T) = {
     inTempDirectory { tmp =>
       val sourceFiles = sources.zipWithIndex.map {
         case (src, id) =>
@@ -124,7 +119,7 @@ class CommentPicklingTest extends ParallelTesting {
 
   private class UnpicklingDriver extends Driver {
     override def initCtx = super.initCtx.addMode(Mode.ReadComments)
-    def unpickle[T](args: Array[String], paths: Seq[Path])(fn: (Seq[tpd.Tree], Context) => T): T = {
+    def unpickle[T](args: Array[String], paths: List[Path])(fn: (List[tpd.Tree], Context) => T): T = {
       implicit val (_, ctx: Context) = setup(args, initCtx)
       ctx.initialize()
       val trees = paths.flatMap { p =>
@@ -146,33 +141,11 @@ class CommentPicklingTest extends ParallelTesting {
     }
   }
 
-  private def getAll(base: Path,
-             pattern: String,
-             maxDepth: Int = Int.MaxValue): Seq[Path] = {
-    val out = collection.mutable.ListBuffer.empty[Path]
+  private def getAll(base: Path, pattern: String): List[Path] = {
+    val paths = Files.walk(base)
     val matcher = FileSystems.getDefault.getPathMatcher(pattern)
-    val visitor = new FileVisitor[Path] {
-      override def preVisitDirectory(directory: Path, attributes: BasicFileAttributes): FileVisitResult = {
-        if (matcher.matches(directory)) out += directory
-        FileVisitResult.CONTINUE
-      }
-
-      override def postVisitDirectory(directory: Path, exception: IOException): FileVisitResult =
-        FileVisitResult.CONTINUE
-
-      override def visitFile(file: Path, attributes: BasicFileAttributes): FileVisitResult = {
-        if (matcher.matches(file)) out += file
-        FileVisitResult.CONTINUE
-      }
-
-      override def visitFileFailed(file: Path, exception: IOException): FileVisitResult =
-        FileVisitResult.CONTINUE
-    }
-    Files.walkFileTree(base,
-                       EnumSet.of(FileVisitOption.FOLLOW_LINKS),
-                       maxDepth,
-                       visitor)
-    out
+    try paths.filter(matcher.matches).iterator().asScala.toList
+    finally paths.close()
   }
 
 }

From afbc23fe40f9e1b76058e9f2ed867f753bf1fdfb Mon Sep 17 00:00:00 2001
From: Martin Duhem <martin.duhem@gmail.com>
Date: Wed, 23 May 2018 10:08:59 +0200
Subject: [PATCH 4/9] Simplify `Driver.setup`

---
 compiler/src/dotty/tools/dotc/Driver.scala | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/Driver.scala b/compiler/src/dotty/tools/dotc/Driver.scala
index 4f3244a7b59a..29109405403a 100644
--- a/compiler/src/dotty/tools/dotc/Driver.scala
+++ b/compiler/src/dotty/tools/dotc/Driver.scala
@@ -52,13 +52,11 @@ class Driver extends DotClass {
   }
 
   def setup(args: Array[String], rootCtx: Context): (List[String], Context) = {
-    val ctx0 = rootCtx.fresh
-    val summary = CompilerCommand.distill(args)(ctx0)
-    ctx0.setSettings(summary.sstate)
+    val ctx = rootCtx.fresh
+    val summary = CompilerCommand.distill(args)(ctx)
+    ctx.setSettings(summary.sstate)
 
-    val ctx =
-      if (shouldAddDocContext(ctx0)) ctx0.setProperty(ContextDoc, new ContextDocstrings)
-      else ctx0
+    if (shouldAddDocContext(ctx)) ctx.setProperty(ContextDoc, new ContextDocstrings)
 
     val fileNames = CompilerCommand.checkUsage(summary, sourcesRequired)(ctx)
     (fileNames, ctx)

From deb9dc8d3017d462144e5f69f5a0e367d06c3275 Mon Sep 17 00:00:00 2001
From: Martin Duhem <martin.duhem@gmail.com>
Date: Thu, 24 May 2018 07:38:59 +0200
Subject: [PATCH 5/9] Keep comments by default, add `-Ydrop-comments`

This commit removes `-Ykeep-comments` and adds a new `-Ydrop-comments`
flag. `-Ydrop-comments` is used to prevent the parser from keeping the
comments, and defaults to `false`.
---
 compiler/src/dotty/tools/dotc/Driver.scala                 | 4 ++--
 compiler/src/dotty/tools/dotc/config/ScalaSettings.scala   | 2 +-
 compiler/src/dotty/tools/dotc/parsing/Scanners.scala       | 2 +-
 compiler/src/dotty/tools/dotc/transform/Pickler.scala      | 2 +-
 compiler/test/dotty/tools/dotc/parsing/DocstringTest.scala | 1 -
 doc-tool/src/dotty/tools/dottydoc/DocDriver.scala          | 1 -
 doc-tool/test/DottyDocTest.scala                           | 1 -
 doc-tool/test/MarkdownTests.scala                          | 1 -
 8 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/Driver.scala b/compiler/src/dotty/tools/dotc/Driver.scala
index 29109405403a..5fb5ba54db6d 100644
--- a/compiler/src/dotty/tools/dotc/Driver.scala
+++ b/compiler/src/dotty/tools/dotc/Driver.scala
@@ -44,11 +44,11 @@ class Driver extends DotClass {
 
   /**
    * Should the `ContextDocstrings` be set for this context? The `ContextDocstrings` is used
-   * to store doc comments when `-Ykeep-comments` is set, or when TASTY is configured to
+   * to store doc comments unless `-Ydrop-comments` is set, or when TASTY is configured to
    * unpickle the doc comments.
    */
   protected def shouldAddDocContext(implicit ctx: Context): Boolean = {
-    ctx.settings.YkeepComments.value || ctx.mode.is(Mode.ReadComments)
+    !ctx.settings.YdropComments.value || ctx.mode.is(Mode.ReadComments)
   }
 
   def setup(args: Array[String], rootCtx: Context): (List[String], Context) = {
diff --git a/compiler/src/dotty/tools/dotc/config/ScalaSettings.scala b/compiler/src/dotty/tools/dotc/config/ScalaSettings.scala
index c420d9e8e23a..d8ba91ed7c24 100644
--- a/compiler/src/dotty/tools/dotc/config/ScalaSettings.scala
+++ b/compiler/src/dotty/tools/dotc/config/ScalaSettings.scala
@@ -116,7 +116,7 @@ class ScalaSettings extends Settings.SettingGroup {
   val YshowPrintErrors = BooleanSetting("-Yshow-print-errors", "don't suppress exceptions thrown during tree printing.")
   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.")
+  val YdropComments = BooleanSetting("-Ydrop-comments", "Drop comments when scanning source files.")
   val YcookComments = BooleanSetting("-Ycook-comments", "Cook the comments (type check `@usecase`, etc.)")
   val YforceSbtPhases = BooleanSetting("-Yforce-sbt-phases", "Run the phases used by sbt for incremental compilation (ExtractDependencies and ExtractAPI) even if the compiler is ran outside of sbt, for debugging.")
   val YdumpSbtInc = BooleanSetting("-Ydump-sbt-inc", "For every compiled foo.scala, output the API representation and dependencies used for sbt incremental compilation in foo.inc, implies -Yforce-sbt-phases.")
diff --git a/compiler/src/dotty/tools/dotc/parsing/Scanners.scala b/compiler/src/dotty/tools/dotc/parsing/Scanners.scala
index 877410ddf2ef..68443e5b2125 100644
--- a/compiler/src/dotty/tools/dotc/parsing/Scanners.scala
+++ b/compiler/src/dotty/tools/dotc/parsing/Scanners.scala
@@ -172,7 +172,7 @@ object Scanners {
   }
 
   class Scanner(source: SourceFile, override val startFrom: Offset = 0)(implicit ctx: Context) extends ScannerCommon(source)(ctx) {
-    val keepComments = ctx.settings.YkeepComments.value
+    val keepComments = !ctx.settings.YdropComments.value
 
     /** All doc comments kept by their end position in a `Map` */
     private[this] var docstringMap: SortedMap[Int, Comment] = SortedMap.empty
diff --git a/compiler/src/dotty/tools/dotc/transform/Pickler.scala b/compiler/src/dotty/tools/dotc/transform/Pickler.scala
index 7868ca188eff..a1772c9ca7f8 100644
--- a/compiler/src/dotty/tools/dotc/transform/Pickler.scala
+++ b/compiler/src/dotty/tools/dotc/transform/Pickler.scala
@@ -60,7 +60,7 @@ class Pickler extends Phase {
       if (tree.pos.exists)
         new PositionPickler(pickler, treePkl.buf.addrOfTree).picklePositions(tree :: Nil)
 
-      if (ctx.settings.YkeepComments.value)
+      if (!ctx.settings.YdropComments.value)
         new CommentPickler(pickler, treePkl.buf.addrOfTree).pickleComment(tree)
 
       // other pickle sections go here.
diff --git a/compiler/test/dotty/tools/dotc/parsing/DocstringTest.scala b/compiler/test/dotty/tools/dotc/parsing/DocstringTest.scala
index 53fddec0bf02..79d7bc055617 100644
--- a/compiler/test/dotty/tools/dotc/parsing/DocstringTest.scala
+++ b/compiler/test/dotty/tools/dotc/parsing/DocstringTest.scala
@@ -6,7 +6,6 @@ import ast.Trees._
 import core.Contexts.Context
 
 trait DocstringTest extends DottyTest {
-  ctx = ctx.fresh.setSetting(ctx.settings.YkeepComments, true)
 
   def checkDocString(actual: Option[String], expected: String): Unit = actual match {
     case Some(str) =>
diff --git a/doc-tool/src/dotty/tools/dottydoc/DocDriver.scala b/doc-tool/src/dotty/tools/dottydoc/DocDriver.scala
index 42eaea8a259b..aacbdcdc3e80 100644
--- a/doc-tool/src/dotty/tools/dottydoc/DocDriver.scala
+++ b/doc-tool/src/dotty/tools/dottydoc/DocDriver.scala
@@ -22,7 +22,6 @@ class DocDriver extends Driver {
     val summary = CompilerCommand.distill(args)(ctx)
 
     ctx.setSettings(summary.sstate)
-    ctx.setSetting(ctx.settings.YkeepComments, true)
     ctx.setSetting(ctx.settings.YcookComments, true)
     ctx.setSetting(ctx.settings.YnoInline, true)
     ctx.setProperty(ContextDoc, new ContextDottydoc)
diff --git a/doc-tool/test/DottyDocTest.scala b/doc-tool/test/DottyDocTest.scala
index df85aeab0f9f..30c7ec2cbf79 100644
--- a/doc-tool/test/DottyDocTest.scala
+++ b/doc-tool/test/DottyDocTest.scala
@@ -23,7 +23,6 @@ trait DottyDocTest extends MessageRendering {
     import base.settings._
     val ctx = base.initialCtx.fresh
     ctx.setSetting(ctx.settings.language, List("Scala2"))
-    ctx.setSetting(ctx.settings.YkeepComments, true)
     ctx.setSetting(ctx.settings.YcookComments, true)
     ctx.setSetting(ctx.settings.Ycheck, "all" :: Nil)
     ctx.setSetting(ctx.settings.YnoInline, true)
diff --git a/doc-tool/test/MarkdownTests.scala b/doc-tool/test/MarkdownTests.scala
index b9314d015a0b..5d392ada1379 100644
--- a/doc-tool/test/MarkdownTests.scala
+++ b/doc-tool/test/MarkdownTests.scala
@@ -15,7 +15,6 @@ class MarkdownTests extends DottyDocTest {
     import base.settings._
     val ctx = base.initialCtx.fresh
     ctx.setSetting(ctx.settings.language, List("Scala2"))
-    ctx.setSetting(ctx.settings.YkeepComments, true)
     ctx.setSetting(ctx.settings.YnoInline, true)
     ctx.setSetting(ctx.settings.Ycheck, "all" :: Nil)
     // No wiki syntax!

From f29e03bce48fa0a493fd87f2b91bcadb31e04401 Mon Sep 17 00:00:00 2001
From: Martin Duhem <martin.duhem@gmail.com>
Date: Wed, 13 Jun 2018 10:34:02 +0200
Subject: [PATCH 6/9] Address review comments

---
 compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala | 4 ++--
 .../src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala    | 2 +-
 compiler/src/dotty/tools/dotc/core/tasty/TastyBuffer.scala    | 4 ----
 compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala    | 4 +++-
 compiler/src/dotty/tools/dotc/core/tasty/TastyReader.scala    | 4 ----
 5 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
index 626e89cae314..ac7b801a4a74 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
@@ -8,7 +8,7 @@ import dotty.tools.dotc.core.tasty.TastyBuffer.Addr
 import java.nio.charset.Charset
 
 class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr])(implicit ctx: Context) {
-  val buf = new TastyBuffer(5000)
+  private[this] val buf = new TastyBuffer(5000)
   pickler.newSection("Comments", buf)
 
   def pickleComment(root: tpd.Tree): Unit =
@@ -21,7 +21,7 @@ class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr]
       buf.writeAddr(addr)
       buf.writeNat(length)
       buf.writeBytes(bytes, length)
-      buf.writeBoolean(cmt.isExpanded)
+      buf.writeByte(if (cmt.isExpanded) 1 else 0)
     case other =>
       ()
   }
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala
index 3baa5d17a3e3..158794dff72c 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala
@@ -19,7 +19,7 @@ class CommentUnpickler(reader: TastyReader) {
       val length = readNat()
       if (length > 0) {
         val bytes = readBytes(length)
-        val expanded = readBoolean()
+        val expanded = readByte() == 1
         val rawComment = new String(bytes, Charset.forName("UTF-8"))
         comments(addr) = Comment(Positions.NoPosition, rawComment, expanded = expanded)
       }
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TastyBuffer.scala b/compiler/src/dotty/tools/dotc/core/tasty/TastyBuffer.scala
index e263f732d7b4..4c092c22e56d 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TastyBuffer.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TastyBuffer.scala
@@ -48,10 +48,6 @@ class TastyBuffer(initialSize: Int) {
 
   // -- Output routines --------------------------------------------
 
-  /** Write a boolean value. */
-  def writeBoolean(b: Boolean): Unit =
-    writeByte(if (b) 1 else 0)
-
   /** Write a byte of data. */
   def writeByte(b: Int): Unit = {
     if (length >= bytes.length)
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala b/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
index a557c375e5f4..f1237189459c 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
@@ -224,7 +224,9 @@ Standard Section: "Positions" Assoc*
 
 Standard Section: "Comments" Comment*
 
-  Comment       = Length Bytes              // Raw comment's bytes encoded as UTF-8
+  Comment       = Length Bytes Byte         // Raw comment's bytes encoded as UTF-8, plus a byte indicating
+                                            // whether the comment is expanded (1) or not expanded (0)
+
 
 **************************************************************************************/
 
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TastyReader.scala b/compiler/src/dotty/tools/dotc/core/tasty/TastyReader.scala
index 9b9886c970be..938c85e10d1c 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TastyReader.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TastyReader.scala
@@ -41,10 +41,6 @@ class TastyReader(val bytes: Array[Byte], start: Int, end: Int, val base: Int =
   def subReader(start: Addr, end: Addr): TastyReader =
     new TastyReader(bytes, index(start), index(end), base)
 
-  /** Read a boolean value. */
-  def readBoolean(): Boolean =
-    readByte() == 1
-
   /** Read a byte of data. */
   def readByte(): Int = {
     val result = bytes(bp) & 0xff

From 236b931f58fbd1d8a673cba4c0fb587885a78c76 Mon Sep 17 00:00:00 2001
From: Martin Duhem <martin.duhem@gmail.com>
Date: Wed, 13 Jun 2018 10:53:52 +0200
Subject: [PATCH 7/9] Unpickle comments in TastyPrinter

---
 .../dotty/tools/dotc/core/tasty/TastyPrinter.scala | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TastyPrinter.scala b/compiler/src/dotty/tools/dotc/core/tasty/TastyPrinter.scala
index c4ee8cffcf05..3a0442d7cbc6 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TastyPrinter.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TastyPrinter.scala
@@ -34,6 +34,7 @@ class TastyPrinter(bytes: Array[Byte])(implicit ctx: Context) {
     println("Trees:")
     unpickle(new TreeSectionUnpickler)
     unpickle(new PositionSectionUnpickler)
+    unpickle(new CommentSectionUnpickler)
   }
 
   class TreeSectionUnpickler extends SectionUnpickler[Unit]("ASTs") {
@@ -120,6 +121,19 @@ class TastyPrinter(bytes: Array[Byte])(implicit ctx: Context) {
     }
   }
 
+  class CommentSectionUnpickler extends SectionUnpickler[Unit]("Comments") {
+    def unpickle(reader: TastyReader, tastyName: NameTable): Unit = {
+      print(s" ${reader.endAddr.index - reader.currentAddr.index}")
+      val comments = new CommentUnpickler(reader).comments
+      println(s" comment bytes:")
+      val sorted = comments.toSeq.sortBy(_._1.index)
+      for ((addr, cmt) <- sorted) {
+        print(treeColor("%10d".format(addr.index)))
+        println(s": ${cmt.raw} (expanded = ${cmt.isExpanded})")
+      }
+    }
+  }
+
   private def nameColor(str: String): String = Magenta(str).show
   private def treeColor(str: String): String = Yellow(str).show
   private def lengthColor(str: String): String = Cyan(str).show

From 4eed3f1d3907768cefe510668c32e5b05e859ab4 Mon Sep 17 00:00:00 2001
From: Martin Duhem <martin.duhem@gmail.com>
Date: Tue, 19 Jun 2018 09:04:56 -0400
Subject: [PATCH 8/9] Address review comments

---
 compiler/src/dotty/tools/dotc/Driver.scala    | 13 +++---------
 .../dotc/core/quoted/TastyUnpickler.scala     | 10 +++++-----
 .../dotc/core/tasty/CommentPickler.scala      | 16 +++++++--------
 .../tools/dotc/core/tasty/TreeUnpickler.scala |  6 +++---
 .../dotc/core/tasty/CommentPicklingTest.scala | 20 +++++++------------
 5 files changed, 26 insertions(+), 39 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/Driver.scala b/compiler/src/dotty/tools/dotc/Driver.scala
index 5fb5ba54db6d..a1e986394fa6 100644
--- a/compiler/src/dotty/tools/dotc/Driver.scala
+++ b/compiler/src/dotty/tools/dotc/Driver.scala
@@ -42,21 +42,14 @@ class Driver extends DotClass {
 
   protected def sourcesRequired = true
 
-  /**
-   * Should the `ContextDocstrings` be set for this context? The `ContextDocstrings` is used
-   * to store doc comments unless `-Ydrop-comments` is set, or when TASTY is configured to
-   * unpickle the doc comments.
-   */
-  protected def shouldAddDocContext(implicit ctx: Context): Boolean = {
-    !ctx.settings.YdropComments.value || ctx.mode.is(Mode.ReadComments)
-  }
-
   def setup(args: Array[String], rootCtx: Context): (List[String], Context) = {
     val ctx = rootCtx.fresh
     val summary = CompilerCommand.distill(args)(ctx)
     ctx.setSettings(summary.sstate)
 
-    if (shouldAddDocContext(ctx)) ctx.setProperty(ContextDoc, new ContextDocstrings)
+    if (!ctx.settings.YdropComments.value(ctx) || ctx.mode.is(Mode.ReadComments)) {
+      ctx.setProperty(ContextDoc, new ContextDocstrings)
+    }
 
     val fileNames = CompilerCommand.checkUsage(summary, sourcesRequired)(ctx)
     (fileNames, ctx)
diff --git a/compiler/src/dotty/tools/dotc/core/quoted/TastyUnpickler.scala b/compiler/src/dotty/tools/dotc/core/quoted/TastyUnpickler.scala
index 4bbe18662a4e..ca58798c4630 100644
--- a/compiler/src/dotty/tools/dotc/core/quoted/TastyUnpickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/quoted/TastyUnpickler.scala
@@ -4,14 +4,14 @@ import dotty.tools.dotc.core.tasty._
 import dotty.tools.dotc.core.tasty.TastyUnpickler.NameTable
 
 object TastyUnpickler {
-  class QuotedTreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], commentUnpickler: Option[CommentUnpickler], splices: Seq[Any])
-    extends DottyUnpickler.TreeSectionUnpickler(posUnpickler, commentUnpickler) {
+  class QuotedTreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], splices: Seq[Any])
+    extends DottyUnpickler.TreeSectionUnpickler(posUnpickler, None) {
     override def unpickle(reader: TastyReader, nameAtRef: NameTable) =
-      new TreeUnpickler(reader, nameAtRef, posUnpickler, commentUnpickler, splices)
+      new TreeUnpickler(reader, nameAtRef, posUnpickler, None, splices)
   }
 }
 
-/** A class for unpickling quoted Tasty trees and symbols.
+/** A class for unpickling quoted Tasty trees and symbols. Comments are never unpickled.
  *  @param bytes         the bytearray containing the Tasty file from which we unpickle
  *  @param splices       splices that will fill the holes in the quote
  */
@@ -20,5 +20,5 @@ class TastyUnpickler(bytes: Array[Byte], splices: Seq[Any]) extends DottyUnpickl
   import TastyUnpickler._
 
   protected override def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler =
-    new QuotedTreeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt, splices)
+    new QuotedTreeSectionUnpickler(posUnpicklerOpt, splices)
 }
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
index ac7b801a4a74..b32efeb7b4f2 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala
@@ -1,7 +1,7 @@
 package dotty.tools.dotc.core.tasty
 
 import dotty.tools.dotc.ast.tpd
-import dotty.tools.dotc.core.Comments.{Comment, CommentsContext}
+import dotty.tools.dotc.core.Comments.{Comment, CommentsContext, ContextDocstrings}
 import dotty.tools.dotc.core.Contexts.Context
 import dotty.tools.dotc.core.tasty.TastyBuffer.Addr
 
@@ -11,8 +11,10 @@ class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr]
   private[this] val buf = new TastyBuffer(5000)
   pickler.newSection("Comments", buf)
 
-  def pickleComment(root: tpd.Tree): Unit =
-    new Traverser().traverse(root)
+  def pickleComment(root: tpd.Tree): Unit = {
+    assert(ctx.docCtx.isDefined, "Trying to pickle comments, but there's no `docCtx`.")
+    new Traverser(ctx.docCtx.get).traverse(root)
+  }
 
   def pickleComment(addrOfTree: Option[Addr], comment: Option[Comment]): Unit = (addrOfTree, comment) match {
     case (Some(addr), Some(cmt)) =>
@@ -26,14 +28,12 @@ class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr]
       ()
   }
 
-  private class Traverser extends tpd.TreeTraverser {
+  private class Traverser(docCtx: ContextDocstrings) extends tpd.TreeTraverser {
     override def traverse(tree: tpd.Tree)(implicit ctx: Context): Unit =
       tree match {
         case md: tpd.MemberDef =>
-          ctx.docCtx.foreach { docCtx =>
-            val comment = docCtx.docstring(md.symbol)
-            pickleComment(addrOfTree(md), comment)
-          }
+          val comment = docCtx.docstring(md.symbol)
+          pickleComment(addrOfTree(md), comment)
           traverseChildren(md)
         case _ =>
           traverseChildren(tree)
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
index 04333602e9a4..a032f1dfdb84 100644
--- a/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
+++ b/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
@@ -825,10 +825,10 @@ class TreeUnpickler(reader: TastyReader,
         sym.registerIfChild(late = true)
 
       if (ctx.mode.is(Mode.ReadComments)) {
-        for { docCtx <- ctx.docCtx
-              commentUnpickler <- commentUnpicklerOpt } {
+        assert(ctx.docCtx.isDefined, "Mode is `ReadComments`, but no `docCtx` is set.")
+        commentUnpicklerOpt.foreach { commentUnpickler =>
           val comment = commentUnpickler.commentAt(start)
-          docCtx.addDocstring(tree.symbol, comment)
+          ctx.docCtx.get.addDocstring(tree.symbol, comment)
           tree.setComment(comment)
         }
       }
diff --git a/compiler/test/dotty/tools/dotc/core/tasty/CommentPicklingTest.scala b/compiler/test/dotty/tools/dotc/core/tasty/CommentPicklingTest.scala
index 461f6a610d8e..8b2a7be3daed 100644
--- a/compiler/test/dotty/tools/dotc/core/tasty/CommentPicklingTest.scala
+++ b/compiler/test/dotty/tools/dotc/core/tasty/CommentPicklingTest.scala
@@ -2,7 +2,7 @@ package dotty.tools.dotc.core.tasty
 
 import dotty.tools.dotc.ast.tpd
 import dotty.tools.dotc.ast.tpd.TreeOps
-import dotty.tools.dotc.Driver
+import dotty.tools.dotc.{Driver, Main}
 import dotty.tools.dotc.core.Comments.CommentsContext
 import dotty.tools.dotc.core.Contexts.Context
 import dotty.tools.dotc.core.Decorators.PreNamedString
@@ -78,16 +78,11 @@ class CommentPicklingTest {
     }
   }
 
-  private def findTreeNamed(name: Name)(trees: List[tpd.Tree], ctx: Context): Option[tpd.NameTree] = {
-    val acc = new tpd.TreeAccumulator[Option[tpd.NameTree]] {
-      override def apply(x: Option[tpd.NameTree], tree: tpd.Tree)(implicit ctx: Context): Option[tpd.NameTree] = {
-        x.orElse(tree match {
-          case md: tpd.NameTree if md.name == name => Some(md)
-          case other => foldOver(None, other)
-        })
-      }
-    }
-    acc(None, trees)(ctx)
+  private def findTreeNamed(name: Name)(trees: List[tpd.Tree], ctx: Context): Option[tpd.MemberDef] = {
+    implicit val _ctx: Context = ctx
+    trees.flatMap { _.find { case md: tpd.MemberDef => md.name == name; case _ => false }
+      .map(_.asInstanceOf[tpd.MemberDef]).toList
+    }.headOption
   }
 
   private def compileAndUnpickle[T](sources: List[String])(fn: (List[tpd.Tree], Context) => T) = {
@@ -103,9 +98,8 @@ class CommentPicklingTest {
       Files.createDirectories(out)
 
       val options = compileOptions.and("-d", out.toAbsolutePath.toString).and(sourceFiles: _*)
-      val driver = new Driver
       val reporter = TestReporter.reporter(System.out, logLevel = ERROR)
-      driver.process(options.all, reporter)
+      Main.process(options.all, reporter)
       assertFalse("Compilation failed.", reporter.hasErrors)
 
       val tastyFiles = getAll(tmp, "glob:**.tasty")

From 3f15ee05b2208b5af14bf10dcf81968f4744085c Mon Sep 17 00:00:00 2001
From: Martin Duhem <martin.duhem@gmail.com>
Date: Tue, 26 Jun 2018 10:14:41 +0200
Subject: [PATCH 9/9] Add `DocContext` in `DottyTest`

Now that we check that the `DocContext` is set, we need to set it in the
tests (or add `-Ydiscard-comments`, but this better reflects how Dotty
is going to be used).
---
 compiler/test/dotty/tools/DottyTest.scala | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/compiler/test/dotty/tools/DottyTest.scala b/compiler/test/dotty/tools/DottyTest.scala
index 6d2076d0e85c..5b356e169567 100644
--- a/compiler/test/dotty/tools/DottyTest.scala
+++ b/compiler/test/dotty/tools/DottyTest.scala
@@ -2,6 +2,7 @@ package dotty
 package tools
 
 import dotc.core._
+import dotc.core.Comments.{ContextDoc, ContextDocstrings}
 import dotc.core.Contexts._
 import dotc.core.Symbols._
 import dotc.core.Flags._
@@ -39,6 +40,7 @@ trait DottyTest extends ContextEscapeDetection {
   protected def initializeCtx(fc: FreshContext): Unit = {
     fc.setSetting(fc.settings.encoding, "UTF8")
     fc.setSetting(fc.settings.classpath, Jars.dottyLib)
+    fc.setProperty(ContextDoc, new ContextDocstrings)
   }
 
   private def compilerWithChecker(phase: String)(assertion: (tpd.Tree, Context) => Unit) = new Compiler {