Optimize classpath implementation to speed up SBT

Adapted from scala commits b9520211b22e53d4f80801b0f64cdb6275ac55cf and
bd7341754f8aec4ddbf455544eb393ff6cf4a43a by Jason Zaugg and Lukas Rytz
in scala/scala#5956:

 - Rather than `findAll.filter(_.name == name)`, just lookup the entry
   with the right name to start with.
 - Avoid a linear scan of aggregate classpath by building and index
   keyed on
   package names

Author: allanrenucci

compiler/src/dotty/tools/dotc/classpath/AggregateClassPath.scala

@@ -19,32 +19,28 @@ import dotty.tools.io.{ AbstractFile, ClassPath, ClassRepresentation }
  */
 case class AggregateClassPath(aggregates: Seq[ClassPath]) extends ClassPath {
   override def findClassFile(className: String): Option[AbstractFile] = {
-    @tailrec
-    def find(aggregates: Seq[ClassPath]): Option[AbstractFile] =
-      if (aggregates.nonEmpty) {
-        val classFile = aggregates.head.findClassFile(className)
-        if (classFile.isDefined) classFile
-        else find(aggregates.tail)
-      } else None
-
-    find(aggregates)
+    val (pkg, _) = PackageNameUtils.separatePkgAndClassNames(className)
+    aggregatesForPackage(pkg).iterator.map(_.findClassFile(className)).collectFirst {
+      case Some(x) => x
+    }
+  }
+  private[this] val packageIndex: collection.mutable.Map[String, Seq[ClassPath]] = collection.mutable.Map()
+  private def aggregatesForPackage(pkg: String): Seq[ClassPath] = packageIndex.synchronized {
+    packageIndex.getOrElseUpdate(pkg, aggregates.filter(_.hasPackage(pkg)))
   }
 
   override def findClass(className: String): Option[ClassRepresentation] = {
-    @tailrec
-    def findEntry(aggregates: Seq[ClassPath], isSource: Boolean): Option[ClassRepresentation] =
-      if (aggregates.nonEmpty) {
-        val entry = aggregates.head.findClass(className) match {
-          case s @ Some(_: SourceFileEntry) if isSource => s
-          case s @ Some(_: ClassFileEntry) if !isSource => s
-          case _ => None
-        }
-        if (entry.isDefined) entry
-        else findEntry(aggregates.tail, isSource)
-      } else None
-
-    val classEntry = findEntry(aggregates, isSource = false)
-    val sourceEntry = findEntry(aggregates, isSource = true)
+    val (pkg, _) = PackageNameUtils.separatePkgAndClassNames(className)
+
+    def findEntry(isSource: Boolean): Option[ClassRepresentation] = {
+      aggregatesForPackage(pkg).iterator.map(_.findClass(className)).collectFirst {
+        case Some(s: SourceFileEntry) if isSource => s
+        case Some(s: ClassFileEntry) if !isSource => s
+      }
+    }
+
+    val classEntry = findEntry(isSource = false)
+    val sourceEntry = findEntry(isSource = true)
 
     (classEntry, sourceEntry) match {
       case (Some(c: ClassFileEntry), Some(s: SourceFileEntry)) => Some(ClassAndSourceFilesEntry(c.file, s.file))
@@ -70,6 +66,7 @@ case class AggregateClassPath(aggregates: Seq[ClassPath]) extends ClassPath {
   override private[dotty] def sources(inPackage: String): Seq[SourceFileEntry] =
     getDistinctEntries(_.sources(inPackage))
 
+  override private[dotty] def hasPackage(pkg: String) = aggregates.exists(_.hasPackage(pkg))
   override private[dotty] def list(inPackage: String): ClassPathEntries = {
     val (packages, classesAndSources) = aggregates.map { cp =>
       try {

compiler/src/dotty/tools/dotc/classpath/DirectoryClassPath.scala

@@ -45,6 +45,8 @@ trait DirectoryLookup[FileEntryType <: ClassRepresentation] extends ClassPath {
     }
   }
 
+  override private[dotty] def hasPackage(pkg: String) = getDirectory(pkg).isDefined
+
   private[dotty] def packages(inPackage: String): Seq[PackageEntry] = {
     val dirForPackage = getDirectory(inPackage)
     val nestedDirs: Array[F] = dirForPackage match {
@@ -161,6 +163,9 @@ final class JrtClassPath(fs: java.nio.file.FileSystem) extends ClassPath with No
     ps.map(p => (p.toString.stripPrefix("/packages/"), lookup(p))).toMap
   }
 
+  /** Empty string represents root package */
+  override private[dotty] def hasPackage(pkg: String) = packageToModuleBases.contains(pkg)
+
   override private[dotty] def packages(inPackage: String): Seq[PackageEntry] = {
     def matches(packageDottedName: String) =
       if (packageDottedName.contains("."))

compiler/src/dotty/tools/dotc/classpath/ZipAndJarFileLookupFactory.scala

@@ -6,8 +6,7 @@ package dotty.tools.dotc.classpath
 import java.io.File
 import java.net.URL
 import scala.annotation.tailrec
-import dotty.tools.io.{AbstractFile, ClassPath, FileZipArchive, ManifestResources}
-import dotty.tools.dotc.config.Settings
+import dotty.tools.io.{AbstractFile, ClassPath, ClassRepresentation, FileZipArchive, ManifestResources}
 import dotty.tools.dotc.core.Contexts.Context
 import FileUtils._
 
@@ -50,7 +49,13 @@ object ZipAndJarClassPathFactory extends ZipAndJarFileLookupFactory {
 
     override def findClassFile(className: String): Option[AbstractFile] = {
       val (pkg, simpleClassName) = PackageNameUtils.separatePkgAndClassNames(className)
-      classes(pkg).find(_.name == simpleClassName).map(_.file)
+      file(pkg, simpleClassName + ".class").map(_.file)
+    }
+
+    // This method is performance sensitive as it is used by SBT's ExtractDependencies phase.
+    override def findClass(className: String): Option[ClassRepresentation] = {
+      val (pkg, simpleClassName) = PackageNameUtils.separatePkgAndClassNames(className)
+      file(pkg, simpleClassName + ".class")
     }
 
     override private[dotty] def classes(inPackage: String): Seq[ClassFileEntry] = files(inPackage)
@@ -133,6 +138,7 @@ object ZipAndJarClassPathFactory extends ZipAndJarFileLookupFactory {
         (for (file <- pkg if file.isClass) yield ClassFileEntryImpl(file))(collection.breakOut)
     }
 
+    override private[dotty] def hasPackage(pkg: String) = cachedPackages.contains(pkg)
     override private[dotty] def list(inPackage: String): ClassPathEntries = ClassPathEntries(packages(inPackage), classes(inPackage))
   }
 

compiler/src/dotty/tools/dotc/classpath/ZipArchiveFileLookup.scala

@@ -39,6 +39,14 @@ trait ZipArchiveFileLookup[FileEntryType <: ClassRepresentation] extends ClassPa
       entry <- dirEntry.iterator if isRequiredFileType(entry)
     } yield createFileEntry(entry)
 
+  protected def file(inPackage: String, name: String): Option[FileEntryType] =
+    for {
+      dirEntry <- findDirEntry(inPackage)
+      entry <- Option(dirEntry.lookupName(name, directory = false))
+      if isRequiredFileType(entry)
+    } yield createFileEntry(entry)
+
+  override private[dotty] def hasPackage(pkg: String) = findDirEntry(pkg).isDefined
   override private[dotty] def list(inPackage: String): ClassPathEntries = {
     val foundDirEntry = findDirEntry(inPackage)
 

compiler/src/dotty/tools/io/ClassPath.scala

@@ -21,17 +21,42 @@ trait ClassPath {
   import dotty.tools.dotc.classpath._
   def asURLs: Seq[URL]
 
-  /** Empty string represents root package */
+  /*
+   * These methods are mostly used in the ClassPath implementation to implement the `list` and
+   * `findX` methods below.
+   *
+   * However, there are some other uses in the compiler, to implement `invalidateClassPathEntries`,
+   * which is used by the repl's `:require` (and maybe the spark repl, https://github.com/scala/scala/pull/4051).
+   * Using these methods directly is more efficient than calling `list`.
+   *
+   * The `inPackage` string is a full package name, e.g. "" or "scala.collection".
+   */
+
+  private[dotty] def hasPackage(pkg: String): Boolean
   private[dotty] def packages(inPackage: String): Seq[PackageEntry]
   private[dotty] def classes(inPackage: String): Seq[ClassFileEntry]
   private[dotty] def sources(inPackage: String): Seq[SourceFileEntry]
 
-  /** Allows to get entries for packages and classes merged with sources possibly in one pass. */
+  /**
+    * Returns packages and classes (source or classfile) that are members of `inPackage` (not
+    * recursively). The `inPackage` string is a full package name, e.g., "scala.collection".
+    *
+    * This is the main method uses to find classes, see class `PackageLoader`. The
+    * `rootMirror.rootLoader` is created with `inPackage = ""`.
+    */
   private[dotty] def list(inPackage: String): ClassPathEntries
 
   /**
-    * It returns both classes from class file and source files (as our base ClassRepresentation).
-    * So note that it's not so strictly related to findClassFile.
+    * Returns the class file and / or source file for a given external name, e.g., "java.lang.String".
+    * If there is both a class file and source file, the compiler can decide whether to read the
+    * class file or compile the source file.
+    *
+    * Internally this seems to be used only by `ScriptRunner`, but only to call `.isDefined`. That
+    * could probably be implemented differently.
+    *
+    * Externally, it is used by sbt's compiler interface:
+    * https://github.com/sbt/sbt/blob/v0.13.15/compile/interface/src/main/scala/xsbt/CompilerInterface.scala#L249
+    * Jason has some improvements for that in the works (https://github.com/scala/bug/issues/10289#issuecomment-310022699)
     */
   def findClass(className: String): Option[ClassRepresentation] = {
     // A default implementation which should be overridden, if we can create the more efficient
@@ -43,6 +68,16 @@ trait ClassPath {
 
     foundClassFromClassFiles orElse findClassInSources
   }
+
+ /**
+   * Returns the classfile for an external name, e.g., "java.lang.String". This method does not
+   * return source files.
+   *
+   * This method is used by the classfile parser. When parsing a Java class, its own inner classes
+   * are entered with a `ClassfileLoader` that parses the classfile returned by this method.
+   * It is also used in the backend, by the inliner, to obtain the bytecode when inlining from the
+   * classpath. It's also used by scalap.
+   */
   def findClassFile(className: String): Option[AbstractFile]
 
   def asClassPathStrings: Seq[String]

compiler/src/dotty/tools/io/ZipArchive.scala

@@ -24,7 +24,7 @@ import scala.annotation.tailrec
  *  ''Note:  This library is considered experimental and should not be used unless you know what you are doing.''
  */
 object ZipArchive {
-  private[io] val closeZipFile = sys.props.get("scala.classpath.closeZip").map(_.toBoolean).getOrElse(false)
+  private[io] val closeZipFile = sys.props.get("scala.classpath.closeZip").exists(_.toBoolean)
 
   /**
    * @param   file  a File
@@ -82,8 +82,8 @@ abstract class ZipArchive(override val file: JFile) extends AbstractFile with Eq
     override def isDirectory = true
     override def iterator: Iterator[Entry] = entries.valuesIterator
     override def lookupName(name: String, directory: Boolean): Entry = {
-      if (directory) entries(name + "/")
-      else entries(name)
+      if (directory) entries.get(name + "/").orNull
+      else entries.get(name).orNull
     }
   }