Suppress returns that cross different stack sizes.

This is needed to avoid verify errors

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

@@ -28,7 +28,7 @@ object DropBreaks:
     var otherRefs: Int = 0
 
   private val LabelUsages = new Property.Key[Map[Symbol, LabelUsage]]
-  private val LabelsShadowedByTry = new Property.Key[Set[Symbol]]
+  private val ShadowedLabels = new Property.Key[Set[Symbol]]
 
 /** Rewrites local Break throws to labeled returns.
  *  Drops `try` statements on breaks if no other uses of its label remain.
@@ -38,7 +38,7 @@ object DropBreaks:
  *    - the throw and the boundary are in the same method, and
  *    - there is no try expression inside the boundary that encloses the throw.
  */
-class DropBreaks extends MiniPhase:
+class DropBreaks extends MiniPhase, RecordStackChange:
   import DropBreaks.*
 
   import tpd._
@@ -99,6 +99,31 @@ class DropBreaks extends MiniPhase:
         None
   end BreakBoundary
 
+  private object Break:
+
+    private def isBreak(sym: Symbol)(using Context): Boolean =
+      sym.name == nme.apply && sym.owner == defn.breakModule.moduleClass
+
+    /** `(local, arg)` provided `tree` matches
+     *
+     *     break[...](arg)(local)
+     *
+     *  or `(local, ())` provided `tree` matches
+     *
+     *     break()(local)
+     */
+    def unapply(tree: Tree)(using Context): Option[(Symbol, Tree)] = tree match
+      case Apply(Apply(fn, args), id :: Nil)
+      if isBreak(fn.symbol) =>
+        stripInlined(id) match
+          case id: Ident =>
+            val arg = (args: @unchecked) match
+              case arg :: Nil => arg
+              case Nil => Literal(Constant(())).withSpan(tree.span)
+            Some((id.symbol, arg))
+          case _ => None
+      case _ => None
+
   /** The LabelUsage data associated with `lbl` in the current context */
   private def labelUsage(lbl: Symbol)(using Context): Option[LabelUsage] =
     for
@@ -117,18 +142,34 @@ class DropBreaks extends MiniPhase:
     case _ =>
       ctx
 
-  /** If `tree` is not a LabeledTry, include all enclosing labels in the
-   *  `LabelsShadowedByTry` context property. This means that breaks to these
-   *  labels will not be translated to labeled returns in the body of the try.
+  /** Include all enclosing labels in the `ShadowedLabels` context property.
+   *  This means that breaks to these labels will not be translated to labeled
+   *  returns while this context is valid.
    */
-  override def prepareForTry(tree: Try)(using Context): Context = tree match
-    case LabelTry(_, _) => ctx
-    case _ => ctx.property(LabelUsages) match
+  private def shadowLabels(using Context): Context =
+    ctx.property(LabelUsages) match
       case Some(usesMap) =>
-        val setSoFar = ctx.property(LabelsShadowedByTry).getOrElse(Set.empty)
-        ctx.fresh.setProperty(LabelsShadowedByTry, setSoFar ++ usesMap.keysIterator)
+        val setSoFar = ctx.property(ShadowedLabels).getOrElse(Set.empty)
+        ctx.fresh.setProperty(ShadowedLabels, setSoFar ++ usesMap.keysIterator)
       case _ => ctx
 
+  /** Need to suppress labeled returns if the stack can change between
+   *  source and target of the jump
+   */
+  protected def stackChange(using Context) = shadowLabels
+
+  /** Need to suppress labeled returns if there is an intervening try
+   */
+  override def prepareForTry(tree: Try)(using Context): Context = tree match
+    case LabelTry(_, _) => ctx
+    case _ => shadowLabels
+
+  override def prepareForApply(tree: Apply)(using Context): Context = tree match
+    case Break(_, _) => ctx
+    case _ => stackChange
+
+  // other stack changing operations are handled in RecordStackChange
+
   /** If `tree` is a BreakBoundary, transform it as follows:
    *   - Wrap it in a labeled block if its label has local uses
    *   - Drop the try/catch if its label has no other uses
@@ -147,29 +188,9 @@ class DropBreaks extends MiniPhase:
     case _ =>
       tree
 
-  private def isBreak(sym: Symbol)(using Context): Boolean =
-    sym.name == nme.apply && sym.owner == defn.breakModule.moduleClass
-
-  private def transformBreak(tree: Tree, arg: Tree, lbl: Symbol)(using Context): Tree =
-    report.log(i"transform break $tree/$arg/$lbl")
-    labelUsage(lbl) match
-      case Some(uses: LabelUsage)
-      if uses.enclMeth == ctx.owner.enclosingMethod
-          && !ctx.property(LabelsShadowedByTry).getOrElse(Set.empty).contains(lbl)
-        =>
-        uses.otherRefs -= 1
-        uses.returnRefs += 1
-        Return(arg, ref(uses.goto)).withSpan(arg.span)
-      case _ =>
-        tree
-
-
-  /** Rewrite a break call
-   *
-   *     break.apply[...](value)(using lbl)
-   *
+  /** Rewrite a break with argument `arg` and label `lbl`
    *  where `lbl` is a label defined in the current method and is not included in
-   *  LabelsShadowedByTry to
+   *  ShadowedLabels to
    *
    *     return[target] arg
    *
@@ -179,22 +200,29 @@ class DropBreaks extends MiniPhase:
    *  binding containing `local` is dropped.
    */
   override def transformApply(tree: Apply)(using Context): Tree = tree match
-    case Apply(Apply(fn, args), (id: Ident) :: Nil) if isBreak(fn.symbol) =>
-      val arg = (args: @unchecked) match
-        case arg :: Nil => arg
-        case Nil => Literal(Constant(())).withSpan(tree.span)
-      transformBreak(tree, arg, id.symbol)
-    case _ =>
-      tree
+    case Break(lbl, arg) =>
+      labelUsage(lbl) match
+        case Some(uses: LabelUsage)
+        if uses.enclMeth == ctx.owner.enclosingMethod
+            && !ctx.property(ShadowedLabels).getOrElse(Set.empty).contains(lbl)
+          =>
+          uses.otherRefs -= 1
+          uses.returnRefs += 1
+          Return(arg, ref(uses.goto)).withSpan(arg.span)
+        case _ => tree
+    case _ => tree
 
   /** If `tree` refers to an enclosing label, increase its non local recount.
    *  This increase is corrected in `transformInlined` if the reference turns
    *  out to be part of a BreakThrow to a local, non-shadowed label.
    */
   override def transformIdent(tree: Ident)(using Context): Tree =
-    if tree.symbol.name == nme.local then
-      for uses <- labelUsage(tree.symbol) do
-        uses.otherRefs += 1
+    for uses <- labelUsage(tree.symbol) do
+      uses.otherRefs += 1
     tree
 
+  //override def transformReturn(tree: Return)(using Context): Tree =
+  //  if !tree.from.isEmpty && tree.expr.tpe.isExactlyNothing then tree.expr
+  //  else tree
+
 end DropBreaks

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

@@ -6,8 +6,8 @@ import core.DenotTransformers._
 import core.Symbols._
 import core.Contexts._
 import core.Types._
-import core.Flags._
 import core.Decorators._
+import core.Flags._
 import core.NameKinds.LiftedTreeName
 import NonLocalReturns._
 import util.Store
@@ -27,7 +27,7 @@ import util.Store
  *  after an exception, so the fact that values on the stack are 'lost' does not matter
  *  (copied from https://github.com/scala/scala/pull/922).
  */
-class LiftTry extends MiniPhase with IdentityDenotTransformer { thisPhase =>
+class LiftTry extends MiniPhase, IdentityDenotTransformer, RecordStackChange { thisPhase =>
   import ast.tpd._
 
   override def phaseName: String = LiftTry.name
@@ -40,35 +40,14 @@ class LiftTry extends MiniPhase with IdentityDenotTransformer { thisPhase =>
   override def initContext(ctx: FreshContext): Unit =
     NeedLift = ctx.addLocation(false)
 
-  private def liftingCtx(p: Boolean)(using Context) =
+  private def liftingCtx(p: Boolean)(using Context): Context =
     if (needLift == p) ctx else ctx.fresh.updateStore(NeedLift, p)
 
-  override def prepareForApply(tree: Apply)(using Context): Context =
-    liftingCtx(true)
+  protected def stackChange(using Context): Context = liftingCtx(true)
 
   override def prepareForDefDef(tree: DefDef)(using Context): Context =
     liftingCtx(false)
 
-  override def prepareForValDef(tree: ValDef)(using Context): Context =
-    if !tree.symbol.exists
-       || tree.symbol.isSelfSym
-       || tree.symbol.owner == ctx.owner.enclosingMethod
-          && !tree.symbol.is(Lazy)
-            // The current implementation wraps initializers of lazy vals in
-            // calls to an initialize method, which means that a `try` in the
-            // initializer needs to be lifted. Note that the new scheme proposed
-            // in #6979 would avoid this.
-    then ctx
-    else liftingCtx(true)
-
-  override def prepareForAssign(tree: Assign)(using Context): Context =
-    if (tree.lhs.symbol.maybeOwner == ctx.owner.enclosingMethod) ctx
-    else liftingCtx(true)
-
-  override def prepareForReturn(tree: Return)(using Context): Context =
-    if (!isNonLocalReturn(tree)) ctx
-    else liftingCtx(true)
-
   override def prepareForTemplate(tree: Template)(using Context): Context =
     liftingCtx(false)
 

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

@@ -0,0 +1,41 @@
+package dotty.tools.dotc
+package transform
+
+import MegaPhase.MiniPhase
+import core.*
+import Contexts.*
+import Flags.Lazy
+import NonLocalReturns.isNonLocalReturn
+
+/** A trait shared between LiftTry and DropBreaks.
+ *  Overrides prepare methods for trees that push to the stack before some of their elements
+ *  are evaluated.
+ */
+trait RecordStackChange extends MiniPhase:
+  import ast.tpd.*
+
+  /** The method to call to record a stack change */
+  protected def stackChange(using Context): Context
+
+  override def prepareForApply(tree: Apply)(using Context): Context =
+    stackChange
+
+  override def prepareForValDef(tree: ValDef)(using Context): Context =
+    if !tree.symbol.exists
+       || tree.symbol.isSelfSym
+       || tree.symbol.owner == ctx.owner.enclosingMethod
+          && !tree.symbol.is(Lazy)
+            // The current implementation wraps initializers of lazy vals in
+            // calls to an initialize method, which means that a `try` in the
+            // initializer needs to be lifted. Note that the new scheme proposed
+            // in #6979 would avoid this.
+    then ctx
+    else stackChange
+
+  override def prepareForAssign(tree: Assign)(using Context): Context =
+    if (tree.lhs.symbol.maybeOwner == ctx.owner.enclosingMethod) ctx
+    else stackChange
+
+  override def prepareForReturn(tree: Return)(using Context): Context =
+    if (!isNonLocalReturn(tree)) ctx
+    else stackChange

tests/run/break-opt.scala

@@ -57,6 +57,21 @@ object breakOpt:
         if x == 0 then break()
     y
 
+  def test7(x0: Int): Option[Int] =
+    boundary:
+      Some(
+        if x0 < 0 then break(None) // no jump possible, since stacksize changes
+        else x0 + 1
+      )
+
+
+  def test8(x0: Int): Option[Int] =
+    boundary:
+      lazy val x =
+        if x0 < 0 then break(None) // no jump possible, since stacksize changes
+        else x0 + 1
+      Some(x)
+
 @main def Test =
   import breakOpt.*
   assert(test1(0) == 0)
@@ -68,5 +83,7 @@ object breakOpt:
   test4(-1)
   assert(test5(2) == 1)
   assert(test6(3) == 18)
+  assert(test7(3) == Some(4))
+  assert(test7(-3) == None)