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

@@ -54,12 +54,24 @@ object PickledQuotes {
 
   /** Unpickle the tree contained in the TastyExpr */
   def unpickleTerm(pickled: String | List[String], typeHole: (Int, Seq[Any]) => scala.quoted.Type[?], termHole: (Int, Seq[Any], scala.quoted.Quotes) => scala.quoted.Expr[?])(using Context): Tree = {
+    val t0 = System.nanoTime()
     val unpickled = withMode(Mode.ReadPositions)(unpickle(pickled, isType = false))
     val Inlined(call, Nil, expnasion) = unpickled
+    val t1 = System.nanoTime()
     val inlineCtx = inlineContext(call)
+    val t2 = System.nanoTime()
     val expansion1 = spliceTypes(expnasion, typeHole, termHole)(using inlineCtx)
+    val t3 = System.nanoTime()
     val expansion2 = spliceTerms(expansion1, typeHole, termHole)(using inlineCtx)
-    cpy.Inlined(unpickled)(call, Nil, expansion2)
+    val t4 = System.nanoTime()
+    val res = cpy.Inlined(unpickled)(call, Nil, expansion2)
+    val t5 = System.nanoTime()
+    // println(s"""
+    //         |unpickle> ${   (t1 - t0).toDouble / 1000 }μs
+    //         |spliceTypes> ${(t3 - t2).toDouble / 1000 }μs
+    //         |spliceTerms> ${(t4 - t3).toDouble / 1000 }μs
+    //         |""".stripMargin)
+    res
   }
 
   /** Unpickle the tree contained in the TastyType */
@@ -181,25 +193,39 @@ 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")
+        tree
+      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)}")
 
-    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
+        var cacheable = true // TODO: can we remove this?
 
-    // 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 =
+            tree match
+              case _: DefTree => cacheable = false
+              case _ =>
+            traverseChildren(tree)
+        }.traverse(tree)
 
-    quotePickling.println(i"**** unpickled quote\n$tree")
-    tree
+        quotePickling.println(i"**** unpickled quote\n$tree")
+        if cacheable then
+          QuotesCache(pickled) = 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) =>