Contextualize owners in reflection API

* Introduce an `Owner` abstraction
* Replace `Tree.changeOwner` with contextualized `Tree.adaptOwner`
* Contextualize owner in `Lambda.apply`
* Contextualize owner in `Term.etaExpand`
* Contextualize owner in `Symbol.{newVal, newMethod, newBind}`
* Recontextualize `TreeMap`, `TreeAccumulator`, `TreeTaverser`

Author: nicolasstucki

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

@@ -199,7 +199,7 @@ object Matcher {
             def bodyFn(lambdaArgs: List[Tree]): Tree = {
               val argsMap = args.map(_.symbol).zip(lambdaArgs.asInstanceOf[List[Term]]).toMap
               new TreeMap {
-                override def transformTerm(tree: Term)(using ctx: Context): Term =
+                override def transformTerm(tree: Term)(using owner: Owner): Term =
                   tree match
                     case tree: Ident => summon[Env].get(tree.symbol).flatMap(argsMap.get).getOrElse(tree)
                     case tree => super.transformTerm(tree)
@@ -211,7 +211,7 @@ object Matcher {
             }
             val argTypes = args.map(x => x.tpe.widenTermRefExpr)
             val resType = pattern.tpe
-            val res = Lambda(Symbol.currentOwner, MethodType(names)(_ => argTypes, _ => resType), (meth, x) => bodyFn(x).changeOwner(meth))
+            val res = Lambda(MethodType(names)(_ => argTypes, _ => resType), x => bodyFn(x).adaptOwner)
             matched(res.asExpr)
 
           //
@@ -354,7 +354,7 @@ object Matcher {
       /** Return all free variables of the term defined in the pattern (i.e. defined in `Env`) */
       def freePatternVars(term: Term)(using ctx: Context, env: Env): Set[Symbol] =
         val accumulator = new TreeAccumulator[Set[Symbol]] {
-          def foldTree(x: Set[Symbol], tree: Tree)(using ctx: Context): Set[Symbol] =
+          def foldTree(x: Set[Symbol], tree: Tree)(using owner: Owner): Set[Symbol] =
             tree match
               case tree: Ident if env.contains(tree.symbol) => foldOverTree(x + tree.symbol, tree)
               case _ => foldOverTree(x, tree)

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

@@ -122,7 +122,7 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
       end extension
 
       extension [ThisTree <: Tree](self: ThisTree):
-        def changeOwner(newOwner: Symbol): ThisTree =
+        def adaptOwner(using newOwner: Symbol): ThisTree =
           tpd.TreeOps(self).changeNonLocalOwners(newOwner).asInstanceOf[ThisTree]
       end extension
 
@@ -365,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(owner: Symbol): Term = self.tpe.widen match {
+        def etaExpand(using owner: Owner): 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(owner, nme.ANON_FUN, Synthetic | Method, closureTpe)
-            tpd.Closure(closureMethod, tss => new tpd.TreeOps(self).appliedToArgs(tss.head).etaExpand(closureMethod))
+            tpd.Closure(closureMethod, tss => new tpd.TreeOps(self).appliedToArgs(tss.head).etaExpand(using closureMethod))
           case _ => self
         }
 
@@ -735,9 +735,9 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
     end ClosureMethodsImpl
 
     object Lambda extends LambdaModule:
-      def apply(owner: Symbol, tpe: MethodType, rhsFn: (Symbol, List[Tree]) => Tree): Block =
+      def apply(tpe: MethodType, rhsFn: Owner ?=> List[Tree] => Tree)(using owner: Owner): Block =
         val meth = dotc.core.Symbols.newSymbol(owner, nme.ANON_FUN, Synthetic | Method, tpe)
-        tpd.Closure(meth, tss => yCheckedOwners(rhsFn(meth, tss.head), meth))
+        tpd.Closure(meth, tss => yCheckedOwners(rhsFn(using 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, _))
@@ -2200,22 +2200,33 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
           case _ => None
     end AmbiguousImplicitsTypeTest
 
+    type Owner = dotc.core.Symbols.Symbol
+
+    object Owner extends OwnerModule:
+      def apply(sym: Symbol): Owner = sym
+      def spliceOwner: Owner = ctx.owner
+    end Owner
+
+    object OwnerMethodsImpl extends OwnerMethods:
+      extension (self: Owner):
+        def symbol: Symbol = self
+    end OwnerMethodsImpl
+
     type Symbol = dotc.core.Symbols.Symbol
 
     object Symbol extends SymbolModule:
-      def currentOwner(using ctx: Context): Symbol = ctx.owner
       def requiredPackage(path: String): Symbol = dotc.core.Symbols.requiredPackage(path)
       def requiredClass(path: String): Symbol = dotc.core.Symbols.requiredClass(path)
       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(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 =
+      def newMethod(name: String, tpe: TypeRepr)(using owner: Owner): Symbol =
+        newMethod(name, tpe, Flags.EmptyFlags, noSymbol)
+      def newMethod(name: String, tpe: TypeRepr, flags: Flags, privateWithin: Symbol)(using owner: Owner): 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 =
+      def newVal(name: String, tpe: TypeRepr, flags: Flags, privateWithin: Symbol)(using owner: Owner): Symbol =
         dotc.core.Symbols.newSymbol(owner, name.toTermName, flags, tpe, privateWithin)
-      def newBind(owner: Symbol, name: String, flags: Flags, tpe: TypeRepr): Symbol =
+      def newBind(name: String, flags: Flags, tpe: TypeRepr)(using owner: Owner): Symbol =
         dotc.core.Symbols.newSymbol(owner, name.toTermName, flags | Case, tpe)
       def noSymbol: Symbol = dotc.core.Symbols.NoSymbol
     end Symbol
@@ -2224,6 +2235,9 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
       extension (self: Symbol):
         def owner: Symbol = self.denot.owner
         def maybeOwner: Symbol = self.denot.maybeOwner
+        def asOwner: Owner =
+          assert(self.exists, "Cannot convert NoSymbol into an Owner")
+          self
         def flags: Flags = self.denot.flags
 
         def privateWithin: Option[TypeRepr] =
@@ -2593,6 +2607,7 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, QuoteUnpickl
                    |which has the AST representation
                    |${TreeMethods.showExtractors(tree)}
                    |
+                   |Owners can be adapted using `Tree.adaptOwner`.
                    |""".stripMargin)
             case _ => traverseChildren(t)
       }.traverse(tree)

library/src-bootstrapped/scala/quoted/ExprMap.scala

@@ -10,8 +10,7 @@ trait ExprMap:
     import qctx.reflect._
     final class MapChildren() {
 
-      def transformStatement(tree: Statement)(using ctx: Context): Statement = {
-        def localCtx(definition: Definition): Context = definition.symbol.localContext
+      def transformStatement(tree: Statement)(using owner: Owner): Statement = {
         tree match {
           case tree: Term =>
             transformTerm(tree, TypeRepr.of[Any])
@@ -22,15 +21,14 @@ trait ExprMap:
         }
       }
 
-      def transformDefinition(tree: Definition)(using ctx: Context): Definition = {
-        def localCtx(definition: Definition): Context = definition.symbol.localContext
+      def transformDefinition(tree: Definition)(using owner: Owner): Definition = {
         tree match {
           case tree: ValDef =>
-            given Context = localCtx(tree)
+            given Owner = Owner(tree.symbol)
             val rhs1 = tree.rhs.map(x => transformTerm(x, tree.tpt.tpe))
             ValDef.copy(tree)(tree.name, tree.tpt, rhs1)
           case tree: DefDef =>
-            given Context = localCtx(tree)
+            given Owner = Owner(tree.symbol)
             DefDef.copy(tree)(tree.name, tree.typeParams, tree.paramss, tree.returnTpt, tree.rhs.map(x => transformTerm(x, tree.returnTpt.tpe)))
           case tree: TypeDef =>
             tree
@@ -40,7 +38,7 @@ trait ExprMap:
         }
       }
 
-      def transformTermChildren(tree: Term, tpe: TypeRepr)(using ctx: Context): Term = tree match {
+      def transformTermChildren(tree: Term, tpe: TypeRepr)(using owner: Owner): Term = tree match {
         case Ident(name) =>
           tree
         case Select(qualifier, name) =>
@@ -94,7 +92,7 @@ trait ExprMap:
           Inlined.copy(tree)(call, transformDefinitions(bindings), transformTerm(expansion, tpe)/*()call.symbol.localContext)*/)
       }
 
-      def transformTerm(tree: Term, tpe: TypeRepr)(using ctx: Context): Term =
+      def transformTerm(tree: Term, tpe: TypeRepr)(using owner: Owner): Term =
         tree match
           case _: Closure =>
             tree
@@ -109,39 +107,39 @@ trait ExprMap:
           case _ =>
             transformTermChildren(tree, tpe)
 
-      def transformTypeTree(tree: TypeTree)(using ctx: Context): TypeTree = tree
+      def transformTypeTree(tree: TypeTree)(using owner: Owner): TypeTree = tree
 
-      def transformCaseDef(tree: CaseDef, tpe: TypeRepr)(using ctx: Context): CaseDef =
+      def transformCaseDef(tree: CaseDef, tpe: TypeRepr)(using owner: Owner): CaseDef =
         CaseDef.copy(tree)(tree.pattern, tree.guard.map(x => transformTerm(x, TypeRepr.of[Boolean])), transformTerm(tree.rhs, tpe))
 
-      def transformTypeCaseDef(tree: TypeCaseDef)(using ctx: Context): TypeCaseDef = {
+      def transformTypeCaseDef(tree: TypeCaseDef)(using owner: Owner): TypeCaseDef = {
         TypeCaseDef.copy(tree)(transformTypeTree(tree.pattern), transformTypeTree(tree.rhs))
       }
 
-      def transformStats(trees: List[Statement])(using ctx: Context): List[Statement] =
+      def transformStats(trees: List[Statement])(using owner: Owner): List[Statement] =
         trees mapConserve (transformStatement(_))
 
-      def transformDefinitions(trees: List[Definition])(using ctx: Context): List[Definition] =
+      def transformDefinitions(trees: List[Definition])(using owner: Owner): List[Definition] =
         trees mapConserve (transformDefinition(_))
 
-      def transformTerms(trees: List[Term], tpes: List[TypeRepr])(using ctx: Context): List[Term] =
+      def transformTerms(trees: List[Term], tpes: List[TypeRepr])(using owner: Owner): List[Term] =
         var tpes2 = tpes // TODO use proper zipConserve
         trees mapConserve { x =>
           val tpe :: tail = tpes2
           tpes2 = tail
           transformTerm(x, tpe)
         }
 
-      def transformTerms(trees: List[Term], tpe: TypeRepr)(using ctx: Context): List[Term] =
+      def transformTerms(trees: List[Term], tpe: TypeRepr)(using owner: Owner): List[Term] =
         trees.mapConserve(x => transformTerm(x, tpe))
 
-      def transformTypeTrees(trees: List[TypeTree])(using ctx: Context): List[TypeTree] =
+      def transformTypeTrees(trees: List[TypeTree])(using owner: Owner): List[TypeTree] =
         trees mapConserve (transformTypeTree(_))
 
-      def transformCaseDefs(trees: List[CaseDef], tpe: TypeRepr)(using ctx: Context): List[CaseDef] =
+      def transformCaseDefs(trees: List[CaseDef], tpe: TypeRepr)(using owner: Owner): List[CaseDef] =
         trees mapConserve (x => transformCaseDef(x, tpe))
 
-      def transformTypeCaseDefs(trees: List[TypeCaseDef])(using ctx: Context): List[TypeCaseDef] =
+      def transformTypeCaseDefs(trees: List[TypeCaseDef])(using owner: Owner): List[TypeCaseDef] =
         trees mapConserve (transformTypeCaseDef(_))
 
     }