Re-design Env of the global object init checker (#24107)

When the global object init checker models lazy fields and by-name
parameters, initially it creates an environment for each symbol of the
lazy field/by-name parameter. For by-name parameters, the environment
includes all the possible environments that pass to the parameter as
outer environments. This PR modifies the design of environments to take
the tree to evaluate as the argument. Therefore, the same of argument
tree passed to a by-name parameter will share the same environment, and
there will be multiple environments for the same by-name parameter. This
facilitates the process of finding local variables in the argument tree.
Now the process of evaluating by-name args in `multiple-by-name.scala`
will behave similarly to evaluating closure arguments in
`multiple-func-arg.scala`

Author: olhotak

compiler/src/dotty/tools/dotc/transform/init/Objects.scala

@@ -93,11 +93,11 @@ class Objects(using Context @constructorOnly):
    *      | UnknownValue                                                 // values whose source are unknown at compile time
    * vs ::= ValueSet(Set(ve))                                            // set of abstract values
    * Value ::= ve | vs | Package
-   * Ref ::= ObjectRef | InstanceRef | ArrayRef                               // values that represent a reference to some (global or instance) object
+   * Ref ::= ObjectRef | InstanceRef | ArrayRef                          // values that represent a reference to some (global or instance) object
    * RefSet ::= Set(ref)                                                 // set of refs
    * Bottom ::= RefSet(Empty)                                            // unreachable code
    * ThisValue ::= Ref | RefSet                                          // possible values for 'this'
-   * EnvRef(meth, ownerObject)                                         // represents environments for methods or functions
+   * EnvRef(tree, ownerObject)                                           // represents environments for evaluating methods, functions, or lazy/by-name values
    * EnvSet ::= Set(EnvRef)
    * InstanceBody ::= (valsMap: Map[Symbol, Value],
                        outersMap: Map[ClassSymbol, Value],
@@ -405,13 +405,18 @@ class Objects(using Context @constructorOnly):
 
   /** Environment for parameters */
   object Env:
-    /** Local environments can be deeply nested, therefore we need `outer`.
-     *
-     *  For local variables in rhs of class field definitions, the `meth` is the primary constructor.
+    /** Represents environments for evaluating methods, functions, or lazy/by-name values
+     *  For methods or closures, `tree` is the DefDef of the method.
+     *  For lazy/by-name values, `tree` is the rhs of the definition or the argument passed to by-name param
      */
-    case class EnvRef(meth: Symbol, owner: ClassSymbol)(using Trace) extends Scope:
+    case class EnvRef(tree: Tree, owner: ClassSymbol)(using Trace) extends Scope:
+      override def equals(that: Any): Boolean =
+        that.isInstanceOf[EnvRef] &&
+        (that.asInstanceOf[EnvRef].tree eq tree) &&
+        (that.asInstanceOf[EnvRef].owner == owner)
+
       def show(using Context) =
-        "meth: " + meth.show + "\n" +
+        "tree: " + tree.show + "\n" +
         "owner: " + owner.show
 
       def valValue(sym: Symbol)(using EnvMap.EnvMapMutableData): Value = EnvMap.readVal(this, sym)
@@ -454,36 +459,49 @@ class Objects(using Context @constructorOnly):
 
     val NoEnv = EnvSet(Set.empty)
 
-    /** An empty environment can be used for non-method environments, e.g., field initializers.
-     *
-     *  The owner for the local environment for field initializers is the primary constructor of the
-     *  enclosing class.
-     */
-    def emptyEnv(meth: Symbol)(using Context, State.Data, EnvMap.EnvMapMutableData, Trace): EnvRef =
-      _of(Map.empty, meth, Bottom, NoEnv)
-
     def valValue(x: Symbol)(using env: EnvRef, ctx: Context, trace: Trace, envMap: EnvMap.EnvMapMutableData): Value =
       if env.hasVal(x) then
         env.valValue(x)
       else
         report.warning("[Internal error] Value not found " + x.show + "\nenv = " + env.show + ". " + Trace.show, Trace.position)
         Bottom
 
-    private[Env] def _of(argMap: Map[Symbol, Value], meth: Symbol, thisV: ThisValue, outerEnv: EnvSet)
+    /** The method of creating an Env that evaluates `tree` */
+    private[Env] def _of(argMap: Map[Symbol, Value], tree: Tree, thisV: ThisValue, outerEnv: EnvSet)
                         (using State.Data, EnvMap.EnvMapMutableData, Trace): EnvRef =
-      val env = EnvRef(meth, State.currentObject)
+      val env = EnvRef(tree, State.currentObject)
       argMap.foreach(env.initVal(_, _))
       env.initThisV(thisV)
       env.initOuterEnvs(outerEnv)
       env
 
+    /**
+     * Creates an environment that evaluates the body of a method or the body of a closure
+     */
+    def ofDefDef(ddef: DefDef, args: List[Value], thisV: ThisValue, outerEnv: EnvSet)
+                (using State.Data, EnvMap.EnvMapMutableData, Trace): EnvRef =
+      val params = ddef.termParamss.flatten.map(_.symbol)
+      assert(args.size == params.size, "arguments = " + args.size + ", params = " + params.size)
+      // assert(ddef.symbol.owner.is(Method) ^ (outerEnv == NoEnv), "ddef.owner = " + ddef.symbol.owner.show + ", outerEnv = " + outerEnv + ", " + ddef.source)
+      _of(params.zip(args).toMap, ddef, thisV, outerEnv)
+
+
+    /**
+     * Creates an environment that evaluates a lazy val with `tree` as rhs
+     * or evaluates a by-name parameter where `tree` is the argument tree
+     */
+    def ofByName(sym: Symbol, tree: Tree, thisV: ThisValue, outerEnv: EnvSet)
+                (using State.Data, EnvMap.EnvMapMutableData, Trace): EnvRef =
+      assert((sym.is(Flags.Param) && sym.info.isInstanceOf[ExprType]) || sym.is(Flags.Lazy));
+      _of(Map.empty, tree, thisV, outerEnv)
+
     /**
      * The main procedure for searching through the outer chain
      * @param target The symbol to search for if `bySymbol = true`; otherwise the method symbol of the target environment
      * @param scopeSet The set of scopes as starting point
      * @return The scopes that contains symbol `target` or whose method is `target`,
      *         and the value for `C.this` where C is the enclosing class of the result scopes
-    */
+     */
     private[Env] def resolveEnvRecur(
         target: Symbol, envSet: EnvSet, bySymbol: Boolean = true)
         : Contextual[Option[EnvSet]] = log("Resolving environment, target = " + target + ", envSet = " + envSet, printer) {
@@ -493,7 +511,7 @@ class Objects(using Context @constructorOnly):
           if bySymbol then
             envSet.envs.filter(_.hasVal(target))
           else
-            envSet.envs.filter(_.meth == target)
+            envSet.envs.filter(env => env.tree.isInstanceOf[DefDef] && env.tree.asInstanceOf[DefDef].symbol == target)
 
         assert(filter.isEmpty || filter.size == envSet.envs.size, "Either all scopes or no scopes contain " + target)
         if (!filter.isEmpty) then
@@ -515,19 +533,6 @@ class Objects(using Context @constructorOnly):
             resolveEnvRecur(target, outerEnvsOfThis, bySymbol)
     }
 
-
-    def ofDefDef(ddef: DefDef, args: List[Value], thisV: ThisValue, outerEnv: EnvSet)
-                (using State.Data, EnvMap.EnvMapMutableData, Trace): EnvRef =
-      val params = ddef.termParamss.flatten.map(_.symbol)
-      assert(args.size == params.size, "arguments = " + args.size + ", params = " + params.size)
-      // assert(ddef.symbol.owner.is(Method) ^ (outerEnv == NoEnv), "ddef.owner = " + ddef.symbol.owner.show + ", outerEnv = " + outerEnv + ", " + ddef.source)
-      _of(params.zip(args).toMap, ddef.symbol, thisV, outerEnv)
-
-    def ofByName(byNameParam: Symbol, thisV: ThisValue, outerEnv: EnvSet)
-                (using State.Data, EnvMap.EnvMapMutableData, Trace): EnvRef =
-      assert(byNameParam.is(Flags.Param) && byNameParam.info.isInstanceOf[ExprType]);
-      _of(Map.empty, byNameParam, thisV, outerEnv)
-
     def setLocalVal(x: Symbol, value: Value)(using scope: Scope, ctx: Context, heap: Heap.MutableData, envMap: EnvMap.EnvMapMutableData): Unit =
       assert(!x.isOneOf(Flags.Param | Flags.Mutable), "Only local immutable variable allowed")
       scope match
@@ -692,6 +697,11 @@ class Objects(using Context @constructorOnly):
     def setHeap(newHeap: Data)(using mutable: MutableData): Unit = mutable.heap = newHeap
   end Heap
 
+  /**
+     *  Local environments can be deeply nested, therefore we need `outerEnvs`, which stores the immediate outer environment.
+     *  If the immediate enclosing scope of an environment is a template, then `outerEnvs` is empty in EnvMap.
+     *  We can restore outerEnvs of `this` in the heap.
+     */
   object EnvMap:
     private case class EnvBody(
       valsMap: Map[Symbol, Value],
@@ -1191,12 +1201,12 @@ class Objects(using Context @constructorOnly):
 
     case ref: Ref =>
       val target = if needResolve then resolve(ref.klass, field) else field
-      if target.is(Flags.Lazy) then
-        given Scope = Env.emptyEnv(target.owner.asInstanceOf[ClassSymbol].primaryConstructor)
+      if target.is(Flags.Lazy) then // select a lazy field
         if ref.hasVal(target) then
           ref.valValue(target)
         else if target.hasSource then
           val rhs = target.defTree.asInstanceOf[ValDef].rhs
+          given Scope = Env.ofByName(target, rhs, ref, Env.NoEnv)
           val result = eval(rhs, ref, target.owner.asClass, cacheResult = true)
           ref.initVal(target, result)
           result
@@ -1358,8 +1368,8 @@ class Objects(using Context @constructorOnly):
     def evalByNameParam(value: Value): Contextual[Value] = value match
       case Fun(code, thisV, klass, scope) =>
         val byNameEnv = scope match {
-          case ref: Ref => Env.ofByName(sym, thisV, Env.NoEnv)
-          case env: Env.EnvRef => Env.ofByName(sym, thisV, Env.EnvSet(Set(env)))
+          case ref: Ref => Env.ofByName(sym, code, thisV, Env.NoEnv) // for by-name arguments of constructors
+          case env: Env.EnvRef => Env.ofByName(sym, code, thisV, Env.EnvSet(Set(env))) // for by-name arguments of methods/functions
         }
         given Scope = byNameEnv
         eval(code, thisV, klass, cacheResult = true)
@@ -1389,7 +1399,7 @@ class Objects(using Context @constructorOnly):
       else
         if sym.is(Flags.Lazy) then
           val outerThis = envSet.joinThisV
-          given Scope = Env.ofByName(sym, outerThis, envSet)
+          given Scope = Env.ofByName(sym, sym.defTree, outerThis, envSet)
           val rhs = sym.defTree.asInstanceOf[ValDef].rhs
           eval(rhs, outerThis, sym.enclosingClass.asClass, cacheResult = true)
         else
@@ -2144,7 +2154,7 @@ class Objects(using Context @constructorOnly):
       thisV
     else
       // `target` must enclose `klass`
-      assert(klass.enclosingClassNamed(target.name) != NoSymbol, target.show + " does not enclose " + klass.show)
+      assert(klass.ownersIterator.contains(target), target.show + " does not enclose " + klass.show)
       val outerThis = thisV match {
         case ref: Ref => ref.outerValue(klass)
         case refSet: RefSet => refSet.joinOuters(klass)

tests/init-global/pos/lazy-field.scala

@@ -0,0 +1,5 @@
+object O {
+  lazy val f1 = this
+  val f2 = 5
+  val f3 = f1.f2 + 5
+}

tests/init-global/pos/multiple-by-name.scala

@@ -10,6 +10,18 @@ object O {
     p.bar()
   }
 
+  def foo2(q: => X) = foo(q)
+  def foo3(r: => X) = foo(r)
+
   val a = foo(new X)
   val b = foo(new Y)
-}
+  val c = foo2(new Y)
+  val d = foo3(new Y)
+}
+
+/**
+  * Pass arg to by-name parameter: create a Fun where body is the argument expression
+  * Read value of by-name parameter: call 'apply' on every possible Fun value of the by-name parameter
+  * Solution: Add special EnvRefs for by-name params;
+  *           differentiate these EnvRefs by the arg tree passed to the by-name param
+  */

tests/init-global/pos/multiple-func-arg.scala

@@ -0,0 +1,20 @@
+class X {
+  def bar(): Int = 5
+}
+class Y extends X {
+  override def bar(): Int = 6
+}
+
+object O {
+  def foo(p: () => X) = {
+    p().bar() // flow: p <- [Fun(() => new X), Fun(() => new Y), Fun(() => q()), Fun(() => r())]
+  }
+
+  def foo2(q: () => X) = foo(() => q()) // flow: q <- [Fun(() => new Y)]
+  def foo3(r: () => X) = foo(() => r()) // flow: r <- [Fun(() => new Y)]
+
+  val a = foo(() => new X)
+  val b = foo(() => new Y)
+  val c = foo2(() => new Y)
+  val d = foo3(() => new Y)
+}

tests/init-global/warn/lazy-field.scala

@@ -0,0 +1,5 @@
+object O {
+  lazy val f1 = this
+  val f2 = 5
+  val f3: Int = f1.f2 + f1.f3 // warn
+}