Improve handling of owners in reflection API

* Add `changeOwner`
* Add symbol to `Lambda.apply`
* Ycheck the owners durring macro expansion
* Fix #10151
* Fix #10211

Author: nicolasstucki

compiler/src/scala/quoted/internal/impl/Matcher.scala

@@ -211,7 +211,7 @@ object Matcher {
             }
             val argTypes = args.map(x => x.tpe.widenTermRefExpr)
             val resType = pattern.tpe
-            val res = Lambda(MethodType(names)(_ => argTypes, _ => resType), bodyFn)
+            val res = Lambda(Symbol.currentOwner, MethodType(names)(_ => argTypes, _ => resType), (meth, x) => bodyFn(x).changeOwner(meth))
             matched(res.asExpr)
 
           //

compiler/src/scala/quoted/internal/impl/QuoteContextImpl.scala

@@ -45,6 +45,9 @@ object QuoteContextImpl {
 
 class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickler, QuoteMatching:
 
+  private val yCheck: Boolean =
+    ctx.settings.Ycheck.value(using ctx).exists(x => x == "all" || x == "macros")
+
   extension [T](self: scala.quoted.Expr[T]):
     def show: String =
       reflect.TreeMethodsImpl.show(reflect.Term.of(self))
@@ -118,6 +121,11 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
           QuoteContextImpl.this.asExprOf[T](self.asExpr)(using tp)
       end extension
 
+      extension [ThisTree <: Tree](self: ThisTree):
+        def changeOwner(newOwner: Symbol): ThisTree =
+          tpd.TreeOps(self).changeNonLocalOwners(newOwner).asInstanceOf[ThisTree]
+      end extension
+
     end TreeMethodsImpl
 
     type PackageClause = tpd.PackageDef
@@ -238,9 +246,9 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
 
     object DefDef extends DefDefModule:
       def apply(symbol: Symbol, rhsFn: List[TypeRepr] => List[List[Term]] => Option[Term]): DefDef =
-        withDefaultPos(tpd.polyDefDef(symbol.asTerm, tparams => vparamss => rhsFn(tparams)(vparamss).getOrElse(tpd.EmptyTree)))
+        withDefaultPos(tpd.polyDefDef(symbol.asTerm, tparams => vparamss => yCheckedOwners(rhsFn(tparams)(vparamss), symbol).getOrElse(tpd.EmptyTree)))
       def copy(original: Tree)(name: String, typeParams: List[TypeDef], paramss: List[List[ValDef]], tpt: TypeTree, rhs: Option[Term]): DefDef =
-        tpd.cpy.DefDef(original)(name.toTermName, typeParams, paramss, tpt, rhs.getOrElse(tpd.EmptyTree))
+        tpd.cpy.DefDef(original)(name.toTermName, typeParams, paramss, tpt, yCheckedOwners(rhs, original.symbol).getOrElse(tpd.EmptyTree))
       def unapply(ddef: DefDef): Option[(String, List[TypeDef], List[List[ValDef]], TypeTree, Option[Term])] =
         Some((ddef.name.toString, ddef.typeParams, ddef.paramss, ddef.tpt, optional(ddef.rhs)))
     end DefDef
@@ -264,9 +272,9 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
 
     object ValDef extends ValDefModule:
       def apply(symbol: Symbol, rhs: Option[Term]): ValDef =
-        tpd.ValDef(symbol.asTerm, rhs.getOrElse(tpd.EmptyTree))
+        tpd.ValDef(symbol.asTerm, yCheckedOwners(rhs, symbol).getOrElse(tpd.EmptyTree))
       def copy(original: Tree)(name: String, tpt: TypeTree, rhs: Option[Term]): ValDef =
-        tpd.cpy.ValDef(original)(name.toTermName, tpt, rhs.getOrElse(tpd.EmptyTree))
+        tpd.cpy.ValDef(original)(name.toTermName, tpt, yCheckedOwners(rhs, original.symbol).getOrElse(tpd.EmptyTree))
       def unapply(vdef: ValDef): Option[(String, TypeTree, Option[Term])] =
         Some((vdef.name.toString, vdef.tpt, optional(vdef.rhs)))
 
@@ -727,9 +735,9 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
     end ClosureMethodsImpl
 
     object Lambda extends LambdaModule:
-      def apply(tpe: MethodType, rhsFn: List[Tree] => Tree): Block =
-        val meth = dotc.core.Symbols.newSymbol(ctx.owner, nme.ANON_FUN, Synthetic | Method, tpe)
-        tpd.Closure(meth, tss => changeOwnerOfTree(rhsFn(tss.head), meth))
+      def apply(owner: Symbol, tpe: MethodType, rhsFn: (Symbol, List[Tree]) => Tree): Block =
+        val meth = dotc.core.Symbols.newSymbol(owner, nme.ANON_FUN, Synthetic | Method, tpe)
+        tpd.Closure(meth, tss => yCheckedOwners(rhsFn(meth, tss.head), meth))
 
       def unapply(tree: Block): Option[(List[ValDef], Term)] = tree match {
         case Block((ddef @ DefDef(_, _, params :: Nil, _, Some(body))) :: Nil, Closure(meth, _))
@@ -2201,14 +2209,14 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
       def requiredModule(path: String): Symbol = dotc.core.Symbols.requiredModule(path)
       def requiredMethod(path: String): Symbol = dotc.core.Symbols.requiredMethod(path)
       def classSymbol(fullName: String): Symbol = dotc.core.Symbols.requiredClass(fullName)
-      def newMethod(parent: Symbol, name: String, tpe: TypeRepr): Symbol =
-        newMethod(parent, name, tpe, Flags.EmptyFlags, noSymbol)
-      def newMethod(parent: Symbol, name: String, tpe: TypeRepr, flags: Flags, privateWithin: Symbol): Symbol =
-        dotc.core.Symbols.newSymbol(parent, name.toTermName, flags | dotc.core.Flags.Method, tpe, privateWithin)
-      def newVal(parent: Symbol, name: String, tpe: TypeRepr, flags: Flags, privateWithin: Symbol): Symbol =
-        dotc.core.Symbols.newSymbol(parent, name.toTermName, flags, tpe, privateWithin)
-      def newBind(parent: Symbol, name: String, flags: Flags, tpe: TypeRepr): Symbol =
-        dotc.core.Symbols.newSymbol(parent, name.toTermName, flags | Case, tpe)
+      def newMethod(owner: Symbol, name: String, tpe: TypeRepr): Symbol =
+        newMethod(owner, name, tpe, Flags.EmptyFlags, noSymbol)
+      def newMethod(owner: Symbol, name: String, tpe: TypeRepr, flags: Flags, privateWithin: Symbol): Symbol =
+        dotc.core.Symbols.newSymbol(owner, name.toTermName, flags | dotc.core.Flags.Method, tpe, privateWithin)
+      def newVal(owner: Symbol, name: String, tpe: TypeRepr, flags: Flags, privateWithin: Symbol): Symbol =
+        dotc.core.Symbols.newSymbol(owner, name.toTermName, flags, tpe, privateWithin)
+      def newBind(owner: Symbol, name: String, flags: Flags, tpe: TypeRepr): Symbol =
+        dotc.core.Symbols.newSymbol(owner, name.toTermName, flags | Case, tpe)
       def noSymbol: Symbol = dotc.core.Symbols.NoSymbol
     end Symbol
 
@@ -2541,6 +2549,53 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
     private def withDefaultPos[T <: Tree](fn: Context ?=> T): T =
       fn(using ctx.withSource(Position.ofMacroExpansion.source)).withSpan(Position.ofMacroExpansion.span)
 
+    /** Checks that all definitions in this tree have the expected owner.
+     *  Nested definitions are ignored and assumed to be correct by construction.
+     */
+    private def yCheckedOwners(tree: Option[Tree], owner: Symbol): tree.type =
+      if yCheck then
+        tree match
+          case Some(tree) =>
+            yCheckOwners(tree, owner)
+          case _ =>
+      tree
+
+    /** Checks that all definitions in this tree have the expected owner.
+     *  Nested definitions are ignored and assumed to be correct by construction.
+     */
+    private def yCheckedOwners(tree: Tree, owner: Symbol): tree.type =
+      if yCheck then
+        yCheckOwners(tree, owner)
+      tree
+
+    /** Checks that all definitions in this tree have the expected owner.
+     *  Nested definitions are ignored and assumed to be correct by construction.
+     */
+    private def yCheckOwners(tree: Tree, owner: Symbol): Unit =
+      new tpd.TreeTraverser {
+        def traverse(t: Tree)(using Context): Unit =
+          t match
+            case t: tpd.DefTree =>
+              val defOwner = t.symbol.owner
+              assert(defOwner == owner,
+                s"""Tree had an unexpected owner for ${t.symbol}
+                   |Expected: $owner (${owner.fullName})
+                   |But was: $defOwner (${defOwner.fullName})
+                   |
+                   |
+                   |The code of the definition of ${t.symbol} is
+                   |${TreeMethods.show(t)}
+                   |
+                   |which was found in the code
+                   |${TreeMethods.show(tree)}
+                   |
+                   |which has the AST representation
+                   |${TreeMethods.showExtractors(tree)}
+                   |
+                   |""".stripMargin)
+            case _ => traverseChildren(t)
+      }.traverse(tree)
+
   end reflect
 
   def unpickleExpr[T](pickled: String | List[String], typeHole: (Int, Seq[Any]) => scala.quoted.Type[?], termHole: (Int, Seq[Any], scala.quoted.QuoteContext) => scala.quoted.Expr[?]): scala.quoted.Expr[T] =

library/src/scala/quoted/QuoteContext.scala

@@ -182,7 +182,7 @@ trait QuoteContext { self: internal.QuoteUnpickler & internal.QuoteMatching =>
     // CONTEXTS //
     //////////////
 
-    /** Compilation context */
+    /** Context containing information on the current owner */
     type Context <: AnyRef
 
     /** Context of the macro expansion */
@@ -238,6 +238,12 @@ trait QuoteContext { self: internal.QuoteUnpickler & internal.QuoteMatching =>
       /** Convert this tree to an `quoted.Expr[T]` if the tree is a valid expression or throws */
       extension [T](self: Tree)
         def asExprOf(using scala.quoted.Type[T]): scala.quoted.Expr[T]
+
+      extension [ThisTree <: Tree](self: ThisTree):
+        /** Changes the owner of the symbols in the tree */
+        def changeOwner(newOwner: Symbol): ThisTree
+      end extension
+
     }
 
     /** Tree representing a pacakage clause in the source code */
@@ -962,8 +968,24 @@ trait QuoteContext { self: internal.QuoteUnpickler & internal.QuoteMatching =>
     val Lambda: LambdaModule
 
     trait LambdaModule { this: Lambda.type =>
+      /** Matches a lambda definition of the form
+       *  ```
+       *  Block((DefDef(_, _, params :: Nil, _, Some(body))) :: Nil, Closure(meth, _))
+       *  ```
+       *  Extracts the parameter definitions and body.
+       *
+       */
       def unapply(tree: Block): Option[(List[ValDef], Term)]
-      def apply(tpe: MethodType, rhsFn: List[Tree] => Tree): Block
+
+      /** Generates a lambda with the given method type.
+       *  ```
+       *  Block((DefDef(_, _, params :: Nil, _, Some(rhsFn(meth, paramRefs)))) :: Nil, Closure(meth, _))
+       *  ```
+       * @param owner: owner of the generated `meth` symbol
+       * @param tpe: Type of the definition
+       * @param rhsFn: Funtion that recieves the `meth` symbol and the a list of references to the `params`
+       */
+      def apply(owner: Symbol, tpe: MethodType, rhsFn: (Symbol, List[Tree]) => Tree): Block
     }
 
     given TypeTest[Tree, If] = IfTypeTest

tests/pos-macros/i10151/Macro_1.scala

@@ -46,7 +46,7 @@ object X:
             val paramTypes = params.map(_.tpt.tpe)
             val paramNames = params.map(_.name)
             val mt = MethodType(paramNames)(_ => paramTypes, _ => TypeRepr.of[CB].appliedTo(body.tpe.widen) )
-            val r = Lambda(mt, args => changeArgs(params,args,transform(body)) )
+            val r = Lambda(Symbol.currentOwner, mt, (newMeth, args) => changeArgs(params,args,transform(body).changeOwner(newMeth)) )
             r
           case _ =>
             throw RuntimeException("lambda expected")