Fix QuotesCache caching quoted symbol definitions with incorrect owners

Previously we would always cache and reuse unpickled trees, which was a
problem for quoted code which included symbol definitions. In those
cases, even if those quotes were being created within another context
(which should dictate owners of symbol definitions), it would be
ignored, and previous potentially incorrect symbol definitions would be
reused.

Now we include the quote symbol owner while caching, which is only taken
into account if the quoted code contains a symbol definition.

Author: jchyb

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

@@ -235,7 +235,21 @@ object PickledQuotes {
 
   /** Unpickle TASTY bytes into it's tree */
   private def unpickle(pickled: String | List[String], isType: Boolean)(using Context): Tree = {
-    QuotesCache.getTree(pickled) match
+    val unpicklingContext =
+      if ctx.owner.isClass then
+        // When a quote is unpickled with a Quotes context that that has a class `spliceOwner`
+        // we need to use a dummy owner to unpickle it. Otherwise any definitions defined
+        // in the quoted block would be accidentally entered in the class.
+        // When splicing this expression, this owner is replaced with the correct owner (see `quotedExprToTree` and `quotedTypeToTree` above).
+        // On the other hand, if the expression is used as a reflect term, the user must call `changeOwner` (same as with other expressions used within a nested owner).
+        // `-Xcheck-macros` will check for inconsistent owners and provide the users hints on how to improve them.
+        //
+        // Quotes context that that has a class `spliceOwner` can come from a macro annotation
+        // or a user setting it explicitly using `Symbol.asQuotes`.
+        ctx.withOwner(newSymbol(ctx.owner, "$quoteOwnedByClass$".toTermName, Private, defn.AnyType, NoSymbol))
+      else ctx
+
+    QuotesCache.getTree(pickled, unpicklingContext.owner) match
       case Some(tree) =>
         quotePickling.println(s"**** Using cached quote for TASTY\n$tree")
         treeOwner(tree) match
@@ -250,20 +264,6 @@ object PickledQuotes {
           case pickled: String => TastyString.unpickle(pickled)
           case pickled: List[String] => TastyString.unpickle(pickled)
 
-        val unpicklingContext =
-          if ctx.owner.isClass then
-            // When a quote is unpickled with a Quotes context that that has a class `spliceOwner`
-            // we need to use a dummy owner to unpickle it. Otherwise any definitions defined
-            // in the quoted block would be accidentally entered in the class.
-            // When splicing this expression, this owner is replaced with the correct owner (see `quotedExprToTree` and `quotedTypeToTree` above).
-            // On the other hand, if the expression is used as a reflect term, the user must call `changeOwner` (same as with other expressions used within a nested owner).
-            // `-Xcheck-macros` will check for inconsistent owners and provide the users hints on how to improve them.
-            //
-            // Quotes context that that has a class `spliceOwner` can come from a macro annotation
-            // or a user setting it explicitly using `Symbol.asQuotes`.
-            ctx.withOwner(newSymbol(ctx.owner, "$quoteOwnedByClass$".toTermName, Private, defn.AnyType, NoSymbol))
-            else ctx
-
         inContext(unpicklingContext) {
 
           quotePickling.println(s"**** unpickling quote from TASTY\n${TastyPrinter.showContents(bytes, ctx.settings.color.value == "never", isBestEffortTasty = false)}")
@@ -273,10 +273,26 @@ object PickledQuotes {
           unpickler.enter(Set.empty)
 
           val tree = unpickler.tree
-          QuotesCache(pickled) = tree
 
+          var includesSymbolDefinition = false
           // Make sure trees and positions are fully loaded
-          tree.foreachSubTree(identity)
+          new TreeTraverser {
+            def traverse(tree: Tree)(using Context): Unit =
+              tree match
+                case _: DefTree =>
+                  if !tree.symbol.hasAnnotation(defn.QuotedRuntime_SplicedTypeAnnot)
+                  && !tree.symbol.hasAnnotation(defn.QuotedRuntimePatterns_patternTypeAnnot)
+                  then
+                    includesSymbolDefinition = true
+                case _ =>
+              traverseChildren(tree)
+          }.traverse(tree)
+
+          // We cache with the context symbol owner only if we need to
+          val symbolOwnerMaybe = 
+            if (includesSymbolDefinition) Some(ctx.owner)
+            else None
+          QuotesCache(pickled, symbolOwnerMaybe) = tree
 
           quotePickling.println(i"**** unpickled quote\n$tree")
 

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

@@ -1,24 +1,43 @@
 package dotty.tools.dotc.quoted
 
 import dotty.tools.dotc.core.Contexts.*
+import dotty.tools.dotc.core.Symbols.Symbol
 import dotty.tools.dotc.util.Property
 import dotty.tools.dotc.ast.tpd
 
 
 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]]
-
+  /** Only used when the cached tree includes symbol definition. 
+   * Represents a mapping from the symbol owning the context of the quote to the unpickled tree. */
+  private type OwnerCache = collection.mutable.Map[Symbol, 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)
+  /** 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], Either[Tree, OwnerCache]]]
+
+
+  /** Get the cached tree of the quote. 
+   *  quoteOwner is taken into account only if the unpickled quote includes a symbol definition */
+  def getTree(pickled: String | List[String], quoteOwner: Symbol)(using Context): Option[Tree] =
+    ctx.property(QuotesCacheKey).get.get(pickled).flatMap {
+      case Left(tree: Tree) => Some(tree)
+      case Right(map) => map.get(quoteOwner)
+    }
+
+  /** Update the cached tree of the quote.
+   *  quoteOwner is applicable only if the quote includes a symbol definition, otherwise should be None */
+  def update(pickled: String | List[String], quoteOwner: Option[Symbol], tree: Tree)(using Context): Unit =
+    val previousValueMaybe = ctx.property(QuotesCacheKey).get.get(pickled)
+    val updatedValue: Either[Tree, OwnerCache] =
+      (previousValueMaybe, quoteOwner) match
+        case (None, Some(owner)) =>
+          Right(collection.mutable.Map((owner, tree)))
+        case (Some(map: OwnerCache), Some(owner)) =>
+          map.update(owner, tree)
+          Right(map)
+        case _ => Left(tree)
+    ctx.property(QuotesCacheKey).get.update(pickled, updatedValue)
 
   /** Context with a cache for quote trees and tasty bytes */
   def init(ctx: FreshContext): ctx.type =

tests/pos-macros/i20471/Macro_1.scala

@@ -0,0 +1,63 @@
+import scala.annotation.experimental
+import scala.quoted.*
+import scala.annotation.tailrec
+
+object FlatMap {
+  @experimental inline def derived[F[_]]: FlatMap[F] = MacroFlatMap.derive
+}
+trait FlatMap[F[_]]{
+  def tailRecM[A, B](a: A)(f: A => F[Either[A, B]]): F[B]
+}
+
+@experimental
+object MacroFlatMap:
+
+  inline def derive[F[_]]: FlatMap[F] = ${ flatMap }
+
+  def flatMap[F[_]: Type](using Quotes): Expr[FlatMap[F]] = '{
+    new FlatMap[F]:
+      def tailRecM[A, B](a: A)(f: A => F[Either[A, B]]): F[B] =
+        ${ deriveTailRecM('{ a }, '{ f }) }
+  }
+
+  def deriveTailRecM[F[_]: Type, A: Type, B: Type](
+      a: Expr[A],
+      f: Expr[A => F[Either[A, B]]]
+  )(using q: Quotes): Expr[F[B]] =
+    import quotes.reflect.*
+
+    val body: PartialFunction[(Symbol, TypeRepr), Term] = {
+        case (method, tpe) => {
+          given q2: Quotes = method.asQuotes
+          '{
+            def step(x: A): B = ???
+            ???
+          }.asTerm
+        }
+      }
+
+    val term = '{ $f($a) }.asTerm
+    val name = Symbol.freshName("$anon")
+    val parents = List(TypeTree.of[Object], TypeTree.of[F[B]])
+  
+    extension (sym: Symbol) def overridableMembers: List[Symbol] = 
+      val member1 = sym.methodMember("abstractEffect")(0)
+      val member2 = sym.methodMember("concreteEffect")(0)
+      def meth(member: Symbol) = Symbol.newMethod(sym, member.name, This(sym).tpe.memberType(member), Flags.Override, Symbol.noSymbol)
+      List(meth(member1), meth(member2))
+    
+    val cls = Symbol.newClass(Symbol.spliceOwner, name, parents.map(_.tpe), _.overridableMembers, None)
+
+    def transformDef(method: DefDef)(argss: List[List[Tree]]): Option[Term] =
+      val sym = method.symbol
+      Some(body.apply((sym, method.returnTpt.tpe)))
+
+    val members = cls.declarations
+      .filterNot(_.isClassConstructor)
+      .map: sym =>
+        sym.tree match
+          case method: DefDef => DefDef(sym, transformDef(method))
+          case _ => report.errorAndAbort(s"Not supported: $sym in ${sym.owner}")
+
+    val newCls = New(TypeIdent(cls)).select(cls.primaryConstructor).appliedToNone
+    Block(ClassDef(cls, parents, members) :: Nil, newCls).asExprOf[F[B]]

tests/pos-macros/i20471/Main_2.scala

@@ -0,0 +1,7 @@
+import scala.annotation.experimental
+
+@experimental
+object autoFlatMapTests:
+  trait TestAlgebra[T] derives FlatMap:
+    def abstractEffect(a: String): T
+    def concreteEffect(a: String): T = abstractEffect(a + " concreteEffect")