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 */
@@ -179,27 +191,41 @@ object PickledQuotes {
     pickled
   }
 
+  private val cache = collection.mutable.Map.empty[String | List[String], Tree]
+
   /** 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.get(pickled) match
+      case Some(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
 
-    // Make sure trees and positions are fully loaded
-    new TreeTraverser {
-      def traverse(tree: Tree)(using Context): Unit = traverseChildren(tree)
-    }.traverse(tree)
+        var holeWithCapture = false // FIXME: remove this. TypeTreeMap is probably not changing the references in the Hole
 
-    quotePickling.println(i"**** unpickled quote\n$tree")
-    tree
+        // Make sure trees and positions are fully loaded
+        new TreeTraverser {
+          def traverse(tree: Tree)(using Context): Unit =
+            tree match
+            case Hole(_, _, _ :: _) =>
+              holeWithCapture = true
+            case _ =>
+              traverseChildren(tree)
+        }.traverse(tree)
+
+        quotePickling.println(i"**** unpickled quote\n$tree")
+        if !holeWithCapture then
+          QuotesCache(pickled) = tree
+        tree
   }
 
 }

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

@@ -0,0 +1,26 @@
+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 class loader used for macro expansion */
+  private val QuotesCacheKey = new Property.Key[collection.mutable.Map[String | List[String], Tree]]
+
+  /** Get the macro class loader */
+  def get(pickled: String | List[String])(using Context): Option[Tree] =
+    ctx.property(QuotesCacheKey).get.get(pickled)
+
+  def update(pickled: String | List[String], tree: Tree)(using Context): Unit =
+    ctx.property(QuotesCacheKey).get.update(pickled, tree)
+
+  /** Context with a cache for  */
+  def init(ctx: FreshContext): ctx.type =
+    ctx.setProperty(QuotesCacheKey, collection.mutable.Map.empty)
+}