Mark trailing map receives body to inspect

Author: som-snytt

compiler/src/dotty/tools/dotc/ast/Desugar.scala

@@ -2236,7 +2236,7 @@ object desugar {
           case (Tuple(ts1), Tuple(ts2)) => ts1.corresponds(ts2)(deepEquals)
           case _ => false
 
-      def markTrailingMap(aply: Apply, gen: GenFrom, selectName: TermName): Unit =
+      def markTrailingMap(aply: Apply, gen: GenFrom, selectName: TermName, body: Tree): Unit =
         if sourceVersion.enablesBetterFors
           && selectName == mapName
           && gen.checkMode != GenCheckMode.Filtered // results of withFilter have the wrong type
@@ -2247,9 +2247,8 @@ object desugar {
       enums match {
         case Nil if sourceVersion.enablesBetterFors => body
         case (gen: GenFrom) :: Nil =>
-          val aply = Apply(rhsSelect(gen, mapName), makeLambda(gen, body))
-          markTrailingMap(aply, gen, mapName)
-          aply
+          Apply(rhsSelect(gen, mapName), makeLambda(gen, body))
+            .tap(markTrailingMap(_, gen, mapName, body))
         case (gen: GenFrom) :: (rest @ (GenFrom(_, _, _) :: _)) =>
           val cont = makeFor(mapName, flatMapName, rest, body)
           Apply(rhsSelect(gen, flatMapName), makeLambda(gen, cont))
@@ -2266,7 +2265,7 @@ object desugar {
               if suffix.exists(_.isInstanceOf[GenFrom]) then flatMapName
               else mapName
             Apply(rhsSelect(gen, selectName), makeLambda(gen, cont))
-              .tap(markTrailingMap(_, gen, selectName))
+              .tap(markTrailingMap(_, gen, selectName, cont))
           else
             val (pats, rhss) = valeqs.map { case GenAlias(pat, rhs) => (pat, rhs) }.unzip
             val (defpat0, id0) = makeIdPat(gen.pat)

tests/run/i24673.scala

@@ -0,0 +1,13 @@
+@main def Test = {
+  def result = for {
+    a <- Option(2)
+    _ = if (true) {
+      sys.error("err")
+    }
+  } yield a
+
+  try
+    result
+    ???
+  catch case e: RuntimeException => assert(e.getMessage == "err")
+}