Improve handling of owners in reflection API

* Add `changeOwner`
* Add `changeNonLocalOwners`
* Add symbol to `Lambda.unapply`
* Ycheck the owners durring macro expansion
* Fix scala#10151
* Fix scala#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(MethodType(names)(_ => argTypes, _ => resType), x => bodyFn(x).changeNonLocalOwners(Symbol.currentOwner))
             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,13 @@ 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(from: Symbol, to: Symbol): ThisTree =
+          tpd.TreeOps(self).changeOwner(from, to)
+        def changeNonLocalOwners(newOwner: Symbol): ThisTree =
+          tpd.TreeOps(self).changeNonLocalOwners(newOwner).asInstanceOf[ThisTree]
+      end extension
+
     end TreeMethodsImpl
 
     type PackageClause = tpd.PackageDef
@@ -238,9 +248,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 +274,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)))
 
@@ -729,12 +739,12 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
     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))
+        tpd.Closure(meth, tss => tpd.TreeOps(yCheckedOwners(rhsFn(tss.head), ctx.owner)).changeOwner(ctx.owner, meth))
 
-      def unapply(tree: Block): Option[(List[ValDef], Term)] = tree match {
+      def unapply(tree: Block): Option[(Symbol, List[ValDef], Term)] = tree match {
         case Block((ddef @ DefDef(_, _, params :: Nil, _, Some(body))) :: Nil, Closure(meth, _))
         if ddef.symbol == meth.symbol =>
-          Some((params, body))
+          Some((meth.symbol, params, body))
         case _ => None
       }
     end Lambda
@@ -2541,6 +2551,44 @@ 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)
 
+    private def yCheckedOwners(tree: Option[Tree], owner: Symbol): tree.type =
+      if yCheck then
+        tree match
+          case Some(tree) =>
+            yCheckOwners(tree, owner)
+          case _ =>
+      tree
+
+    private def yCheckedOwners(tree: Tree, owner: Symbol): tree.type =
+      if yCheck then
+        yCheckOwners(tree, owner)
+      tree
+
+    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) // TODO add a good error message
+            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] =

compiler/src/scala/quoted/internal/impl/printers/SourceCode.scala

@@ -456,7 +456,7 @@ object SourceCode {
         this += " = "
         printTree(rhs)
 
-      case tree @ Lambda(params, body) =>  // must come before `Block`
+      case tree @ Lambda(_, params, body) =>  // must come before `Block`
         inParens {
           printArgsDefs(params)
           this += (if tree.tpe.isContextFunctionType then " ?=> " else  " => ")
@@ -543,7 +543,7 @@ object SourceCode {
     private def flatBlock(stats: List[Statement], expr: Term): (List[Statement], Term) = {
       val flatStats = List.newBuilder[Statement]
       def extractFlatStats(stat: Statement): Unit = stat match {
-        case Lambda(_, _) =>   // must come before `Block`
+        case Lambda(_, _, _) =>   // must come before `Block`
           flatStats += stat
         case Block(stats1, expr1) =>
           val it = stats1.iterator
@@ -559,7 +559,7 @@ object SourceCode {
         case stat => flatStats += stat
       }
       def extractFlatExpr(term: Term): Term = term match {
-        case Lambda(_, _) =>   // must come before `Block`
+        case Lambda(_, _, _) =>   // must come before `Block`
           term
         case Block(stats1, expr1) =>
           val it = stats1.iterator
@@ -601,7 +601,7 @@ object SourceCode {
       def printSeparator(next: Tree): Unit = {
         // Avoid accidental application of opening `{` on next line with a double break
         def rec(next: Tree): Unit = next match {
-          case Lambda(_, _) => this += lineBreak()
+          case Lambda(_, _, _) => this += lineBreak()
           case Block(stats, _) if stats.nonEmpty => this += doubleLineBreak()
           case Inlined(_, bindings, _) if bindings.nonEmpty => this += doubleLineBreak()
           case Select(qual, _) => rec(qual)

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,14 @@ 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):
+        /** TODO */
+        def changeOwner(from: Symbol, to: Symbol): ThisTree
+        /** TODO */
+        def changeNonLocalOwners(newOwner: Symbol): ThisTree
+      end extension
+
     }
 
     /** Tree representing a pacakage clause in the source code */
@@ -962,7 +970,12 @@ trait QuoteContext { self: internal.QuoteUnpickler & internal.QuoteMatching =>
     val Lambda: LambdaModule
 
     trait LambdaModule { this: Lambda.type =>
-      def unapply(tree: Block): Option[(List[ValDef], Term)]
+      def unapply(tree: Block): Option[(Symbol, List[ValDef], Term)]
+
+      /* Generates a lambda with the given method type.
+       *
+       * Any definition in `rhs` is expected to be owned by the current owner in context.
+       */
       def apply(tpe: MethodType, rhsFn: List[Tree] => Tree): Block
     }
 

tests/neg-staging/i5941/macro_1.scala

@@ -25,7 +25,7 @@ object Lens {
         None, Nil,
         Block(
           DefDef(_, Nil, (param :: Nil) :: Nil, _, Some(Select(o, field))) :: Nil,
-          Lambda(meth, _)
+          Lambda(_, meth, _)
         )
       ) if o.symbol == param.symbol =>
         '{

tests/pos-macros/i10151/Macro_1.scala

@@ -42,11 +42,11 @@ object X:
 
    def shiftLambda(term:Term): Term =
         term match
-          case lt@Lambda(params, body) =>
+          case lt@Lambda(meth, params, body) =>
             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(mt, args => changeArgs(params,args,transform(body).changeOwner(meth, Symbol.currentOwner)) )
             r
           case _ =>
             throw RuntimeException("lambda expected")

tests/pos-macros/i10151/Test_2.scala

@@ -40,7 +40,7 @@ object Lens {
 
     object Function {
       def unapply(t: Term): Option[(List[ValDef], Term)] = t match {
-        case Inlined(None, Nil, Lambda(params, body)) => Some((params, body))
+        case Inlined(None, Nil, Lambda(_, params, body)) => Some((params, body))
         case _ => None
       }
     }

tests/pos-macros/i10211/Macro_1.scala

@@ -9,7 +9,7 @@ object lib {
     import util._
 
     Term.of(cond).underlyingArgument match {
-      case t @ Apply(Select(lhs, op), Lambda(param :: Nil, Apply(Select(a, "=="), b :: Nil)) :: Nil)
+      case t @ Apply(Select(lhs, op), Lambda(_, param :: Nil, Apply(Select(a, "=="), b :: Nil)) :: Nil)
       if a.symbol == param.symbol || b.symbol == param.symbol =>
         '{ scala.Predef.assert($cond) }
     }