Try special casing the resolveEnv function for locals

Author: q-ata

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

@@ -410,6 +410,7 @@ class Objects(using Context @constructorOnly):
       new LocalEnv(argMap, meth, outer)(valsMap = mutable.Map.empty, varsMap = mutable.Map.empty)
 
     def ofDefDef(ddef: DefDef, args: List[Value], outer: Data)(using Context): Data =
+      // println(ddef.show)
       val params = ddef.termParamss.flatten.map(_.symbol)
       assert(args.size == params.size, "arguments = " + args.size + ", params = " + params.size)
       assert(ddef.symbol.owner.isClass ^ (outer != NoEnv), "ddef.owner = " + ddef.symbol.owner.show + ", outer = " + outer + ", " + ddef.source)
@@ -452,17 +453,43 @@ class Objects(using Context @constructorOnly):
      *
      * @return the environment and value for `this` owned by the given method.
      */
-    def resolveEnv(meth: Symbol, thisV: ThisValue, env: Data)(using Context): Option[(ThisValue, Data)] = log("Resolving env for " + meth.show + ", this = " + thisV.show + ", env = " + env.show, printer) {
+    def _resolveEnvByValue(target: Symbol, thisV: ThisValue, env: Data)(using Context): Option[(ThisValue, Data)] = log("Resolving env for " + target.show + ", this = " + thisV.show + ", env = " + env.show, printer) {
+      env match
+      case localEnv: LocalEnv =>
+        if localEnv.getVal(target).isDefined then Some(thisV -> localEnv)
+        else if localEnv.getVar(target).isDefined then Some(thisV -> localEnv)
+        else _resolveEnvByValue(target, thisV, localEnv.outer)
+      case NoEnv =>
+        thisV match
+        case ref: OfClass =>
+          ref.outer match
+          case outer : ThisValue =>
+            _resolveEnvByValue(target, outer, ref.env)
+          case _ =>
+            // TODO: properly handle the case where ref.outer is ValueSet
+            None
+        case _ =>
+          None
+    }
+
+    def resolveEnvByValue(target: Symbol, thisV: ThisValue, env: Data)(using Context): Option[(ThisValue, Data)] =
+      // println("the target: " + target.show + "its owner: " + target.owner.show + " the env: " + env.show)
+      val e = _resolveEnvByValue(target, thisV, env)
+      if e.isDefined then println(e.get._2.show)
+      else println("NONE")
+      e
+
+    def resolveEnvByOwner(meth: Symbol, thisV: ThisValue, env: Data)(using Context): Option[(ThisValue, Data)] = log("Resolving env for " + meth.show + ", this = " + thisV.show + ", env = " + env.show, printer) {
       env match
       case localEnv: LocalEnv =>
         if localEnv.meth == meth then Some(thisV -> env)
-        else resolveEnv(meth, thisV, localEnv.outer)
+        else resolveEnvByOwner(meth, thisV, localEnv.outer)
       case NoEnv =>
         thisV match
         case ref: OfClass =>
           ref.outer match
           case outer : ThisValue =>
-            resolveEnv(meth, outer, ref.env)
+            resolveEnvByOwner(meth, outer, ref.env)
           case _ =>
             // TODO: properly handle the case where ref.outer is ValueSet
             None
@@ -725,7 +752,8 @@ class Objects(using Context @constructorOnly):
             if meth.owner.isClass then
               (ref, Env.NoEnv)
             else
-              Env.resolveEnv(meth.owner.enclosingMethod, ref, summon[Env.Data]).getOrElse(Cold -> Env.NoEnv)
+              // println("1: " + meth.show + " env : " + summon[Env.Data].show)
+              Env.resolveEnvByOwner(meth.owner.enclosingMethod, ref, summon[Env.Data]).getOrElse(Cold -> Env.NoEnv)
 
           val env2 = Env.ofDefDef(ddef, args.map(_.value), outerEnv)
           extendTrace(ddef) {
@@ -963,7 +991,8 @@ class Objects(using Context @constructorOnly):
                   (thisV.widenRefOrCold(1), Env.NoEnv)
               else
                 // klass.enclosingMethod returns its primary constructor
-                Env.resolveEnv(klass.owner.enclosingMethod, thisV, summon[Env.Data]).getOrElse(Cold -> Env.NoEnv)
+                // println("callsite 2: " + klass.show + " env : " + summon[Env.Data].show)
+                Env.resolveEnvByOwner(klass.owner.enclosingMethod, thisV, summon[Env.Data]).getOrElse(Cold -> Env.NoEnv)
 
         val instance = OfClass(klass, outerWidened, ctor, args.map(_.value), envWidened)
         callConstructor(instance, ctor, args)
@@ -983,7 +1012,6 @@ class Objects(using Context @constructorOnly):
       Env.setLocalVar(sym, addr)
       Heap.writeJoin(addr, value)
     else
-      System.out.println("initLocal sym: " + sym.show + " val: " + value.show + " data: " + data.show)
       Env.setLocalVal(sym, value)
   }
 
@@ -994,7 +1022,8 @@ class Objects(using Context @constructorOnly):
    */
   def readLocal(thisV: ThisValue, sym: Symbol): Contextual[Value] = log("reading local " + sym.show, printer, (_: Value).show) {
     def isByNameParam(sym: Symbol) = sym.is(Flags.Param) && sym.info.isInstanceOf[ExprType]
-    Env.resolveEnv(sym.enclosingMethod, thisV, summon[Env.Data]) match
+    // println("callsite 3: " + sym.show + " env : " + summon[Env.Data].show)
+    Env.resolveEnvByValue(sym, thisV, summon[Env.Data]) match
     case Some(thisV -> env) =>
       if sym.is(Flags.Mutable) then
         // Assume forward reference check is doing a good job
@@ -1017,15 +1046,12 @@ class Objects(using Context @constructorOnly):
           val rhs = sym.defTree.asInstanceOf[ValDef].rhs
           eval(rhs, thisV, sym.enclosingClass.asClass, cacheResult = true)
         else
-          System.out.println("env: " + env.show)
-          System.out.println("sym: " + sym.show)
           // Assume forward reference check is doing a good job
           val value = Env.valValue(sym)
           if isByNameParam(sym) then
             value match
             case fun: Fun =>
               given Env.Data = Env.ofByName(sym, fun.env)
-              System.out.println("created new env " + fun.code.show)
               eval(fun.code, fun.thisV, fun.klass)
             case Cold =>
               report.warning("Calling cold by-name alias. " + Trace.show, Trace.position)
@@ -1052,8 +1078,8 @@ class Objects(using Context @constructorOnly):
    */
   def writeLocal(thisV: ThisValue, sym: Symbol, value: Value): Contextual[Value] = log("write local " + sym.show + " with " + value.show, printer, (_: Value).show) {
     assert(sym.is(Flags.Mutable), "Writing to immutable variable " + sym.show)
-
-    Env.resolveEnv(sym.enclosingMethod, thisV, summon[Env.Data]) match
+    // println("callsite 4: " + sym.show + " env : " + summon[Env.Data].show)
+    Env.resolveEnvByValue(sym, thisV, summon[Env.Data]) match
     case Some(thisV -> env) =>
       given Env.Data = env
       Env.getVar(sym) match
@@ -1189,6 +1215,7 @@ class Objects(using Context @constructorOnly):
             withTrace(trace2) { call(receiver, ref.symbol, args, receiver = qual.tpe, superType = NoType) }
 
         case id: Ident =>
+          // println("ID : " + id.show)
           id.tpe match
           case TermRef(NoPrefix, _) =>
             // resolve this for the local method
@@ -1258,6 +1285,7 @@ class Objects(using Context @constructorOnly):
         Fun(ddef, thisV, klass, summon[Env.Data])
 
       case Block(stats, expr) =>
+        // println("BLOCK DATA : " + expr.show)
         evalExprs(stats, thisV, klass)
         eval(expr, thisV, klass)