From 533ee8aae91501c6a6a721fc79bc7629c16697a1 Mon Sep 17 00:00:00 2001
From: Fengyun Liu <fengyun.liu.cs@gmail.com>
Date: Thu, 29 Aug 2024 21:49:36 +0200
Subject: [PATCH] Widen values in assignment

---
 .../tools/dotc/transform/init/Objects.scala   | 45 +++++++++++--------
 1 file changed, 26 insertions(+), 19 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/transform/init/Objects.scala b/compiler/src/dotty/tools/dotc/transform/init/Objects.scala
index 1050fbe85ef2..1598e58ad767 100644
--- a/compiler/src/dotty/tools/dotc/transform/init/Objects.scala
+++ b/compiler/src/dotty/tools/dotc/transform/init/Objects.scala
@@ -1226,11 +1226,12 @@ class Objects(using Context @constructorOnly):
               extendTrace(id) { evalType(prefix, thisV, klass) }
 
         val value = eval(rhs, thisV, klass)
+        val widened = widenEscapedValue(value, rhs)
 
         if isLocal then
-          writeLocal(thisV, lhs.symbol, value)
+          writeLocal(thisV, lhs.symbol, widened)
         else
-          withTrace(trace2) { assign(receiver, lhs.symbol, value, rhs.tpe) }
+          withTrace(trace2) { assign(receiver, lhs.symbol, widened, rhs.tpe) }
 
       case closureDef(ddef) =>
         Fun(ddef, thisV, klass, summon[Env.Data])
@@ -1568,6 +1569,28 @@ class Objects(using Context @constructorOnly):
         throw new Exception("unexpected type: " + tp + ", Trace:\n" + Trace.show)
   }
 
+  /** Widen the escaped value (a method argument or rhs of an assignment)
+   *
+   *  The default widening is 1 for most values, 2 for function values.
+   *  User-specified widening annotations are repected.
+   */
+  def widenEscapedValue(value: Value, expr: Tree): Contextual[Value] =
+    expr.tpe.getAnnotation(defn.InitWidenAnnot) match
+      case Some(annot) =>
+        annot.argument(0).get match
+          case arg @ Literal(c: Constants.Constant) =>
+            val height = c.intValue
+            if height < 0 then
+              report.warning("The argument should be positive", arg)
+              value.widen(1)
+            else
+              value.widen(c.intValue)
+          case arg =>
+            report.warning("The argument should be a constant integer value", arg)
+            value.widen(1)
+      case _ =>
+        if value.isInstanceOf[Fun] then value.widen(2) else value.widen(1)
+
   /** Evaluate arguments of methods and constructors */
   def evalArgs(args: List[Arg], thisV: ThisValue, klass: ClassSymbol): Contextual[List[ArgInfo]] =
     val argInfos = new mutable.ArrayBuffer[ArgInfo]
@@ -1578,23 +1601,7 @@ class Objects(using Context @constructorOnly):
         else
           eval(arg.tree, thisV, klass)
 
-      val widened =
-        arg.tree.tpe.getAnnotation(defn.InitWidenAnnot) match
-        case Some(annot) =>
-          annot.argument(0).get match
-          case arg @ Literal(c: Constants.Constant) =>
-            val height = c.intValue
-            if height < 0 then
-              report.warning("The argument should be positive", arg)
-              res.widen(1)
-            else
-              res.widen(c.intValue)
-          case arg =>
-            report.warning("The argument should be a constant integer value", arg)
-            res.widen(1)
-        case _ =>
-          if res.isInstanceOf[Fun] then res.widen(2) else res.widen(1)
-
+      val widened = widenEscapedValue(res, arg.tree)
       argInfos += ArgInfo(widened, trace.add(arg.tree), arg.tree)
     }
     argInfos.toList