Merge pull request #12242 from dotty-staging/add-cache-to-quote-unpickling

Cache quote unpickling

Author: nicolasstucki

compiler/src/dotty/tools/dotc/quoted/MacroExpansion.scala

@@ -14,6 +14,6 @@ object MacroExpansion {
     ctx.property(MacroExpansionPosition)
 
   def context(inlinedFrom: tpd.Tree)(using Context): Context =
-    ctx.fresh.setProperty(MacroExpansionPosition, SourcePosition(inlinedFrom.source, inlinedFrom.span)).setTypeAssigner(new Typer).withSource(inlinedFrom.source)
+    QuotesCache.init(ctx.fresh).setProperty(MacroExpansionPosition, SourcePosition(inlinedFrom.source, inlinedFrom.span)).setTypeAssigner(new Typer).withSource(inlinedFrom.source)
 }
 

compiler/src/dotty/tools/dotc/quoted/PickledQuotes.scala

@@ -181,25 +181,38 @@ object PickledQuotes {
 
   /** Unpickle TASTY bytes into it's tree */
   private def unpickle(pickled: String | List[String], isType: Boolean)(using Context): Tree = {
-    val bytes = pickled match
-      case pickled: String => TastyString.unpickle(pickled)
-      case pickled: List[String] => TastyString.unpickle(pickled)
+    QuotesCache.getTree(pickled) match
+      case Some(tree) =>
+        quotePickling.println(s"**** Using cached quote for TASTY\n$tree")
+        treeOwner(tree) match
+          case Some(owner) =>
+            // Copy the cached tree to make sure the all definitions are unique.
+            TreeTypeMap(oldOwners = List(owner), newOwners = List(owner)).apply(tree)
+          case _ =>
+            tree
 
-    quotePickling.println(s"**** unpickling quote from TASTY\n${TastyPrinter.show(bytes)}")
+      case _ =>
+        val bytes = pickled match
+          case pickled: String => TastyString.unpickle(pickled)
+          case pickled: List[String] => TastyString.unpickle(pickled)
+
+        quotePickling.println(s"**** unpickling quote from TASTY\n${TastyPrinter.show(bytes)}")
 
-    val mode = if (isType) UnpickleMode.TypeTree else UnpickleMode.Term
-    val unpickler = new DottyUnpickler(bytes, mode)
-    unpickler.enter(Set.empty)
+        val mode = if (isType) UnpickleMode.TypeTree else UnpickleMode.Term
+        val unpickler = new DottyUnpickler(bytes, mode)
+        unpickler.enter(Set.empty)
 
-    val tree = unpickler.tree
+        val tree = unpickler.tree
+        QuotesCache(pickled) = tree
 
-    // Make sure trees and positions are fully loaded
-    new TreeTraverser {
-      def traverse(tree: Tree)(using Context): Unit = traverseChildren(tree)
-    }.traverse(tree)
+        // Make sure trees and positions are fully loaded
+        new TreeTraverser {
+          def traverse(tree: Tree)(using Context): Unit = traverseChildren(tree)
+        }.traverse(tree)
 
-    quotePickling.println(i"**** unpickled quote\n$tree")
-    tree
+        quotePickling.println(i"**** unpickled quote\n$tree")
+
+        tree
   }
 
 }

compiler/src/dotty/tools/dotc/quoted/QuotesCache.scala

@@ -0,0 +1,28 @@
+package dotty.tools.dotc.quoted
+
+import dotty.tools.dotc.core.Contexts._
+import dotty.tools.dotc.util.Property
+import dotty.tools.dotc.reporting.trace
+import dotty.tools.dotc.ast.tpd
+
+import scala.collection.mutable
+
+object QuotesCache {
+  import tpd._
+
+  /** A key to be used in a context property that caches the unpickled trees */
+  private val QuotesCacheKey = new Property.Key[collection.mutable.Map[String | List[String], Tree]]
+
+
+  /** Get the cached tree of the quote */
+  def getTree(pickled: String | List[String])(using Context): Option[Tree] =
+    ctx.property(QuotesCacheKey).get.get(pickled)
+
+  /** Update the cached tree of the quote */
+  def update(pickled: String | List[String], tree: Tree)(using Context): Unit =
+    ctx.property(QuotesCacheKey).get.update(pickled, tree)
+
+  /** Context with a cache for quote trees and tasty bytes */
+  def init(ctx: FreshContext): ctx.type =
+    ctx.setProperty(QuotesCacheKey, collection.mutable.Map.empty)
+}

staging/src/scala/quoted/staging/QuoteDriver.scala

@@ -4,6 +4,7 @@ package staging
 import dotty.tools.dotc.ast.tpd
 import dotty.tools.dotc.Driver
 import dotty.tools.dotc.core.Contexts.{Context, ContextBase, FreshContext}
+import dotty.tools.dotc.quoted.QuotesCache
 import dotty.tools.io.{AbstractFile, Directory, PlainDirectory, VirtualDirectory}
 import dotty.tools.repl.AbstractFileClassLoader
 import dotty.tools.dotc.reporting._
@@ -33,8 +34,11 @@ private class QuoteDriver(appClassloader: ClassLoader) extends Driver:
           new VirtualDirectory("<quote compilation output>")
     end outDir
 
-    val ctx0 = setup(settings.compilerArgs.toArray :+ "dummy.scala", initCtx.fresh).get._2
-    val ctx = setCompilerSettings(ctx0.fresh.setSetting(ctx0.settings.outputDir, outDir), settings)
+    val ctx = {
+      val ctx0 = QuotesCache.init(initCtx.fresh)
+      val ctx1 = setup(settings.compilerArgs.toArray :+ "dummy.scala", ctx0).get._2
+      setCompilerSettings(ctx1.fresh.setSetting(ctx1.settings.outputDir, outDir), settings)
+    }
 
     new QuoteCompiler().newRun(ctx).compileExpr(exprBuilder) match
       case Right(value) =>