Improve handling of owners in reflection API

* Add `changeOwner`
* Add owner and symbol to `Lambda.apply`
* Add owner to etaExpand for the creation of the closure
* 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)))
 
@@ -357,15 +365,15 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
         def tpe: TypeRepr = self.tpe
         def underlyingArgument: Term = new tpd.TreeOps(self).underlyingArgument
         def underlying: Term = new tpd.TreeOps(self).underlying
-        def etaExpand: Term = self.tpe.widen match {
+        def etaExpand(owner: Symbol): Term = self.tpe.widen match {
           case mtpe: Types.MethodType if !mtpe.isParamDependent =>
             val closureResType = mtpe.resType match {
               case t: Types.MethodType => t.toFunctionType()
               case t => t
             }
             val closureTpe = Types.MethodType(mtpe.paramNames, mtpe.paramInfos, closureResType)
-            val closureMethod = dotc.core.Symbols.newSymbol(ctx.owner, nme.ANON_FUN, Synthetic | Method, closureTpe)
-            tpd.Closure(closureMethod, tss => new tpd.TreeOps(self).appliedToArgs(tss.head).etaExpand)
+            val closureMethod = dotc.core.Symbols.newSymbol(owner, nme.ANON_FUN, Synthetic | Method, closureTpe)
+            tpd.Closure(closureMethod, tss => new tpd.TreeOps(self).appliedToArgs(tss.head).etaExpand(closureMethod))
           case _ => self
         }
 
@@ -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 */
@@ -477,7 +483,7 @@ trait QuoteContext { self: internal.QuoteUnpickler & internal.QuoteMatching =>
         def underlying: Term
 
         /** Converts a partally applied term into a lambda expression */
-        def etaExpand: Term
+        def etaExpand(owner: Symbol): Term
 
         /** A unary apply node with given argument: `tree(arg)` */
         def appliedTo(arg: Term): Term
@@ -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")

tests/pos-macros/i10151/Test_2.scala

@@ -70,7 +70,7 @@ object UnsafeExpr {
   }
   private def paramsAndBody[R](using qctx: QuoteContext)(f: Expr[Any]): (List[qctx.reflect.ValDef], Expr[R]) = {
     import qctx.reflect._
-    val Block(List(DefDef("$anonfun", Nil, List(params), _, Some(body))), Closure(Ident("$anonfun"), None)) = Term.of(f).etaExpand
+    val Block(List(DefDef("$anonfun", Nil, List(params), _, Some(body))), Closure(Ident("$anonfun"), None)) = Term.of(f).etaExpand(Symbol.currentOwner)
     (params, body.asExpr.asInstanceOf[Expr[R]])
   }
 

tests/pos-macros/i10211/Macro_1.scala

@@ -82,7 +82,7 @@ object UnsafeExpr {
   }
   private def paramsAndBody[R](using qctx: QuoteContext)(f: Expr[Any]): (List[qctx.reflect.ValDef], Expr[R]) = {
     import qctx.reflect._
-    val Block(List(DefDef("$anonfun", Nil, List(params), _, Some(body))), Closure(Ident("$anonfun"), None)) = Term.of(f).etaExpand
+    val Block(List(DefDef("$anonfun", Nil, List(params), _, Some(body))), Closure(Ident("$anonfun"), None)) = Term.of(f).etaExpand(Symbol.currentOwner)
     (params, body.asExpr.asInstanceOf[Expr[R]])
   }
 

tests/pos-macros/i10211/Test_2.scala

@@ -34,7 +34,7 @@ object UnsafeExpr {
   }
   private def paramsAndBody[R](using qctx: QuoteContext)(f: Expr[Any]): (List[qctx.reflect.ValDef], Expr[R]) = {
     import qctx.reflect._
-    val Block(List(DefDef("$anonfun", Nil, List(params), _, Some(body))), Closure(Ident("$anonfun"), None)) = Term.of(f).etaExpand
+    val Block(List(DefDef("$anonfun", Nil, List(params), _, Some(body))), Closure(Ident("$anonfun"), None)) = Term.of(f).etaExpand(Symbol.currentOwner)
     (params, body.asExpr.asInstanceOf[Expr[R]])
   }
 

tests/pos-macros/i9894/Macro_1.scala

@@ -14,10 +14,10 @@ object Asserts {
         fn.tpe.widen match {
           case _: MethodType =>
             args.size match {
-              case 0 => Expr.betaReduce('{ ${fn.etaExpand.asExprOf[() => Int]}() })
-              case 1 => Expr.betaReduce('{ ${fn.etaExpand.asExprOf[Int => Int]}(0) })
-              case 2 => Expr.betaReduce('{ ${fn.etaExpand.asExprOf[(Int, Int) => Int]}(0, 0) })
-              case 3 => Expr.betaReduce('{ ${fn.etaExpand.asExprOf[(Int, Int, Int) => Int]}(0, 0, 0) })
+              case 0 => Expr.betaReduce('{ ${fn.etaExpand(Symbol.currentOwner).asExprOf[() => Int]}() })
+              case 1 => Expr.betaReduce('{ ${fn.etaExpand(Symbol.currentOwner).asExprOf[Int => Int]}(0) })
+              case 2 => Expr.betaReduce('{ ${fn.etaExpand(Symbol.currentOwner).asExprOf[(Int, Int) => Int]}(0, 0) })
+              case 3 => Expr.betaReduce('{ ${fn.etaExpand(Symbol.currentOwner).asExprOf[(Int, Int, Int) => Int]}(0, 0, 0) })
             }
         }
       case _ => x
@@ -35,10 +35,10 @@ object Asserts {
       case Apply(fn, args) =>
         val pre = rec(fn)
         args.size match {
-          case 0 => Term.of(Expr.betaReduce('{ ${pre.etaExpand.asExprOf[() => Any]}() }))
-          case 1 => Term.of(Expr.betaReduce('{ ${pre.etaExpand.asExprOf[Int => Any]}(0) }))
-          case 2 => Term.of(Expr.betaReduce('{ ${pre.etaExpand.asExprOf[(Int, Int) => Any]}(0, 0) }))
-          case 3 => Term.of(Expr.betaReduce('{ ${pre.etaExpand.asExprOf[(Int, Int, Int) => Any]}(0, 0, 0) }))
+          case 0 => Term.of(Expr.betaReduce('{ ${pre.etaExpand(Symbol.currentOwner).asExprOf[() => Any]}() }))
+          case 1 => Term.of(Expr.betaReduce('{ ${pre.etaExpand(Symbol.currentOwner).asExprOf[Int => Any]}(0) }))
+          case 2 => Term.of(Expr.betaReduce('{ ${pre.etaExpand(Symbol.currentOwner).asExprOf[(Int, Int) => Any]}(0, 0) }))
+          case 3 => Term.of(Expr.betaReduce('{ ${pre.etaExpand(Symbol.currentOwner).asExprOf[(Int, Int, Int) => Any]}(0, 0, 0) }))
         }
       case _ => term
     }