Merge pull request #2910 from dotty-staging/link-from-tasty

Link from tasty

Author: nicolasstucki

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

@@ -5,6 +5,9 @@ import dotty.tools.dotc.core.Types.Type
 import dotty.tools.dotc.core.tasty.{TastyUnpickler, TastyBuffer, TastyPickler}
 import util.SourceFile
 import ast.{tpd, untpd}
+import dotty.tools.dotc.ast.tpd.{ Tree, TreeTraverser }
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.core.SymDenotations.ClassDenotation
 import dotty.tools.dotc.core.Symbols._
 
 class CompilationUnit(val source: SourceFile) {
@@ -20,3 +23,20 @@ class CompilationUnit(val source: SourceFile) {
   /** Pickled TASTY binaries, indexed by class. */
   var pickled: Map[ClassSymbol, Array[Byte]] = Map()
 }
+
+object CompilationUnit {
+
+  /** Make a compilation unit for top class `clsd` with the contends of the `unpickled` */
+  def mkCompilationUnit(clsd: ClassDenotation, unpickled: Tree, forceTrees: Boolean)(implicit ctx: Context): CompilationUnit = {
+    val unit1 = new CompilationUnit(new SourceFile(clsd.symbol.sourceFile, Seq()))
+    unit1.tpdTree = unpickled
+    if (forceTrees)
+      force.traverse(unit1.tpdTree)
+    unit1
+  }
+
+  /** Force the tree to be loaded */
+  private object force extends TreeTraverser {
+    def traverse(tree: Tree)(implicit ctx: Context): Unit = traverseChildren(tree)
+  }
+}

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),            // Reload compilation units from TASTY for library code (if needed)
       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

@@ -60,10 +60,6 @@ object FromTasty extends Driver {
     override def toString = s"class file $className"
   }
 
-  object force extends TreeTraverser {
-    def traverse(tree: Tree)(implicit ctx: Context): Unit = traverseChildren(tree)
-  }
-
   class ReadTastyTreesFromClasses extends FrontEnd {
 
     override def isTyper = false
@@ -85,13 +81,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) CompilationUnit.mkCompilationUnit(clsd, unpickled, forceTrees = true)
+                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")
             }

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", "Recompile library code with the application.").withAbbreviation("-Xlink-optimize")
 
   /** Dottydoc specific settings */
   val siteRoot = StringSetting(

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

@@ -323,11 +323,15 @@ class ClassfileLoader(val classfile: AbstractFile) extends SymbolLoader {
 
   def load(root: SymDenotation)(implicit ctx: Context): Unit = {
     val (classRoot, moduleRoot) = rootDenots(root.asClass)
-    (new ClassfileParser(classfile, classRoot, moduleRoot)(ctx)).run() match {
-      case Some(unpickler: tasty.DottyUnpickler) if ctx.settings.YretainTrees.value =>
-        classRoot.symbol.asClass.unpickler = unpickler
-        moduleRoot.symbol.asClass.unpickler = unpickler
-      case _ =>
+    val classfileParser = new ClassfileParser(classfile, classRoot, moduleRoot)(ctx)
+    val result = classfileParser.run()
+    if (ctx.settings.YretainTrees.value || ctx.settings.XlinkOptimise.value) {
+      result match {
+        case Some(unpickler: tasty.DottyUnpickler) =>
+          classRoot.symbol.asClass.unpickler = unpickler
+          moduleRoot.symbol.asClass.unpickler = unpickler
+        case _ =>
+      }
     }
   }
 }

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

@@ -555,35 +555,31 @@ object Symbols {
 
     type ThisName = TypeName
 
-    /** If this is a top-level class, and if `-Yretain-trees` is set, return the TypeDef tree
-     *  for this class, otherwise EmptyTree. This will force the info of the class.
-     */
-    def tree(implicit ctx: Context): tpd.Tree /* tpd.TypeDef | tpd.EmptyTree */ = {
+    /** If this is either:
+      *   - a top-level class and `-Yretain-trees` is set
+     *    - a top-level class loaded from TASTY and `-Xlink-optimise` is set
+      * then return the TypeDef tree (possibly wrapped inside PackageDefs) for this class, otherwise EmptyTree.
+      * This will force the info of the class.
+      */
+    def tree(implicit ctx: Context): tpd.Tree /* tpd.PackageDef | tpd.TypeDef | tpd.EmptyTree */ = {
       denot.info
       // TODO: Consider storing this tree like we store lazy trees for inline functions
       if (unpickler != null && !denot.isAbsent) {
         assert(myTree.isEmpty)
-
-        import ast.Trees._
-
-        def findTree(tree: tpd.Tree): Option[tpd.TypeDef] = tree match {
-          case PackageDef(_, stats) =>
-            stats.flatMap(findTree).headOption
-          case tree: tpd.TypeDef if tree.symbol == this =>
-            Some(tree)
-          case _ =>
-              None
-        }
-        val List(unpickledTree) = unpickler.body(ctx.addMode(Mode.ReadPositions))
+        val body = unpickler.body(ctx.addMode(Mode.ReadPositions))
+        myTree = body.headOption.getOrElse(tpd.EmptyTree)
         unpickler = null
-
-        myTree = findTree(unpickledTree).get
       }
       myTree
     }
-    private[dotc] var myTree: tpd.Tree = tpd.EmptyTree
+    private[this] var myTree: tpd.Tree /* tpd.PackageDef | tpd.TypeDef | tpd.EmptyTree */ = tpd.EmptyTree
     private[dotc] var unpickler: tasty.DottyUnpickler = _
 
+    private[dotc] def registerTree(tree: tpd.TypeDef)(implicit ctx: Context): Unit = {
+      if (ctx.settings.YretainTrees.value)
+        myTree = tree
+    }
+
     /** The source or class file from which this class was generated, null if not applicable. */
     override def associatedFile(implicit ctx: Context): AbstractFile =
       if (assocFile != null || (this.owner is PackageClass) || this.isEffectiveRoot) assocFile

compiler/src/dotty/tools/dotc/interactive/SourceTree.scala

@@ -43,13 +43,15 @@ object SourceTree {
         sym.sourceFile == null) // FIXME: We cannot deal with external projects yet
       None
     else {
-      sym.tree match {
-        case tree: tpd.TypeDef =>
+      import ast.Trees._
+      def sourceTreeOfClass(tree: tpd.Tree): Option[SourceTree] = tree match {
+        case PackageDef(_, stats) => stats.flatMap(sourceTreeOfClass).headOption
+        case tree: tpd.TypeDef if tree.symbol == sym =>
           val sourceFile = new SourceFile(sym.sourceFile, Codec.UTF8)
           Some(SourceTree(tree, sourceFile))
-        case _ =>
-          None
+        case _ => None
       }
+      sourceTreeOfClass(sym.tree)
     }
   }
 }

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

@@ -161,6 +161,10 @@ class FirstTransform extends MiniPhaseTransform with InfoTransformer with Annota
     cpy.Template(impl)(self = EmptyValDef)
   }
 
+  /** Eliminate empty package definitions that may have been stored in the TASTY trees */
+  override def transformPackageDef(tree: PackageDef)(implicit ctx: Context, info: TransformerInfo): Tree =
+    if (tree.stats.isEmpty) EmptyTree else tree
+
   override def transformDefDef(ddef: DefDef)(implicit ctx: Context, info: TransformerInfo) = {
     if (ddef.symbol.hasAnnotation(defn.NativeAnnot)) {
       ddef.symbol.resetFlag(Deferred)

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

@@ -0,0 +1,80 @@
+package dotty.tools.dotc.transform
+
+import dotty.tools.dotc.CompilationUnit
+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._
+  import LinkAll._
+
+  override def phaseName = "linkAll"
+
+  /** Do not transform the any tree, runOn will traverse the trees and reload compilation units if needed */
+  override def prepareForUnit(tree: tpd.Tree)(implicit ctx: Context): TreeTransform = NoTransform
+
+  override def runOn(units: List[CompilationUnit])(implicit ctx: Context): List[CompilationUnit] = {
+    /** Loads and processes new compilation units, possibly loading more units. */
+    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 => loadCompilationUnit(cls))
+        allUnits(processed ++ unprocessed, loadedUnits, loadedClasses ++ classesToLoad)
+      }
+    }
+
+    if (ctx.settings.XlinkOptimise.value) super.runOn(allUnits(Set.empty, units.toSet, Set.empty))
+    else super.runOn(units)
+  }
+
+  /** 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
+    }
+  }
+}
+
+object LinkAll {
+
+  private[LinkAll] def loadCompilationUnit(clsd: ClassDenotation)(implicit ctx: Context): Option[CompilationUnit] = {
+    assert(ctx.settings.XlinkOptimise.value)
+    val tree = clsd.symbol.asClass.tree
+    if (tree.isEmpty) None
+    else {
+      ctx.log("Loading compilation unit for: " + clsd)
+      Some(CompilationUnit.mkCompilationUnit(clsd, tree, forceTrees = false))
+    }
+  }
+
+}

compiler/src/dotty/tools/dotc/typer/Typer.scala

@@ -1427,9 +1427,8 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
       // check value class constraints
       checkDerivedValueClass(cls, body1)
 
-      if (ctx.settings.YretainTrees.value) {
-        cls.myTree = cdef1
-      }
+      cls.registerTree(cdef1)
+
       cdef1
 
       // todo later: check that