Load compilation units from tasty

Author: nicolasstucki

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -47,6 +47,7 @@ class Compiler {
       List(new PostTyper),          // Additional checks and cleanups after type checking
       List(new sbt.ExtractAPI),     // Sends a representation of the API of classes to sbt via callbacks
       List(new Pickler),            // Generate TASTY info
+      List(new LinkAll),            // Link all
       List(new FirstTransform,      // Some transformations to put trees into a canonical form
            new CheckReentrant,      // Internal use only: Check that compiled program has no data races involving global vars
            new ElimJavaPackages),   // Eliminate syntactic references to Java packages

compiler/src/dotty/tools/dotc/FromTasty.scala

@@ -85,13 +85,8 @@ object FromTasty extends Driver {
               case info: ClassfileLoader =>
                 info.load(clsd)
                 val unpickled = clsd.symbol.asClass.tree
-                if (unpickled != null) {
-                    val unit1 = new CompilationUnit(new SourceFile(clsd.symbol.sourceFile, Seq()))
-                    unit1.tpdTree = unpickled
-                    force.traverse(unit1.tpdTree)
-                    unit1
-                } else
-                    cannotUnpickle(s"its class file ${info.classfile} does not have a TASTY attribute")
+                if (unpickled != null) mkUnit(clsd, unpickled)
+                else cannotUnpickle(s"its class file ${info.classfile} does not have a TASTY attribute")
               case info =>
                 cannotUnpickle(s"its info of type ${info.getClass} is not a ClassfileLoader")
             }
@@ -101,4 +96,20 @@ object FromTasty extends Driver {
         }
     }
   }
+
+  def loadCompilationUnit(clsd: ClassDenotation)(implicit ctx: Context): Option[CompilationUnit] = {
+    assert(ctx.settings.XlinkOptimise.value)
+    clsd.dottyUnpickler.flatMap { unpickler =>
+      ctx.log("Loading compilation unit for: " + clsd)
+      val body = unpickler.body(ctx.addMode(Mode.ReadPositions))
+      body.headOption.map(unpickled => mkUnit(clsd, unpickled))
+    }
+  }
+
+  private def mkUnit(clsd: ClassDenotation, unpickled: Tree)(implicit ctx: Context): CompilationUnit = {
+    val unit1 = new CompilationUnit(new SourceFile(clsd.symbol.sourceFile, Seq()))
+    unit1.tpdTree = unpickled
+    force.traverse(unit1.tpdTree)
+    unit1
+  }
 }

compiler/src/dotty/tools/dotc/config/ScalaSettings.scala

@@ -113,6 +113,7 @@ class ScalaSettings extends Settings.SettingGroup {
   val YoptPhases = PhasesSetting("-Yopt-phases", "Restrict the optimisation phases to execute under -optimise.")
   val YoptFuel = IntSetting("-Yopt-fuel", "Maximum number of optimisations performed under -optimise.", -1)
   val optimise = BooleanSetting("-optimise", "Generates faster bytecode by applying local optimisations to the .program") withAbbreviation "-optimize"
+  val XlinkOptimise = BooleanSetting("-Xlink-optimise", "Link class files.").withAbbreviation("-Xlink-optimize")
 
   /** Dottydoc specific settings */
   val siteRoot = StringSetting(

compiler/src/dotty/tools/dotc/core/SymDenotations.scala

@@ -19,6 +19,8 @@ import java.util.WeakHashMap
 import config.Config
 import config.Printers.{completions, incremental, noPrinter}
 
+import dotty.tools.dotc.core.tasty.DottyUnpickler
+
 trait SymDenotations { this: Context =>
   import SymDenotations._
 
@@ -139,6 +141,8 @@ object SymDenotations {
     private[this] var myPrivateWithin: Symbol = initPrivateWithin
     private[this] var myAnnotations: List[Annotation] = Nil
 
+    private[dotc] var dottyUnpickler: Option[DottyUnpickler] = None
+
     /** The owner of the symbol; overridden in NoDenotation */
     def owner: Symbol = ownerIfExists
 

compiler/src/dotty/tools/dotc/core/classfile/ClassfileParser.scala

@@ -15,6 +15,7 @@ import scala.annotation.switch
 import typer.Checking.checkNonCyclic
 import io.AbstractFile
 import scala.util.control.NonFatal
+import tasty.DottyUnpickler
 
 object ClassfileParser {
   /** Marker trait for unpicklers that can be embedded in classfiles. */
@@ -162,6 +163,16 @@ class ClassfileParser(
       }
     }
 
+    if (ctx.settings.XlinkOptimise.value) {
+      // Save references to DottyUnpickler to be able to load compilation units from tasty
+      result match {
+        case result@Some(_: DottyUnpickler) =>
+          classRoot.dottyUnpickler = result.asInstanceOf[Option[DottyUnpickler]]
+          moduleRoot.dottyUnpickler = result.asInstanceOf[Option[DottyUnpickler]]
+        case _ =>
+      }
+    }
+
     // eager load java enum definitions for exhaustivity check of pattern match
     if (isEnum) {
       instanceScope.toList.map(_.ensureCompleted())

compiler/src/dotty/tools/dotc/transform/LinkAll.scala

@@ -0,0 +1,63 @@
+package dotty.tools.dotc.transform
+
+import dotty.tools.dotc.{CompilationUnit, FromTasty}
+import dotty.tools.dotc.ast.Trees._
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Contexts._
+import dotty.tools.dotc.core.SymDenotations.ClassDenotation
+import dotty.tools.dotc.core.Symbols._
+import dotty.tools.dotc.core.Flags._
+import dotty.tools.dotc.transform.TreeTransforms._
+
+/** Loads all potentially reachable trees from tasty.    ▲
+ *  Only performed on whole world optimization mode.    ▲ ▲
+ *
+ *  TODO: Next step is to only load compilation units reachable in the call graph
+ */
+class LinkAll extends MiniPhaseTransform {
+  import tpd._
+
+  override def phaseName = "linkAll"
+
+  override def prepareForUnit(tree: tpd.Tree)(implicit ctx: Context): TreeTransform = NoTransform
+
+  override def runOn(units: List[CompilationUnit])(implicit ctx: Context): List[CompilationUnit] =
+    if (ctx.settings.XlinkOptimise.value) super.runOn(allUnits(Set.empty, units.toSet, Set.empty))
+    else super.runOn(units)
+
+  /** Loads and processes new compilation units, possibly loading more units. */
+  private def allUnits(processed: Set[CompilationUnit], unprocessed: Set[CompilationUnit], loadedClasses: Set[ClassDenotation])(implicit ctx: Context): List[CompilationUnit] = {
+    if (unprocessed.isEmpty) processed.toList
+    else {
+      val accum = new ClassesToLoadAccumulator
+      val classesToLoad = unprocessed.foldLeft(Set.empty[ClassDenotation])((acc, unit) => accum.apply(acc, unit.tpdTree)) -- loadedClasses
+      val loadedUnits = classesToLoad.flatMap(cls => FromTasty.loadCompilationUnit(cls))
+      allUnits(processed ++ unprocessed, loadedUnits, loadedClasses ++ classesToLoad)
+    }
+  }
+
+  /** Collects all class denotations that may need to be loaded. */
+  private class ClassesToLoadAccumulator extends TreeAccumulator[Set[ClassDenotation]] {
+    private var inParents = false
+    override def apply(acc: Set[ClassDenotation], tree: tpd.Tree)(implicit ctx: Context): Set[ClassDenotation] = tree match {
+      case New(tpt) => accum(acc, tpt.tpe.classSymbol)
+      case AppliedTypeTree(tpt, _) if inParents => accum(acc, tpt.symbol)
+      case tree: RefTree if inParents || tree.symbol.is(Module) =>
+        foldOver(accum(acc, tree.symbol), tree)
+      case tree @ Template(constr, parents, self, _) =>
+        val acc1 = this(acc, constr)
+        inParents = true
+        val acc2 = this(acc1, parents)
+        inParents = false
+        this(this(acc2, self), tree.body)
+      case _ => foldOver(acc, tree)
+    }
+
+    /** Accumulate class denotation for `sym` if needed */
+    private def accum(acc: Set[ClassDenotation], sym: Symbol)(implicit ctx: Context): Set[ClassDenotation] = {
+      val topClass = sym.topLevelClass.denot.asClass
+      if (topClass.is(JavaDefined) || topClass.is(Scala2x) || topClass.symbol == defn.ObjectClass) acc
+      else acc + topClass
+    }
+  }
+}