Extract MethodHandles.lookup() into a local scope in the static constructor

Author: jchyb

compiler/src/dotty/tools/dotc/transform/LazyVals.scala

@@ -28,7 +28,7 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
    * The map contains the list of the offset trees.
    */
   class OffsetInfo(var defs: List[Tree], var ord: Int = 0)
-  class VarHandleInfo(var defs: List[Tree])
+  class VarHandleInfo(var methodHandlesLookupDef: Tree, var varHandleDefs: List[Tree])
 
   private val appendOffsetDefs = mutable.Map.empty[Symbol, OffsetInfo]
   private val appendVarHandleDefs = mutable.Map.empty[Symbol, VarHandleInfo]
@@ -122,7 +122,7 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
       appendVarHandleDefs.get(cls) match {
         case None => template
         case Some(data) =>
-          cpy.Template(template)(body = addInFront(data.defs, template.body))
+          cpy.Template(template)(body = addInFront(data.methodHandlesLookupDef +: data.varHandleDefs, template.body))
       }
   }
 
@@ -475,22 +475,28 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
     containerSymbol.addAnnotation(Annotation(defn.VolatileAnnot, containerSymbol.span)) // private @volatile var _x: AnyRef
     containerSymbol.addAnnotations(x.symbol.annotations) // pass annotations from original definition
     containerSymbol.removeAnnotation(defn.ScalaStaticAnnot)
-    val getOffset =
-        Select(ref(defn.LazyValsModule), lazyNme.RLazyVals.getOffsetStatic)
     val containerTree = ValDef(containerSymbol, nullLiteral)
 
+    val varHandleInfo = appendVarHandleDefs.getOrElseUpdate(claz, new VarHandleInfo(EmptyTree, Nil))
+    varHandleInfo.methodHandlesLookupDef match
+      case EmptyTree =>
+        val lookupSym: TermSymbol = newSymbol(claz, (s"${claz.name}${lazyNme.methodHandleLookup}").toTermName, Synthetic, defn.MethodHandlesLookupClass.typeRef).enteredAfter(this)
+        lookupSym.addAnnotation(Annotation(defn.ScalaStaticAnnot, lookupSym.span))
+        varHandleInfo.methodHandlesLookupDef =
+          ValDef(lookupSym, Apply(Select(ref(defn.MethodHandlesClass), defn.MethodHandles_lookup.name), Nil))
+      case _ =>
+
     // create a VarHandle for this lazy val
     val varHandleSymbol: TermSymbol = newSymbol(claz, s"$containerName${lazyNme.lzyHandleSuffix}".toTermName, Synthetic, defn.VarHandleClass.typeRef).enteredAfter(this)
     varHandleSymbol.addAnnotation(Annotation(defn.ScalaStaticAnnot, varHandleSymbol.span))
     val getVarHandle = Apply(
-      Select(Apply(Select(ref(defn.MethodHandlesClass), defn.MethodHandles_lookup.name), Nil), defn.MethodHandlesLookup_FindVarHandle.name),
+      Select(ref(varHandleInfo.methodHandlesLookupDef.symbol), defn.MethodHandlesLookup_FindVarHandle.name),
       List(thizClass, Literal(Constant(containerName.toString)), Literal(Constant(defn.ObjectType)))
     )
     val varHandleTree = ValDef(varHandleSymbol, getVarHandle)
     val varHandle = ref(varHandleSymbol)
 
-    val varHandleInfo = appendVarHandleDefs.getOrElseUpdate(claz, new VarHandleInfo(Nil))
-    varHandleInfo.defs = varHandleTree :: varHandleInfo.defs
+    varHandleInfo.varHandleDefs = varHandleTree :: varHandleInfo.varHandleDefs
     val swapOver =
         This(claz)
 
@@ -696,8 +702,11 @@ object LazyVals {
     val lock: TermName        = "lock".toTermName
     val discard: TermName     = "discard".toTermName
     val lzyHandleSuffix: String = "$$lzyHandle"
+    val methodHandleLookup: String = "$$methodHandleLookup"
   }
 
   extension (sym: Symbol) def isVarHandleForLazy(using Context) =
     sym.name.endsWith(lazyNme.lzyHandleSuffix)
+  extension (sym: Symbol) def isMethodLookupForLazy(using Context) =
+    sym.name.endsWith(lazyNme.methodHandleLookup)
 }

compiler/src/dotty/tools/dotc/transform/MoveStatics.scala

@@ -12,15 +12,18 @@ import SymDenotations.SymDenotation
 import Names.Name
 import StdNames.nme
 import NameOps.*
-import LazyVals.isVarHandleForLazy
+import LazyVals.{isVarHandleForLazy, isMethodLookupForLazy}
 
 import ast.*
 
 import scala.collection.mutable
 
 import MegaPhase.*
 
-/** Move static methods from companion to the class itself */
+/** Move static methods from companion to the class itself. Also create the static constructor.
+ *  VarHandles generated by the compiler (more in LazyVals.scala) are left in the original class.
+ *  MethodSymbols.lookup() generated by the compiler is extracted to be a local static constructor value.
+ */
 class MoveStatics extends MiniPhase with SymTransformer {
   import ast.tpd.*
 
@@ -30,7 +33,8 @@ class MoveStatics extends MiniPhase with SymTransformer {
 
   def transformSym(sym: SymDenotation)(using Context): SymDenotation =
     if (sym.hasAnnotation(defn.ScalaStaticAnnot) && sym.owner.is(Flags.Module) && sym.owner.companionClass.exists &&
-        (sym.is(Flags.Method) || !(sym.isMutableVarOrAccessor && sym.owner.companionClass.is(Flags.Trait)) && !sym.symbol.isVarHandleForLazy)) {
+        (sym.is(Flags.Method) || !(sym.isMutableVarOrAccessor && sym.owner.companionClass.is(Flags.Trait)) &&
+        !sym.symbol.isVarHandleForLazy && !sym.symbol.isMethodLookupForLazy)) {
       sym.owner.asClass.delete(sym.symbol)
       sym.owner.companionClass.asClass.enter(sym.symbol)
       sym.copySymDenotation(owner = sym.owner.companionClass)
@@ -42,7 +46,11 @@ class MoveStatics extends MiniPhase with SymTransformer {
       val (classes, others) = trees.partition(x => x.isInstanceOf[TypeDef] && x.symbol.isClass)
       val pairs = classes.groupBy(_.symbol.name.stripModuleClassSuffix).asInstanceOf[Map[Name, List[TypeDef]]]
 
-      def rebuild(orig: TypeDef, newBody: List[Tree]): Tree = {
+      /** Rebuilds a template with static ValDefs from `newBody`` initialized in a static constructor.
+       *  `localStaticDefs`, used for initialisation of other static values,
+       *   are moved entirely into the static constructor.
+       */
+      def rebuild(orig: TypeDef, newBody: List[Tree], localStaticDefs: List[ValDef]): Tree = {
         val staticFields = newBody.filter(x => x.isInstanceOf[ValDef] && x.symbol.hasAnnotation(defn.ScalaStaticAnnot)).asInstanceOf[List[ValDef]]
         val newBodyWithStaticConstr =
           if (staticFields.nonEmpty) {
@@ -51,8 +59,24 @@ class MoveStatics extends MiniPhase with SymTransformer {
             staticCostructor.addAnnotation(Annotation(defn.ScalaStaticAnnot, staticCostructor.span))
             staticCostructor.entered
 
-            val staticAssigns = staticFields.map(x => Assign(ref(x.symbol), x.rhs.changeOwner(x.symbol, staticCostructor)))
-            tpd.DefDef(staticCostructor, Block(staticAssigns, tpd.unitLiteral)) :: newBody
+            val symbolRemap = localStaticDefs.map { x =>
+              x.symbol.owner.asClass.delete(x.symbol)
+              x.changeOwner(x.symbol.owner, staticCostructor)
+              (x.symbol, x.symbol.copy(owner = staticCostructor))
+            }.toMap
+
+            val localValDefs = localStaticDefs.map ( vDef =>
+              val newSymbol = symbolRemap(vDef.symbol)
+              ValDef(newSymbol.asTerm, vDef.rhs.changeOwner(vDef.symbol, newSymbol))
+            )
+
+            val localValDefMapper = new TreeMap:
+              override def transform(tree: Tree)(using Context): Tree = tree match
+                case ident @ Ident(tp) if symbolRemap.contains(ident.symbol) => ref(symbolRemap(ident.symbol))
+                case _ => super.transform(tree)
+
+            val staticAssigns = staticFields.map(x => Assign(ref(x.symbol), localValDefMapper.transform(x.rhs).changeOwner(x.symbol, staticCostructor)))
+            tpd.DefDef(staticCostructor, Block(localValDefs ++ staticAssigns, tpd.unitLiteral)) :: newBody
           }
           else newBody
 
@@ -69,9 +93,12 @@ class MoveStatics extends MiniPhase with SymTransformer {
 
           val staticDefs = mutable.ListBuffer[Tree]()
           val staticTiedDefs = mutable.ListBuffer[Tree]()
+          val localStaticDefs = mutable.ListBuffer[ValDef]()
           val remainingDefs = mutable.ListBuffer[Tree]()
 
           moduleTmpl.body.foreach {
+            case valDef: ValDef if valDef.symbol.isMethodLookupForLazy && valDef.symbol.isScalaStatic =>
+              localStaticDefs.addOne(valDef)
             case memberDef: MemberDef if memberDef.symbol.isScalaStatic =>
               if memberDef.symbol.isVarHandleForLazy then
                 staticTiedDefs.addOne(memberDef)
@@ -80,7 +107,7 @@ class MoveStatics extends MiniPhase with SymTransformer {
             case other => remainingDefs.addOne(other)
           }
 
-          rebuild(companion, companionTmpl.body ++ staticDefs) :: rebuild(module, staticTiedDefs.toList ++ remainingDefs.toList) :: Nil
+          rebuild(companion, companionTmpl.body ++ staticDefs, Nil) :: rebuild(module, staticTiedDefs.toList ++ remainingDefs.toList, localStaticDefs.toList) :: Nil
         }
       }
       val newPairs =
@@ -89,7 +116,11 @@ class MoveStatics extends MiniPhase with SymTransformer {
             if (classes.tail.isEmpty) {
               val classDef = classes.head
               val tmpl = classDef.rhs.asInstanceOf[Template]
-              rebuild(classDef, tmpl.body) :: Nil
+              val (localStaticDefs, remainingDefs) = tmpl.body.partition {
+                case valDef: ValDef => valDef.symbol.isMethodLookupForLazy && valDef.symbol.isScalaStatic
+                case _ => false
+              }
+              rebuild(classDef, remainingDefs, localStaticDefs.asInstanceOf[List[ValDef]]) :: Nil
             }
             else move(classes.head, classes.tail.head)
       Trees.flatten(newPairs.toList.flatten ++ others)