Re-design Env to take a tree as parameter

Author: EnzeXing

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
+}